import { assert } from "../../../core/assert.js";
import { BodyKind } from "../ecs/BodyKind.js";

const DEFAULT_INITIAL_CAPACITY = 16;
const MIN_CAPACITY = 4;

/**
 * Mask for the 8-bit generation field encoded in a packed body id.
 * @type {number}
 */
const GENERATION_MASK = 0xFF;

/**
 * Bit width of the generation field. The remaining bits are the body index;
 * 24-bit indices give us 16M live bodies which is comfortably above the design
 * target of "millions".
 * @type {number}
 */
const GENERATION_BITS = 8;

/**
 * Sentinel returned by {@link BodyStorage#awake_position_of} and friends when a
 * body is not in the awake set.
 * @type {number}
 */
export const BODY_INDEX_ABSENT = -1;

/**
 * Pack a body index and generation into a single integer handle.
 *
 * @param {number} index
 * @param {number} generation
 * @returns {number}
 */
export function pack_body_id(index, generation) {
    return (index << GENERATION_BITS) | (generation & GENERATION_MASK);
}

/**
 *
 * @param {number} packed
 * @returns {number}
 */
export function body_id_index(packed) {
    return packed >>> GENERATION_BITS;
}

/**
 *
 * @param {number} packed
 * @returns {number}
 */
export function body_id_generation(packed) {
    return packed & GENERATION_MASK;
}

/**
 * Structure-of-arrays pool for rigid bodies.
 *
 * Owns:
 *  - per-body identity (entity, generation, kind, flags),
 *  - a dense list of awake body indices (the simulation hot iteration target),
 *  - a min-heap of free body indices so reuse order is deterministic regardless
 *    of how interleaved allocate / free calls have been.
 *
 * The pool grows by doubling when the high-water mark reaches capacity; arrays
 * are replaced wholesale on grow, so callers must not retain references to the
 * raw typed arrays across allocate-after-grow boundaries.
 *
 * Determinism: allocate always reuses the lowest free index. Two pools given the
 * same sequence of allocate/free calls observe identical body indices.
 *
 * @author Alex Goldring
 * @copyright Company Named Limited (c) 2026
 */
export class BodyStorage {

    /**
     *
     * @param {number} [initial_capacity]
     */
    constructor(initial_capacity = DEFAULT_INITIAL_CAPACITY) {
        assert.isNonNegativeInteger(initial_capacity, 'initial_capacity');

        const cap = Math.max(MIN_CAPACITY, initial_capacity);

        this.__capacity = cap;
        // High-water mark of slot indices ever issued. Slots in [0, count) are
        // either currently allocated or sitting on the free heap.
        this.__count = 0;

        this.__entities = new Int32Array(cap);
        this.__generations = new Uint8Array(cap);
        this.__kinds = new Uint8Array(cap);
        this.__flags = new Uint32Array(cap);

        // 1 = allocated, 0 = free. Cheaper than scanning the free heap.
        this.__alive = new Uint8Array(cap);

        // Awake set: dense list of awake body indices + reverse map.
        this.__awake_list = new Uint32Array(cap);
        this.__awake_pos = new Int32Array(cap);
        this.__awake_count = 0;

        // Min-heap of free body indices for deterministic reuse.
        this.__free_heap = new Uint32Array(cap);
        this.__free_count = 0;

        // Entity → body-index map (one body per entity). Keeps {@link
        // index_of_entity} O(1) instead of an O(N) scan over the slot table on
        // every collider attach / detach and joint link. Maintained on
        // allocate / free; never grows with the typed arrays (a plain Map).
        this.__entity_to_index = new Map();

        // Initialise reverse map to BODY_INDEX_ABSENT.
        this.__awake_pos.fill(BODY_INDEX_ABSENT);
    }

    /**
     * Currently allocated body count (live, regardless of awake/sleeping).
     * @returns {number}
     */
    get size() {
        return this.__count - this.__free_count;
    }

    /**
     * Total slot capacity. Starts at the (power-of-two) default and grows by
     * doubling; a caller-supplied non-power-of-two initial capacity is honoured
     * as-is, so this is not guaranteed to be a power of two.
     * @returns {number}
     */
    get capacity() {
        return this.__capacity;
    }

