import { ceilPowerOfTwo } from "../../../core/binary/operations/ceilPowerOfTwo.js";
import { body_id_index } from "../body/BodyStorage.js";
import { BodyKind } from "../ecs/BodyKind.js";
import { is_sensor } from "../ecs/is_sensor.js";
import { JOINT_WORLD } from "../ecs/Joint.js";
import { uf_find, uf_init, uf_union } from "./union_find.js";

/**
 * Partitions this frame's awake dynamic bodies + touched non-sensor contacts
 * into connected components ("islands") via union-find. The solver iterates
 * each island independently (so impulses converge inside an island without
 * waiting for global passes) and the sleep test will eventually use island
 * granularity for atomic whole-island sleep.
 *
 * Static and Kinematic bodies are not merged into islands — they act as
 * constraint anchors. A 1000-block stack on a static floor is therefore one
 * island of 1000 dynamic bodies, not 1001 — the floor anchors the island but
 * does not enlarge it. Multiple separate piles all resting on the same static
 * floor remain *separate* islands so each can sleep/wake independently.
 *
 * Sensor contacts are skipped entirely: they don't transmit constraint
 * forces, so two bodies linked only by a sensor pair are not in the same
 * structural island.
 *
 * Determinism contract:
 *  - Union-find uses union-by-min-index + path halving, so the canonical
 *    root of any component is the smallest body index of any of its members.
 *  - Islands are emitted sorted ascending by root index.
 *  - Bodies within an island are emitted sorted ascending by body index.
 *  - Contacts within an island are emitted sorted ascending by manifold slot id.
 *
 * Output layout is CSR-like — `body_offsets[i]..body_offsets[i+1]` indexes
 * into `body_data` for island `i`, similarly for contacts.
 *
 * @author Alex Goldring
 * @copyright Company Named Limited (c) 2026
 */
export class IslandBuilder {

    constructor() {
        /**
         * Union-find parent table, indexed by body index (NOT packed id).
         * Sized to `storage.high_water_mark` at `build` time; reallocated by
         * doubling when the body pool grows past current capacity.
         * @type {Uint32Array}
         */
        this.parent = new Uint32Array(16);

        /**
         * Per-body island id. `island_of_body[idx] === -1` means the body
         * is not in any island this frame (static, kinematic, or unallocated).
         * @type {Int32Array}
         */
        this.island_of_body = new Int32Array(16);

        /**
         * Number of islands this frame.
         * @type {number}
         */
        this.island_count = 0;

        /**
         * CSR offsets into `body_data`. Length is `island_count + 1`. Bodies
         * for island `i` live in `body_data[body_offsets[i]..body_offsets[i+1])`.
         * @type {Uint32Array}
         */
        this.body_offsets = new Uint32Array(2);

        /**
         * Flattened body indices, grouped by island, sorted ascending within
         * each island.
         * @type {Uint32Array}
         */
        this.body_data = new Uint32Array(16);

        /**
         * Total number of body entries across all islands (= sum of island sizes).
         * @type {number}
         */
        this.body_total = 0;

        /**
         * CSR offsets into `contact_data`. Length is `island_count + 1`.
         * Manifold slot ids for island `i` live in
         * `contact_data[contact_offsets[i]..contact_offsets[i+1])`.
         * @type {Uint32Array}
         */
        this.contact_offsets = new Uint32Array(2);

        /**
         * Flattened manifold slot ids, grouped by island, sorted ascending
         * within each island.
         * @type {Uint32Array}
         */
        this.contact_data = new Uint32Array(16);

        /**
         * Total number of contact entries across all islands.
         * @type {number}
         */
        this.contact_total = 0;

        /**
         * Scratch: root body index → island id, or -1 if unassigned.
         * Sized in lockstep with `parent`.
         * @private
         * @type {Int32Array}
         */
        this.__root_to_island = new Int32Array(16);

        /**
         * Scratch: list of distinct root indices observed this frame, used
         * to sort them ascending before assigning island ids.
         * @private
         * @type {Uint32Array}
         */
        this.__scratch_roots = new Uint32Array(16);

        /**
         * Scratch write-cursors, one per island.
         * @private
         * @type {Uint32Array}
         */
        this.__cursors = new Uint32Array(2);
    }

