import {
    BVH,
    COLUMN_CHILD_1,
    COLUMN_CHILD_2,
    COLUMN_USER_DATA,
    ELEMENT_WORD_COUNT,
    NULL_NODE
} from "../../../../core/bvh2/bvh3/BVH.js";
import { ebvh_build_for_geometry_morton } from "../../../../core/bvh2/bvh3/ebvh_build_for_geometry_morton.js";
import { ebvh_optimize_treelet } from "../../../../core/bvh2/bvh3/ebvh_optimize_treelet.js";
import { bvh_query_user_data_ray_segment } from "../../../../core/bvh2/bvh3/query/bvh_query_user_data_ray_segment.js";
import { array_copy } from "../../../../core/collection/array/array_copy.js";
import { SCRATCH_UINT32_TRAVERSAL_STACK } from "../../../../core/collection/SCRATCH_UINT32_TRAVERSAL_STACK.js";
import { AABB3 } from "../../../../core/geom/3d/aabb/AABB3.js";
import { aabb3_from_v3_array } from "../../../../core/geom/3d/aabb/aabb3_from_v3_array.js";
import { aabb3_intersects_ray_segment } from "../../../../core/geom/3d/aabb/aabb3_intersects_ray_segment.js";
import { SurfacePoint3 } from "../../../../core/geom/3d/SurfacePoint3.js";

import {
    computeTriangleRayIntersectionBarycentricGeometry
} from "../../../../core/geom/3d/triangle/computeTriangleRayIntersectionBarycentricGeometry.js";
import Vector4 from "../../../../core/geom/Vector4.js";
import { makeGeometryIndexed } from "../../geometry/buffered/makeGeometryIndexed.js";
import { query_bvh_geometry_nearest } from "../../geometry/buffered/query/query_bvh_geometry_nearest.js";
import { computeBoundingSphereFromVertexData } from "../../geometry/computeBoundingSphereFromVertexData.js";
import { construct_ray_hit_from_geometry } from "./geometry/construct_ray_hit_from_geometry.js";

/**
 *
 * @type {number[]|Uint32Array}
 */
const scratch_uint32_array = new Uint32Array(4096);
/**
 *
 * @type {number[]}
 */
const v3_scratch_0 = [];

const stack = SCRATCH_UINT32_TRAVERSAL_STACK;

export class BufferedGeometryBVH {
    /**
     * @type {Uint32Array}
     */
    #morton_codes

    #bounds = new AABB3()

    #bounds_sphere = new Float32Array(4);


    /**
     *
     * @type {BVH}
     * @private
     */
    #bvh = new BVH();

