import { assert } from "../../assert.js";
import { RowFirstTable } from "../../collection/table/RowFirstTable.js";
import { RowFirstTableSpec } from "../../collection/table/RowFirstTableSpec.js";
import { ctz32 } from "../ctz32.js";
import { floatToUint, uintToFloatRoundDown, uintToFloatRoundUp } from "../float/SmallFloat.js";
import { msb_32 } from "../msb_32.js";
import { BinaryDataType } from "../type/BinaryDataType.js";

const NUM_TOP_BINS = 32;
const BINS_PER_LEAF = 8;
const TOP_BINS_INDEX_SHIFT = 3;
const LEAF_BINS_INDEX_MASK = 0x7;
const NUM_LEAF_BINS = NUM_TOP_BINS * BINS_PER_LEAF;
export const ALLOCATOR_NO_SPACE = 0xffffffff;
const NODE_UNUSED = 0xffffffff

/**
 * Can potentially use u16, but number of allocation slots will be limited to 65k, so a severe limitation for little gain
 * @type {BinaryDataType}
 */
const NODE_INDEX_DATA_TYPE = BinaryDataType.Uint32;

const DEBUG_VERBOSE = false;

/**
 * @see https://github.com/sebbbi/OffsetAllocator/blob/3610a7377088b1e8c8f1525f458c96038a4e6fc0/offsetAllocator.hpp#L70
 * @type {RowFirstTableSpec}
 */
const NODE_TABLE_SPEC = RowFirstTableSpec.get([
    BinaryDataType.Uint32, // Data Offset
    BinaryDataType.Uint32, // Data Size
    NODE_INDEX_DATA_TYPE, // Bin List Prev
    NODE_INDEX_DATA_TYPE, // Bin List Next
    NODE_INDEX_DATA_TYPE, // Neighbour Prev
    NODE_INDEX_DATA_TYPE, // Neighbour Next
    BinaryDataType.Uint8, // FLAGS
]);

const NODE_FLAGS_USED = 1;


const NODE_FIELD_DATA_OFFSET = 0;
const NODE_FIELD_DATA_SIZE = 1;
const NODE_FIELD_BIN_LIST_PREV = 2;
const NODE_FIELD_BIN_LIST_NEXT = 3;
const NODE_FIELD_NEIGHBOR_PREV = 4;
const NODE_FIELD_NEIGHBOR_NEXT = 5;
const NODE_FIELD_FLAGS = 6;

class Allocation {
    offset = ALLOCATOR_NO_SPACE;
    metadata = ALLOCATOR_NO_SPACE;

    /**
     *
     * @param {number} offset
     * @param {number} metadata
     * @return {Allocation}
     */
    static from(offset, metadata) {
        const r = new Allocation();

        r.offset = offset;
        r.metadata = metadata;

        return r;
    }

    /**
     * @readonly
     * @type {Allocation}
     */
    static INVALID = Object.freeze(Allocation.from(ALLOCATOR_NO_SPACE, ALLOCATOR_NO_SPACE));
}

/**
 *
 * @param {number} bitMask
 * @param {number} startBitIndex
 * @return {number}
 */
function findLowestSetBitAfter(bitMask, startBitIndex) {
    const maskBeforeStartIndex = (1 << startBitIndex) - 1;
    const maskAfterStartIndex = ~maskBeforeStartIndex;
    const bitsAfter = bitMask & maskAfterStartIndex;

    if (bitsAfter === 0) {
        return ALLOCATOR_NO_SPACE;
    }

    return ctz32(bitsAfter);
}

class StorageReport {
    totalFreeSpace = 0
    largestFreeRegion = 0
}

/*
Largely a port of https://github.com/sebbbi/OffsetAllocator by "Sebastian Aaltonen"
 */

/**
 * Fast hard realtime O(1) offset allocator with minimal fragmentation.
 *
 * Uses 256 bins with 8 bit floating point distribution (3 bit mantissa + 5 bit exponent) and a two level bitfield to find the next available bin using 2x LZCNT instructions to make all operations O(1). Bin sizes following the floating point distribution ensures hard bounds for memory overhead percentage regardless of size class. Pow2 bins would waste up to +100% memory (+50% on average). Our float bins waste up to +12.5% (+6.25% on average).
 *
 * The allocation metadata is stored in a separate data structure, making this allocator suitable for sub-allocating any resources, such as GPU heaps, buffers and arrays. Returns an offset to the first element of the allocated contiguous range.
 *
 * @see "A comparison of memory allocators for real-time applications" by Miguel Masmano et al
 * @example
 *      const allocator = new Allocator(12345);     // Allocator with 12345 contiguous elements in total
 *
 *      const a = allocator.allocate(1337);         // Allocate a 1337 element contiguous range
 *      const offset_a = a.offset;                  // Provides offset to the first element of the range
 *      do_something(offset_a);
 *
 *      const b = allocator.allocate(123);          // Allocate a 123 element contiguous range
 *      const offset_b = b.offset;                  // Provides offset to the first element of the range
 *      do_something(offset_b);
 *
 *      allocator.free(a);                          // Free allocation a
 *      allocator.free(b);                          // Free allocation b
 */