    /**
     * @param {RigidBody} rb
     * @param {Array} collider_list collider list for the body (or undefined)
     * @returns {boolean}
     * @private
     */
    static __body_is_sensor(rb, collider_list) {
        // v1 approximation: the body's primary (first) collider stands in for
        // the whole body, matching the solver's pair_is_sensor / raycast checks.
        const collider = (collider_list !== undefined && collider_list.length > 0)
            ? collider_list[0].collider
            : null;
        return is_sensor(rb, collider);
    }

    /**
     * (Re)build islands from `manifolds`' currently-touched non-sensor slots
     * and the bodies in `storage`'s awake list.
     *
     * @param {BodyStorage} storage
     * @param {ManifoldStore} manifolds
     * @param {RigidBody[]} bodies sparse, indexed by body index
     * @param {Array[]} body_collider_lists sparse, indexed by body index
     * @param {Joint[]} joints live joints (sparse). Jointed dynamic-dynamic
     *   bodies are unioned into the same island so a chain / ragdoll sleeps
     *   and wakes as a unit. Callers with no joints pass an empty array.
     */
    build(storage, manifolds, bodies, body_collider_lists, joints) {
        const hwm = storage.high_water_mark;
        this.__ensure_body_capacity(hwm);

        const parent = this.parent;
        uf_init(parent, hwm);

        // --- Pass 1: union dynamic-dynamic pairs from touched non-sensor manifolds.
        const live_count = manifolds.count;
        for (let i = 0; i < live_count; i++) {
            const slot = manifolds.live_at(i);
            if (!manifolds.is_touched(slot)) continue;
            const idxA = body_id_index(manifolds.bodyA(slot));
            const idxB = body_id_index(manifolds.bodyB(slot));
            const rbA = bodies[idxA];
            const rbB = bodies[idxB];
            if (rbA === undefined || rbB === undefined) continue;
            if (IslandBuilder.__body_is_sensor(rbA, body_collider_lists[idxA])) continue;
            if (IslandBuilder.__body_is_sensor(rbB, body_collider_lists[idxB])) continue;
            if (rbA.kind === BodyKind.Dynamic && rbB.kind === BodyKind.Dynamic) {
                uf_union(parent, idxA, idxB);
            }
        }

        // --- Pass 1b: union dynamic-dynamic bodies connected by a joint, so a
        //     chain / ragdoll forms one island (and so sleeps/wakes as a unit).
        //     Joint-to-world and joint-to-static/kinematic anchor the island
        //     without enlarging it — same rule as a static contact. Stale joint
        //     references (body unlinked / slot reused) are filtered by the
        //     generation-checked `is_valid`. (Empty when the caller has no
        //     joints — the loop is then a no-op.)
        const jn = joints.length;
        for (let i = 0; i < jn; i++) {
            const joint = joints[i];
            if (joint === undefined || joint === null) continue;
            if (joint._bodyIdB === JOINT_WORLD) continue;
            if (!storage.is_valid(joint._bodyIdA) || !storage.is_valid(joint._bodyIdB)) continue;
            const idxA = body_id_index(joint._bodyIdA);
            const idxB = body_id_index(joint._bodyIdB);
            const rbA = bodies[idxA];
            const rbB = bodies[idxB];
            if (rbA === undefined || rbB === undefined) continue;
            if (rbA.kind === BodyKind.Dynamic && rbB.kind === BodyKind.Dynamic) {
                uf_union(parent, idxA, idxB);
            }
        }

        // --- Pass 2: collect distinct roots over awake dynamic bodies.
        const island_of_body = this.island_of_body;
        // Cheap reset: only the indices we may write are below hwm.
        for (let i = 0; i < hwm; i++) island_of_body[i] = -1;

        const root_to_island = this.__root_to_island;
        for (let i = 0; i < hwm; i++) root_to_island[i] = -1;

        const awake_count = storage.awake_count;
        const scratch_roots = this.__scratch_roots;
        let root_count = 0;
        for (let ai = 0; ai < awake_count; ai++) {
            const idx = storage.awake_at(ai);
            const rb = bodies[idx];
            if (rb === undefined) continue;
            if (rb.kind !== BodyKind.Dynamic) continue;
            const r = uf_find(parent, idx);
            if (root_to_island[r] === -1) {
                scratch_roots[root_count++] = r;
                // Tag as collected — actual id assigned after sort.
                root_to_island[r] = -2;
            }
        }

        // Sort distinct roots ascending so island id 0 has the smallest root.
        if (root_count > 1) {
            scratch_roots.subarray(0, root_count).sort();
        }
        for (let i = 0; i < root_count; i++) {
            root_to_island[scratch_roots[i]] = i;
        }
        this.island_count = root_count;

        // Assign island id to every awake dynamic body.
        for (let ai = 0; ai < awake_count; ai++) {
            const idx = storage.awake_at(ai);
            const rb = bodies[idx];
            if (rb === undefined) continue;
            if (rb.kind !== BodyKind.Dynamic) continue;
            const r = uf_find(parent, idx);
            island_of_body[idx] = root_to_island[r];
        }

        this.__build_body_csr(storage, bodies, awake_count, root_count);
        this.__build_contact_csr(manifolds, bodies, body_collider_lists, root_count);
    }