    /**
     * Number of awake bodies.
     * @returns {number}
     */
    get awake_count() {
        return this.__awake_count;
    }

    /**
     * High-water mark of slot indices ever issued. Useful when callers maintain
     * per-slot side-tables.
     * @returns {number}
     */
    get high_water_mark() {
        return this.__count;
    }

    /**
     * Allocate a body slot for `entity`. The new body starts in the awake set
     * with default kind {@link BodyKind.Dynamic} and zero flags; the caller
     * may override these via {@link set_kind} / {@link set_flags}.
     *
     * @param {number} entity
     * @returns {number} packed body id
     */
    allocate(entity) {
        let index;

        if (this.__free_count > 0) {
            index = this.__heap_pop();
        } else {
            if (this.__count === this.__capacity) {
                this.__grow();
            }
            index = this.__count++;
        }

        this.__entities[index] = entity;
        this.__kinds[index] = BodyKind.Dynamic;
        this.__flags[index] = 0;
        this.__alive[index] = 1;
        this.__entity_to_index.set(entity, index);

        // Insert into awake set.
        const awake_pos = this.__awake_count++;
        this.__awake_list[awake_pos] = index;
        this.__awake_pos[index] = awake_pos;

        return pack_body_id(index, this.__generations[index]);
    }

    /**
     * Release a body slot previously returned by {@link allocate}. Bumps the
     * generation so any old packed id becomes stale.
     *
     * @param {number} packed_body_id
     */
    free(packed_body_id) {
        const index = body_id_index(packed_body_id);

        assert.equal(this.is_valid(packed_body_id), true, 'free() called on stale or unknown body id');

        // Remove from awake set if present.
        const awake_pos = this.__awake_pos[index];
        if (awake_pos !== BODY_INDEX_ABSENT) {
            this.__awake_remove_at(awake_pos);
            this.__awake_pos[index] = BODY_INDEX_ABSENT;
        }

        this.__alive[index] = 0;

        // Drop the entity → index mapping (the slot still holds the old entity
        // value until reallocation, so delete by it now while it's valid).
        this.__entity_to_index.delete(this.__entities[index]);

        // Bump generation; wraps mod 256.
        this.__generations[index] = (this.__generations[index] + 1) & GENERATION_MASK;

        this.__heap_push(index);
    }

    /**
     *
     * @param {number} packed_body_id
     * @returns {boolean}
     */
    is_valid(packed_body_id) {
        const index = body_id_index(packed_body_id);
        if (index < 0 || index >= this.__count) {
            return false;
        }
        if (this.__alive[index] !== 1) {
            return false;
        }
        return this.__generations[index] === body_id_generation(packed_body_id);
    }

    /**
     * @param {number} index body index (NOT a packed id)
     * @returns {number} entity for the body, or -1 if the slot is free.
     */
    entity_at(index) {
        if (this.__alive[index] !== 1) {
            return -1;
        }
        return this.__entities[index];
    }

    /**
     * Body index for `entity`, or {@link BODY_INDEX_ABSENT} if no live body owns
     * it. O(1) reverse of {@link entity_at} — the lookup callers use on the
     * link / attach / joint paths instead of scanning the slot table.
     * @param {number} entity
     * @returns {number}
     */
    index_of_entity(entity) {
        const idx = this.__entity_to_index.get(entity);
        return idx === undefined ? BODY_INDEX_ABSENT : idx;
    }

    /**
     * @param {number} index
     * @returns {number}
     */
    generation_at(index) {
        return this.__generations[index];
    }

    /**
     * @param {number} index
     * @returns {BodyKind|number}
     */
    kind_at(index) {
        return this.__kinds[index];
    }

    /**
     * @param {number} index
     * @param {BodyKind|number} kind
     */
    set_kind(index, kind) {
        this.__kinds[index] = kind;
    }

    /**
     * @param {number} index
     * @returns {number}
     */
    flags_at(index) {
        return this.__flags[index];
    }

    /**
     * @param {number} index
     * @param {number} flags
     */
    set_flags(index, flags) {
        this.__flags[index] = flags;
    }