export class OffsetAllocator {
    /**
     * Total managed space
     * uint32
     * @type {number}
     */
    size = 0;

    /**
     * Limit on number of allocations
     * uint32
     * @type {number}
     */
    maxAllocs = 0;

    /**
     * uint32
     * @type {number}
     */
    freeStorage = 0;

    /**
     * uint32
     * @type {number}
     */
    usedBinsTop = 0;

    /**
     *
     * @type {Uint8Array}
     */
    usedBins = new Uint8Array(NUM_TOP_BINS);

    /**
     *
     * @type {Uint32Array}
     */
    binIndices = new Uint32Array(NUM_LEAF_BINS);

    /**
     *
     * @type {RowFirstTable}
     */
    nodes = new RowFirstTable(NODE_TABLE_SPEC)

    /**
     * @type {Uint32Array}
     */
    freeNodes;

    /**
     * uint32
     * @type {number}
     */
    freeOffset = 0;

    /**
     *
     * @param {number} size amount of contiguous elements managed by this allocator (typically bytes)
     * @param {number} [maxAllocs]
     */
    constructor(size, maxAllocs = 128 * 1024) {
        assert.isNonNegativeInteger(size, 'size');
        assert.isNonNegativeInteger(maxAllocs, 'maxAllocs');

        this.size = size;
        this.maxAllocs = maxAllocs;

        this.freeNodes = new Uint32Array(this.maxAllocs);

        this.reset();
    }

    /**
     *
     * @param {number} size
     * @returns {Allocation}
     */
    allocate(size) {
        assert.isNonNegativeInteger(size, 'size');

        // Out of allocations?
        if (this.freeOffset === 0) {
            return Allocation.from(ALLOCATOR_NO_SPACE, ALLOCATOR_NO_SPACE);
        }

        // Round up to bin index to ensure that alloc >= bin
        // Gives us min bin index that fits the size
        let minBinIndex = uintToFloatRoundUp(size);

        const minTopBinIndex = minBinIndex >>> TOP_BINS_INDEX_SHIFT;
        const minLeafBinIndex = minBinIndex & LEAF_BINS_INDEX_MASK;

        let topBinIndex = minTopBinIndex;
        let leafBinIndex = ALLOCATOR_NO_SPACE;

        // If top bin exists, scan its leaf bin. This can fail (NO_SPACE).
        if ((this.usedBinsTop & (1 << topBinIndex)) !== 0) {
            leafBinIndex = findLowestSetBitAfter(this.usedBins[topBinIndex], minLeafBinIndex);
        }

        // If we didn't find space in top bin, we search top bin from +1
        if (leafBinIndex === ALLOCATOR_NO_SPACE) {
            topBinIndex = findLowestSetBitAfter(this.usedBinsTop, minTopBinIndex + 1);

            // Out of space?
            if (topBinIndex === ALLOCATOR_NO_SPACE) {
                return Allocation.from(ALLOCATOR_NO_SPACE, ALLOCATOR_NO_SPACE);
            }

            // All leaf bins here fit the alloc, since the top bin was rounded up. Start leaf search from bit 0.
            // NOTE: This search can't fail since at least one leaf bit was set because the top bit was set.
            leafBinIndex = ctz32(this.usedBins[topBinIndex]);
        }

        let bin_index = ((topBinIndex << TOP_BINS_INDEX_SHIFT) | leafBinIndex) >>> 0; // ">>> 0" shift is to force UINT32

        // Pop the top node of the bin. Bin top = node.next.
        let node_index = this.binIndices[bin_index];

        const nodes = this.nodes;

        const nodeTotalSize = nodes.readCellValue(node_index, NODE_FIELD_DATA_SIZE);
        nodes.writeCellValue(node_index, NODE_FIELD_DATA_SIZE, size);
        nodes.writeCellValue(node_index, NODE_FIELD_FLAGS, NODE_FLAGS_USED); //mark as USED

        const node_bin_list_next = nodes.readCellValue(node_index, NODE_FIELD_BIN_LIST_NEXT);
        this.binIndices[bin_index] = node_bin_list_next;

        if (node_bin_list_next !== NODE_UNUSED) {
            nodes.writeCellValue(node_bin_list_next, NODE_FIELD_BIN_LIST_PREV, NODE_UNUSED);
        }

        this.freeStorage -= nodeTotalSize;

        if (DEBUG_VERBOSE) {
            console.log(`Free storage: ${this.freeStorage} (-${nodeTotalSize}) (allocate)`);
        }

        // Bin empty?
        if (this.binIndices[bin_index] === NODE_UNUSED) {
            // Remove a leaf bin mask bit
            this.usedBins[topBinIndex] &= ~(1 << leafBinIndex);

            // All leaf bins empty?
            if (this.usedBins[topBinIndex] === 0) {
                // Remove a top bin mask bit
                this.usedBinsTop &= ~(1 << topBinIndex);
            }
        }

        // Push back reminder N elements to a lower bin
        const reminderSize = nodeTotalSize - size;

        const node_data_offset = nodes.readCellValue(node_index, NODE_FIELD_DATA_OFFSET);

        if (reminderSize > 0) {
            const new_node_index = this.#insertNodeIntoBin(reminderSize, node_data_offset + size);

            const node_neighbor_next = nodes.readCellValue(node_index, NODE_FIELD_NEIGHBOR_NEXT);

            // Link nodes next to each other so that we can merge them later if both are free
            // And update the old next neighbor to point to the new node (in middle)
            if (node_neighbor_next !== NODE_UNUSED) {
                nodes.writeCellValue(node_neighbor_next, NODE_FIELD_NEIGHBOR_PREV, new_node_index);
            }
            nodes.writeCellValue(new_node_index, NODE_FIELD_NEIGHBOR_PREV, node_index);
            nodes.writeCellValue(new_node_index, NODE_FIELD_NEIGHBOR_NEXT, node_neighbor_next);

            nodes.writeCellValue(node_index, NODE_FIELD_NEIGHBOR_NEXT, new_node_index);
        }


        return Allocation.from(node_data_offset, node_index);
    }