    /**
     * Fill `body_offsets` + `body_data` with awake dynamic bodies grouped
     * by island, sorted ascending within each island.
     * @private
     */
    __build_body_csr(storage, bodies, awake_count, island_count) {
        this.__ensure_island_count_capacity(island_count);
        const offsets = this.body_offsets;
        for (let i = 0; i <= island_count; i++) offsets[i] = 0;

        // Count pass: offsets[i + 1] starts as the number of bodies in island i.
        let total = 0;
        for (let ai = 0; ai < awake_count; ai++) {
            const idx = storage.awake_at(ai);
            const rb = bodies[idx];
            if (rb === undefined) continue;
            if (rb.kind !== BodyKind.Dynamic) continue;
            const isl = this.island_of_body[idx];
            if (isl < 0) continue;
            offsets[isl + 1]++;
            total++;
        }
        // Prefix sum → start offsets.
        for (let i = 0; i < island_count; i++) {
            offsets[i + 1] += offsets[i];
        }
        this.body_total = total;

        this.__ensure_body_data_capacity(total);

        const cursors = this.__cursors;
        for (let i = 0; i < island_count; i++) cursors[i] = offsets[i];

        const data = this.body_data;
        for (let ai = 0; ai < awake_count; ai++) {
            const idx = storage.awake_at(ai);
            const rb = bodies[idx];
            if (rb === undefined) continue;
            if (rb.kind !== BodyKind.Dynamic) continue;
            const isl = this.island_of_body[idx];
            if (isl < 0) continue;
            data[cursors[isl]++] = idx;
        }

        // Sort each island's body slice ascending.
        for (let i = 0; i < island_count; i++) {
            const start = offsets[i];
            const end = offsets[i + 1];
            if (end - start > 1) {
                data.subarray(start, end).sort();
            }
        }
    }