    /**
     * Index of `index` in the awake list, or {@link BODY_INDEX_ABSENT} if asleep.
     * @param {number} index
     * @returns {number}
     */
    awake_position_of(index) {
        return this.__awake_pos[index];
    }

    /**
     *
     * @param {number} index
     * @returns {boolean}
     */
    is_awake(index) {
        return this.__awake_pos[index] !== BODY_INDEX_ABSENT;
    }

    /**
     * Read the body index at position `i` in the dense awake list (0-based).
     * @param {number} i
     * @returns {number}
     */
    awake_at(i) {
        return this.__awake_list[i];
    }

    /**
     * Add `index` to the awake set if not already present.
     * @param {number} index
     */
    mark_awake(index) {
        if (this.__awake_pos[index] !== BODY_INDEX_ABSENT) {
            return;
        }
        const pos = this.__awake_count++;
        this.__awake_list[pos] = index;
        this.__awake_pos[index] = pos;
    }

    /**
     * Remove `index` from the awake set if present.
     * @param {number} index
     */
    mark_sleeping(index) {
        const pos = this.__awake_pos[index];
        if (pos === BODY_INDEX_ABSENT) {
            return;
        }
        this.__awake_remove_at(pos);
        this.__awake_pos[index] = BODY_INDEX_ABSENT;
    }

    /**
     * Swap-with-last removal at position `pos` in the dense awake list.
     * @private
     * @param {number} pos
     */
    __awake_remove_at(pos) {
        const last_pos = --this.__awake_count;
        if (pos !== last_pos) {
            const swapped_index = this.__awake_list[last_pos];
            this.__awake_list[pos] = swapped_index;
            this.__awake_pos[swapped_index] = pos;
        }
    }

    /**
     * @private
     */
    __grow() {
        const new_cap = this.__capacity << 1;
        this.__capacity = new_cap;

        const grow_int32 = (a) => {
            const next = new Int32Array(new_cap);
            next.set(a);
            return next;
        };
        const grow_uint8 = (a) => {
            const next = new Uint8Array(new_cap);
            next.set(a);
            return next;
        };
        const grow_uint32 = (a) => {
            const next = new Uint32Array(new_cap);
            next.set(a);
            return next;
        };

        this.__entities = grow_int32(this.__entities);
        this.__generations = grow_uint8(this.__generations);
        this.__kinds = grow_uint8(this.__kinds);
        this.__flags = grow_uint32(this.__flags);
        this.__alive = grow_uint8(this.__alive);
        this.__awake_list = grow_uint32(this.__awake_list);
        this.__free_heap = grow_uint32(this.__free_heap);

        // Awake position has signed sentinel — needs fill on the grown portion.
        const next_pos = new Int32Array(new_cap);
        next_pos.set(this.__awake_pos);
        next_pos.fill(BODY_INDEX_ABSENT, this.__awake_pos.length);
        this.__awake_pos = next_pos;
    }

    // --- Min-heap of free indices --------------------------------------------

    /**
     * @private
     * @param {number} index
     */
    __heap_push(index) {
        const heap = this.__free_heap;
        let i = this.__free_count++;
        heap[i] = index;

        while (i > 0) {
            const parent = (i - 1) >>> 1;
            if (heap[parent] > heap[i]) {
                const tmp = heap[parent];
                heap[parent] = heap[i];
                heap[i] = tmp;
                i = parent;
            } else {
                break;
            }
        }
    }

    /**
     * @private
     * @returns {number}
     */
    __heap_pop() {
        const heap = this.__free_heap;
        const result = heap[0];
        const new_count = --this.__free_count;
        if (new_count > 0) {
            heap[0] = heap[new_count];
            let i = 0;
            while (true) {
                const l = (i << 1) + 1;
                const r = l + 1;
                let smallest = i;
                if (l < new_count && heap[l] < heap[smallest]) smallest = l;
                if (r < new_count && heap[r] < heap[smallest]) smallest = r;
                if (smallest === i) break;
                const tmp = heap[i];
                heap[i] = heap[smallest];
                heap[smallest] = tmp;
                i = smallest;
            }
        }
        return result;
    }
}