    /**
     * Direct method of releasing an allocation.
     * Allows the user to skip holding an object reference in memory.
     * {@link node_index} can be read from {@link Allocation.metadata}
     *
     * @param {number} node_index
     * @see free
     */
    free_node(node_index) {
        assert.isNonNegativeInteger(node_index, 'node_index');
        assert.notEqual(node_index, ALLOCATOR_NO_SPACE, 'Invalid allocation');

        const nodes = this.nodes;
        const node_flags = nodes.readCellValue(node_index, NODE_FIELD_FLAGS);

        // Double delete check
        assert.notEqual(node_flags & NODE_FLAGS_USED, 0, 'Node is not in use');

        // Merge with neighbors...
        let offset = nodes.readCellValue(node_index, NODE_FIELD_DATA_OFFSET)
        let size = nodes.readCellValue(node_index, NODE_FIELD_DATA_SIZE);

        const node_neighbour_prev = nodes.readCellValue(node_index, NODE_FIELD_NEIGHBOR_PREV);

        if (
            (node_neighbour_prev !== NODE_UNUSED)
            && ((nodes.readCellValue(node_neighbour_prev, NODE_FIELD_FLAGS) & NODE_FLAGS_USED) === 0) // not used
        ) {
            // Previous (contiguous) free node: Change offset to previous node offset. Sum sizes

            const previous_offset = nodes.readCellValue(node_neighbour_prev, NODE_FIELD_DATA_OFFSET);
            const previous_size = nodes.readCellValue(node_neighbour_prev, NODE_FIELD_DATA_SIZE);

            offset = previous_offset;
            size += previous_size;

            // Remove node from the bin linked list and put it in the freelist
            this.#removeNodeFromBin(node_neighbour_prev);

            assert.equal(nodes.readCellValue(node_neighbour_prev, NODE_FIELD_NEIGHBOR_NEXT), node_index);

            nodes.writeCellValue(
                node_index, NODE_FIELD_NEIGHBOR_PREV,
                nodes.readCellValue(node_neighbour_prev, NODE_FIELD_NEIGHBOR_PREV)
            );
        }

        const node_neighbour_next = nodes.readCellValue(node_index, NODE_FIELD_NEIGHBOR_NEXT);

        if (
            (node_neighbour_next !== NODE_UNUSED)
            && ((nodes.readCellValue(node_neighbour_next, NODE_FIELD_FLAGS) & NODE_FLAGS_USED) === 0) // not used
        ) {
            // Next (contiguous) free node: Offset remains the same. Sum sizes.

            size += nodes.readCellValue(node_neighbour_next, NODE_FIELD_DATA_SIZE);

            // Remove node from the bin linked list and put it in the freelist
            this.#removeNodeFromBin(node_neighbour_next);

            assert.equal(nodes.readCellValue(node_neighbour_next, NODE_FIELD_NEIGHBOR_PREV), node_index);

            nodes.writeCellValue(
                node_index, NODE_FIELD_NEIGHBOR_NEXT,
                nodes.readCellValue(node_neighbour_next, NODE_FIELD_NEIGHBOR_NEXT)
            );
        }

        const neighborNext = nodes.readCellValue(node_index, NODE_FIELD_NEIGHBOR_NEXT);
        const neighborPrev = nodes.readCellValue(node_index, NODE_FIELD_NEIGHBOR_PREV);

        // Insert the removed node to freelist
        if (DEBUG_VERBOSE) {
            console.log(`Putting node ${node_index} into freelist[${this.freeOffset + 1}] (free)`);
        }

        this.freeNodes[++this.freeOffset] = node_index;

        // Insert the (combined) free node to bin
        const combinedNodeIndex = this.#insertNodeIntoBin(size, offset);

        // Connect neighbors with the new combined node
        if (neighborNext !== NODE_UNUSED) {
            nodes.writeCellValue(combinedNodeIndex, NODE_FIELD_NEIGHBOR_NEXT, neighborNext);
            nodes.writeCellValue(neighborNext, NODE_FIELD_NEIGHBOR_PREV, combinedNodeIndex);
        }
        if (neighborPrev !== NODE_UNUSED) {
            nodes.writeCellValue(combinedNodeIndex, NODE_FIELD_NEIGHBOR_PREV, neighborPrev);
            nodes.writeCellValue(neighborPrev, NODE_FIELD_NEIGHBOR_NEXT, combinedNodeIndex);
        }
    }