    /**
     * Fill `contact_offsets` + `contact_data` with touched non-sensor manifold
     * slot ids grouped by island. A contact belongs to its dynamic
     * participant's island (if both sides are dynamic they share an island
     * by construction). Static-vs-static contacts are skipped.
     * @private
     */
    __build_contact_csr(manifolds, bodies, body_collider_lists, island_count) {
        const offsets = this.contact_offsets;
        for (let i = 0; i <= island_count; i++) offsets[i] = 0;

        // Count pass.
        const live_count = manifolds.count;
        let total = 0;
        for (let i = 0; i < live_count; i++) {
            const slot = manifolds.live_at(i);
            const isl = this.__island_of_slot(manifolds, slot, bodies, body_collider_lists);
            if (isl < 0) continue;
            offsets[isl + 1]++;
            total++;
        }
        for (let i = 0; i < island_count; i++) {
            offsets[i + 1] += offsets[i];
        }
        this.contact_total = total;

        this.__ensure_contact_data_capacity(total);

        const cursors = this.__cursors;
        for (let i = 0; i < island_count; i++) cursors[i] = offsets[i];

        const data = this.contact_data;
        for (let i = 0; i < live_count; i++) {
            const slot = manifolds.live_at(i);
            const isl = this.__island_of_slot(manifolds, slot, bodies, body_collider_lists);
            if (isl < 0) continue;
            data[cursors[isl]++] = slot;
        }

        // Sort each island's contact slice ascending by slot id.
        for (let i = 0; i < island_count; i++) {
            const start = offsets[i];
            const end = offsets[i + 1];
            if (end - start > 1) {
                data.subarray(start, end).sort();
            }
        }
    }

    /**
     * Resolve which island a manifold slot belongs to.
     *
     * Returns `-1` if the manifold is not touched this frame, is a sensor
     * pair, or both bodies are non-dynamic (static-static, which the solver
     * wouldn't act on anyway).
     *
     * @private
     */
    __island_of_slot(manifolds, slot, bodies, body_collider_lists) {
        if (!manifolds.is_touched(slot)) return -1;
        const idxA = body_id_index(manifolds.bodyA(slot));
        const idxB = body_id_index(manifolds.bodyB(slot));
        const rbA = bodies[idxA];
        const rbB = bodies[idxB];
        if (rbA === undefined || rbB === undefined) return -1;
        if (IslandBuilder.__body_is_sensor(rbA, body_collider_lists[idxA])) return -1;
        if (IslandBuilder.__body_is_sensor(rbB, body_collider_lists[idxB])) return -1;
        // Pick the dynamic side's island. If both are dynamic they share one by
        // construction (pass-1 union). Guard `>= 0`: a dynamic body that is ASLEEP
        // has island_of_body === -1 (islands cover only awake bodies). Returning
        // that -1 would drop the contact from the island, so the live
        // non-penetration constraint against the awake side would never be solved
        // — e.g. an awake body resting on a freshly-slept neighbour falls through.
        // Falling through to side B keeps the awake dynamic side's island.
        const ia = this.island_of_body[idxA];
        if (rbA.kind === BodyKind.Dynamic && ia >= 0) return ia;
        const ib = this.island_of_body[idxB];
        if (rbB.kind === BodyKind.Dynamic && ib >= 0) return ib;
        return -1;
    }

    // --- Resize helpers ------------------------------------------------------

    /**
     * @private
     * @param {number} n
     */
    __ensure_body_capacity(n) {
        if (this.parent.length < n) {
            const cap = ceilPowerOfTwo(n);
            this.parent = new Uint32Array(cap);
            this.island_of_body = new Int32Array(cap);
            this.__root_to_island = new Int32Array(cap);
            this.__scratch_roots = new Uint32Array(cap);
        }
    }

    /**
     * @private
     * @param {number} n
     */
    __ensure_island_count_capacity(n) {
        if (this.body_offsets.length < n + 1) {
            const cap = ceilPowerOfTwo(n + 1);
            this.body_offsets = new Uint32Array(cap);
            this.contact_offsets = new Uint32Array(cap);
            this.__cursors = new Uint32Array(cap);
        }
    }

    /**
     * @private
     * @param {number} n
     */
    __ensure_body_data_capacity(n) {
        if (this.body_data.length < n) {
            this.body_data = new Uint32Array(ceilPowerOfTwo(n));
        }
    }

    /**
     * @private
     * @param {number} n
     */
    __ensure_contact_data_capacity(n) {
        if (this.contact_data.length < n) {
            this.contact_data = new Uint32Array(ceilPowerOfTwo(n));
        }
    }
}