    /**
     *
     * @param {AABB3|number[]} out
     */
    getBounds(out) {
        array_copy(this.#bounds, 0, out, 0, 6);
    }

    constructor() {

        /**
         *
         * @type {THREE.BufferGeometry|null}
         * @private
         */
        this.__geometry = null;

        /**
         *
         * @type {number[]|null}
         * @private
         */
        this.__geometry_index = null;

        /**
         *
         * @type {number[]|null}
         * @private
         */
        this.__geometry_positions = null;


        /**
         *
         * @type {number}
         * @private
         */
        this.__triangle_count = 0;
    }

    /**
     *
     * @param {THREE.BufferGeometry} geo
     */
    build(geo) {
        makeGeometryIndexed(geo);

        this.__geometry = geo;

        const index = geo.getIndex();
        const index_array = index.array;


        const attribute_position = geo.getAttribute('position');
        const array_positions = attribute_position.array;

        this.__geometry_positions = array_positions;
        this.__geometry_index = index_array;

        this.__triangle_count = index_array.length / 3;

        this.#morton_codes = new Uint32Array(this.__triangle_count);

        const bvh = this.#bvh;

        aabb3_from_v3_array(this.#bounds, array_positions, array_positions.length);

        ebvh_build_for_geometry_morton(
            bvh,
            index_array,
            array_positions,
            this.#morton_codes,
            this.#bounds,
            2
        );

        for (let i = 0; i < 3; i++) {
            ebvh_optimize_treelet(bvh);
        }

        //
        // ebvh_build_for_geometry_incremental(bvh, index_array, array_positions, 1);

        // remove any extra unused space
        bvh.trim();

        const r = new Vector4();
        computeBoundingSphereFromVertexData(array_positions, r);

        r.toArray(this.#bounds_sphere);
    }

    /**
     * Tests ray for occlusion
     * Returns true if ray hits anything at all
     * @param {Ray3} ray
     * @returns {boolean}
     */
    occluded(ray) {
        const indices = this.__geometry_index;
        const positions = this.__geometry_positions;

        const origin_x = ray[0];
        const origin_y = ray[1];
        const origin_z = ray[2];

        const direction_x = ray[3];
        const direction_y = ray[4];
        const direction_z = ray[5];

        const max_distance = ray[6];

        const bvh = this.#bvh;

        const root = bvh.root;

        if (root === NULL_NODE) {
            return -1;
        }

        /**
         * Move stack pointer to local variable scope to avoid de-referencing inside the loop
         * @type {number}
         */
        let pointer = stack.pointer;

        /**
         *
         * @type {number}
         */
        const stack_top = pointer;

        stack[pointer++] = root;


        /*
         For performance, we bind data directly to avoid extra copies required to read out AABB
         */
        const float32 = bvh.__data_float32;
        const uint32 = bvh.__data_uint32;

        const inv_direction_x = 1 / direction_x;
        const inv_direction_y = 1 / direction_y;
        const inv_direction_z = 1 / direction_z;

        do {

            --pointer;

            /**
             *
             * @type {number}
             */
            const node = stack[pointer];

            const address = node * ELEMENT_WORD_COUNT;

            // test node against the ray
            const intersects = aabb3_intersects_ray_segment(
                float32[address], float32[address + 1], float32[address + 2],
                float32[address + 3], float32[address + 4], float32[address + 5],
                origin_x, origin_y, origin_z,
                inv_direction_x, inv_direction_y, inv_direction_z,
                0, max_distance
            );

            if (!intersects) {
                continue;
            }

            // get fist child to check if this is a leaf node or not
            const child_1 = uint32[address + COLUMN_CHILD_1];

            if (child_1 !== NULL_NODE) {

                // this is not a leaf node, push children onto traversal stack
                const child_2 = uint32[address + COLUMN_CHILD_2];

                stack[pointer++] = child_2;
                stack[pointer++] = child_1;


            } else {
                // leaf node

                const triangle_index = uint32[address + COLUMN_USER_DATA];

                const intersection_found = computeTriangleRayIntersectionBarycentricGeometry(
                    v3_scratch_0,
                    origin_x, origin_y, origin_z,
                    direction_x, direction_y, direction_z,
                    indices, triangle_index,
                    positions
                );

                if (!intersection_found) {
                    continue;
                }

                const t = v3_scratch_0[0];

                if (t < max_distance && t > 0) {
                    return true;

                }

            }
        } while (pointer > stack_top);


        return false;
    }

    /**
     *
     * @param {SurfacePoint3} result
     * @param {number} x
     * @param {number} y
     * @param {number} z
     * @returns {boolean} true if result is found, only false when geometry is empty
     */
    nearestSurfacePoint(result, x, y, z) {

        const indices = this.__geometry_index;
        const positions = this.__geometry_positions;

        const bvh = this.#bvh;

        return query_bvh_geometry_nearest(result, bvh, positions, indices, x, y, z, Infinity);
    }


    /**
     * Code is largely inlined, to avoid extra checks
     * NOTE: raycast is performed in local coordinate space
     * @param {number[]} output
     * @param {number[]|Ray3} ray
     * @returns {number} distance along the ray, negative if no hit
     */
    raycast2(output, ray) {

        const indices = this.__geometry_index;
        const positions = this.__geometry_positions;

        const origin_x = ray[0];
        const origin_y = ray[1];
        const origin_z = ray[2];

        const direction_x = ray[3];
        const direction_y = ray[4];
        const direction_z = ray[5];

        const max_distance = ray[6];

        const bvh = this.#bvh;

        const root = bvh.root;

        if (root === NULL_NODE) {
            return -1;
        }

        /**
         * Move stack pointer to local variable scope to avoid de-referencing inside the loop
         * @type {number}
         */
        let pointer = stack.pointer;

        let nearest_hit_distance = max_distance;

        let best_index = -1;
        let best_u = 0;
        let best_v = 0;

        /**
         *
         * @type {number}
         */
        const stack_top = pointer;

        stack[pointer++] = root;


        /*
         For performance, we bind data directly to avoid extra copies required to read out AABB
         */
        const float32 = bvh.__data_float32;
        const uint32 = bvh.__data_uint32;


        const inv_direction_x = 1 / direction_x;
        const inv_direction_y = 1 / direction_y;
        const inv_direction_z = 1 / direction_z;

        do {

            --pointer;

            /**
             *
             * @type {number}
             */
            const node = stack[pointer];

            const address = node * ELEMENT_WORD_COUNT;

            // test node against the ray
            const intersects = aabb3_intersects_ray_segment(
                float32[address], float32[address + 1], float32[address + 2],
                float32[address + 3], float32[address + 4], float32[address + 5],
                origin_x, origin_y, origin_z,
                inv_direction_x, inv_direction_y, inv_direction_z,
                0, nearest_hit_distance
            );

            if (!intersects) {
                continue;
            }

            // get fist child to check if this is a leaf node or not
            const child_1 = uint32[address + COLUMN_CHILD_1];

            if (child_1 !== NULL_NODE) {

                // this is not a leaf node, push children onto traversal stack
                const child_2 = uint32[address + COLUMN_CHILD_2];

                stack[pointer++] = child_2;
                stack[pointer++] = child_1;

            } else {
                // leaf node

                const triangle_index = uint32[address + COLUMN_USER_DATA];

                const intersection_found = computeTriangleRayIntersectionBarycentricGeometry(
                    v3_scratch_0,
                    origin_x, origin_y, origin_z,
                    direction_x, direction_y, direction_z,
                    indices, triangle_index,
                    positions
                );

                if (!intersection_found) {
                    continue;
                }

                const t = v3_scratch_0[0];

                if (t < nearest_hit_distance && t > 0) {
                    nearest_hit_distance = t;

                    best_index = triangle_index;
                    best_u = v3_scratch_0[1];
                    best_v = v3_scratch_0[2];

                }

            }
        } while (pointer > stack_top);

        if (nearest_hit_distance === max_distance) {
            // no hit
            return -1;
        }

        construct_ray_hit_from_geometry(output, indices, positions, best_index, nearest_hit_distance, best_u, best_v);

        return nearest_hit_distance;
    }

    /**
     * NOTE: raycast is performed in local coordinate space
     * @param {number[]} output
     * @param {number[]|Ray3} ray
     * @returns {number} distance along the ray, negative if no hit
     */
    raycast(output, ray) {

        const indices = this.__geometry_index;
        const positions = this.__geometry_positions;

        const origin_x = ray[0];
        const origin_y = ray[1];
        const origin_z = ray[2];

        const direction_x = ray[3];
        const direction_y = ray[4];
        const direction_z = ray[5];

        const max_distance = ray[6];

        const bvh = this.#bvh;

        const count = bvh_query_user_data_ray_segment(
            bvh, bvh.root,
            scratch_uint32_array, 0,
            origin_x, origin_y, origin_z,
            direction_x, direction_y, direction_z,
            0, max_distance
        );

        if (count === 0) {
            // no bvh hit, early exit
            return -1;
        }


        let nearest_hit_distance = max_distance;

        let best_index = -1;
        let best_u = 0;
        let best_v = 0;

        // check triangles found via BVH
        for (let i = 0; i < count; i++) {
            const triangle_index = scratch_uint32_array[i];

            const intersection_found = computeTriangleRayIntersectionBarycentricGeometry(
                v3_scratch_0,
                origin_x, origin_y, origin_z,
                direction_x, direction_y, direction_z,
                indices, triangle_index,
                positions
            );

            if (!intersection_found) {
                continue;
            }

            const t = v3_scratch_0[0];

            if (t < nearest_hit_distance && t > 0) {
                nearest_hit_distance = t;

                best_index = triangle_index;
                best_u = v3_scratch_0[1];
                best_v = v3_scratch_0[2];

            }
        }

        if (nearest_hit_distance === max_distance) {
            // no hit
            return -1;
        }

        construct_ray_hit_from_geometry(output, indices, positions, best_index, nearest_hit_distance, best_u, best_v);

        return nearest_hit_distance;
    }
}