    /**
     * @param {Allocation} allocation
     * @see free_node
     */
    free(allocation) {
        assert.notEqual(allocation.metadata, ALLOCATOR_NO_SPACE, 'Invalid allocation');
        // if (!m_nodes) return;

        const node_index = allocation.metadata;

        this.free_node(node_index);
    }

    /**
     *
     * @param {number} size
     * @param {number} data_offset
     * @returns {number}
     */
    #insertNodeIntoBin(size, data_offset) {
        assert.isNonNegativeInteger(size, 'size');
        assert.isNonNegativeInteger(data_offset, 'size');

        // Round down to bin index to ensure that bin >= alloc
        const bin_index = uintToFloatRoundDown(size);

        const top_bin_index = bin_index >>> TOP_BINS_INDEX_SHIFT;
        const leafBinIndex = bin_index & LEAF_BINS_INDEX_MASK;

        // Bin was empty before?
        if (this.binIndices[bin_index] === NODE_UNUSED) {
            // Set bin mask bits
            this.usedBins[top_bin_index] |= 1 << leafBinIndex;
            this.usedBinsTop |= 1 << top_bin_index;
        }

        // Take a freelist node and insert on top of the bin linked list (next = old top)
        const top_node_index = this.binIndices[bin_index];
        const node_index = this.freeNodes[this.freeOffset--];

        if (DEBUG_VERBOSE) {
            console.log(`Getting node ${node_index} from freelist[${this.freeOffset + 1}]`);
        }

        const nodes = this.nodes;

        // initialize node

        nodes.writeCellValue(node_index, NODE_FIELD_DATA_OFFSET, data_offset);
        nodes.writeCellValue(node_index, NODE_FIELD_DATA_SIZE, size);

        nodes.writeCellValue(node_index, NODE_FIELD_BIN_LIST_PREV, NODE_UNUSED);
        nodes.writeCellValue(node_index, NODE_FIELD_BIN_LIST_NEXT, top_node_index);

        nodes.writeCellValue(node_index, NODE_FIELD_NEIGHBOR_PREV, NODE_UNUSED);
        nodes.writeCellValue(node_index, NODE_FIELD_NEIGHBOR_NEXT, NODE_UNUSED);

        nodes.writeCellValue(node_index, NODE_FIELD_FLAGS, 0);

        if (top_node_index !== NODE_UNUSED) {
            nodes.writeCellValue(top_node_index, NODE_FIELD_BIN_LIST_PREV, node_index);
        }

        this.binIndices[bin_index] = node_index;

        this.freeStorage += size;

        if (DEBUG_VERBOSE) {
            console.log(`Free storage: ${this.freeStorage} (+${size}) (insertNodeIntoBin)`);
        }

        return node_index;
    }

    /**
     *
     * @param {number} node_index
     */
    #removeNodeFromBin(node_index) {
        assert.isNonNegativeInteger(node_index, 'node_index');

        const nodes = this.nodes;
        const node_bin_list_prev = nodes.readCellValue(node_index, NODE_FIELD_BIN_LIST_PREV);
        if (node_bin_list_prev !== NODE_UNUSED) {
            // Easy case: We have previous node. Just remove this node from the middle of the list.
            const node_bin_list_next = nodes.readCellValue(node_index, NODE_FIELD_BIN_LIST_NEXT);
            nodes.writeCellValue(
                node_bin_list_prev, NODE_FIELD_BIN_LIST_NEXT,
                node_bin_list_next
            );

            if (node_bin_list_next !== NODE_UNUSED) {
                nodes.writeCellValue(
                    node_bin_list_next, NODE_FIELD_BIN_LIST_PREV,
                    node_bin_list_prev
                )
            }
        } else {
            // Hard case: We are the first node in a bin. Find the bin.

            // Round down to bin index to ensure that bin >= alloc
            const node_data_size = nodes.readCellValue(node_index, NODE_FIELD_DATA_SIZE);
            const binIndex = uintToFloatRoundDown(node_data_size);

            const topBinIndex = binIndex >> TOP_BINS_INDEX_SHIFT;
            const leafBinIndex = binIndex & LEAF_BINS_INDEX_MASK;

            const node_bin_list_next = nodes.readCellValue(node_index, NODE_FIELD_BIN_LIST_NEXT);
            this.binIndices[binIndex] = node_bin_list_next;
            if (node_bin_list_next !== NODE_UNUSED) {
                nodes.writeCellValue(
                    node_bin_list_next, NODE_FIELD_BIN_LIST_PREV,
                    NODE_UNUSED
                );
            }

            // Bin empty?
            if (this.binIndices[binIndex] === NODE_UNUSED) {
                // Remove a leaf bin mask bit
                this.usedBins[topBinIndex] &= ~(1 << leafBinIndex);

                // All leaf bins empty?
                if (this.usedBins[topBinIndex] === 0) {
                    // Remove a top bin mask bit
                    this.usedBinsTop &= ~(1 << topBinIndex);
                }
            }
        }

        // Insert the node to freelist
        if (DEBUG_VERBOSE) {
            console.log(`Putting node ${node_index} into freelist[${this.freeOffset + 1}] (removeNodeFromBin)`);
        }

        this.freeNodes[++this.freeOffset] = node_index;

        const node_data_size = nodes.readCellValue(node_index, NODE_FIELD_DATA_SIZE);
        this.freeStorage -= node_data_size;

        if (DEBUG_VERBOSE) {
            console.log(`Free storage: ${this.freeStorage} (-${node_data_size}) (removeNodeFromBin)`);
        }
    }

    /**
     * Drop all allocations, effectively deallocating everything and restoring allocator into initial state
     */
    reset() {
        this.freeStorage = 0;
        this.usedBinsTop = 0;

        const maxAllocs = this.maxAllocs;

        this.freeOffset = maxAllocs - 1;

        for (let i = 0; i < NUM_TOP_BINS; i++) {
            this.usedBins[i] = 0;
        }
        for (let i = 0; i < NUM_LEAF_BINS; i++) {
            this.binIndices[i] = NODE_UNUSED;
        }

        // drop data from the node table
        this.nodes.length = 0;
        this.nodes.setCapacity(maxAllocs);

        // Freelist is a stack. Nodes in inverse order so that [0] pops first.
        for (let i = 0; i < maxAllocs; i++) {
            this.freeNodes[i] = maxAllocs - i - 1;
        }

        // Start state: Whole storage as one big node
        // Algorithm will split remainders and push them back as smaller nodes
        this.#insertNodeIntoBin(this.size, 0);

    }

    /**
     * Useful for debug and UI
     * @return {StorageReport}
     */
    storageReport() {
        let largestFreeRegion = 0;
        let freeStorage = 0;

        // Out of allocations? -> Zero free space
        if (this.freeOffset > 0) {
            freeStorage = this.freeStorage;

            if (this.usedBinsTop !== 0) {
                const topBinIndex = msb_32(this.usedBinsTop);
                const leafBinIndex = msb_32(this.usedBins[topBinIndex]);

                largestFreeRegion = floatToUint((topBinIndex << TOP_BINS_INDEX_SHIFT) | leafBinIndex);

                assert.greaterThanOrEqual(freeStorage, largestFreeRegion);
            }

        }

        const report = new StorageReport();

        report.largestFreeRegion = largestFreeRegion;
        report.totalFreeSpace = freeStorage;

        return report;
    }
}