import { Dot } from "../../crdt/Dot.js";
import TrieBranch from "./TrieBranch.js";
import TrieLeaf from "./TrieLeaf.js";
import { compareBytes, pendingChildOid } from "./trieCursorHelpers.js";
export default class TrieCompactor {
    #geometry;
    #nibbleBits;
    #loadRootIfNeeded;
    #hasRoot;
    #leafAt;
    #branchAt;
    #ensureChildLoaded;
    #markLeafDirty;
    #markBranchDirty;
    #clearLeafAt;
    #clearBranchAt;
    constructor(init) {
        this.#geometry = init.geometry;
        this.#nibbleBits = init.nibbleBits;
        this.#loadRootIfNeeded = init.loadRootIfNeeded;
        this.#hasRoot = init.hasRoot;
        this.#leafAt = init.leafAt;
        this.#branchAt = init.branchAt;
        this.#ensureChildLoaded = init.ensureChildLoaded;
        this.#markLeafDirty = init.markLeafDirty;
        this.#markBranchDirty = init.markBranchDirty;
        this.#clearLeafAt = init.clearLeafAt;
        this.#clearBranchAt = init.clearBranchAt;
        Object.freeze(this);
    }
    async compact(includedVV) {
        await this.#loadRootIfNeeded();
        if (!this.#hasRoot()) {
            return;
        }
        await this.#compactBelow([], includedVV);
    }
    async #compactBelow(path, includedVV) {
        const leaf = this.#leafAt(path);
        if (leaf !== null) {
            return this.#compactLeaf(path, leaf, includedVV);
        }
        const branch = this.#branchAt(path);
        if (branch === null) {
            return { kind: "absent", changed: false };
        }
        return await this.#compactBranch(path, branch, includedVV);
    }
    #compactLeaf(path, leaf, includedVV) {
        const result = compactLeafEntries(leaf.entries(), includedVV);
        if (!result.changed) {
            return { kind: "leaf", changed: false, leaf };
        }
        if (result.entries.length === 0) {
            this.#clearLeafAt(path);
            return { kind: "absent", changed: true };
        }
        const nextLeaf = new TrieLeaf(result.entries, this.#geometry);
        this.#markLeafDirty(path, nextLeaf);
        return { kind: "leaf", changed: true, leaf: nextLeaf };
    }
    async #compactBranch(path, branch, includedVV) {
        const childStates = new Map();
        for (const [childNibble, childOid] of branch.entries()) {
            await this.#ensureChildLoaded(path, childNibble, childOid);
            childStates.set(childNibble, await this.#compactBelow([...path, childNibble], includedVV));
        }
        let changed = false;
        for (const state of childStates.values()) {
            if (state.changed || state.kind === "absent") {
                changed = true;
            }
        }
        if (this.#mergeSiblingLeaves(path, childStates)) {
            changed = true;
        }
        return this.#rewriteBranch(path, branch, childStates, changed);
    }
    #mergeSiblingLeaves(path, childStates) {
        let merged = false;
        while (true) {
            const candidate = this.#findMergeCandidate(childStates);
            if (candidate === null) {
                return merged;
            }
            this.#markLeafDirty([...path, candidate.keepNibble], candidate.mergedLeaf);
            this.#clearLeafAt([...path, candidate.dropNibble]);
            childStates.set(candidate.keepNibble, {
                kind: "leaf",
                changed: true,
                leaf: candidate.mergedLeaf,
            });
            childStates.delete(candidate.dropNibble);
            merged = true;
        }
    }
    #findMergeCandidate(childStates) {
        const nibbles = [...childStates.keys()].sort((a, b) => a - b);
        for (const nibble of nibbles) {
            const state = childStates.get(nibble);
            if (state?.kind !== "leaf") {
                continue;
            }
            if (!state.leaf.requiresMerge(this.#geometry)) {
                continue;
            }
            for (const siblingNibble of nibbles) {
                if (siblingNibble === nibble) {
                    continue;
                }
                const sibling = childStates.get(siblingNibble);
                if (sibling?.kind !== "leaf") {
                    continue;
                }
                const combinedSize = state.leaf.size() + sibling.leaf.size();
                if (combinedSize > this.#geometry.leafCapacity) {
                    continue;
                }
                const keepNibble = Math.min(nibble, siblingNibble);
                const dropNibble = Math.max(nibble, siblingNibble);
                const leftLeaf = childStates.get(keepNibble);
                const rightLeaf = childStates.get(dropNibble);
                if (leftLeaf?.kind !== "leaf" || rightLeaf?.kind !== "leaf") {
                    continue;
                }
                return {
                    keepNibble,
                    dropNibble,
                    mergedLeaf: mergeLeaves(leftLeaf.leaf, rightLeaf.leaf, this.#geometry),
                };
            }
        }
        return null;
    }
    #rewriteBranch(path, branch, childStates, childChanged) {
        const nextChildren = new Map();
        let changed = childChanged;
        for (const [childNibble, childOid] of branch.entries()) {
            const state = childStates.get(childNibble);
            if (state === undefined || state.kind === "absent") {
                changed = true;
                continue;
            }
            if (state.changed) {
                nextChildren.set(childNibble, pendingChildOid(path, childNibble));
                changed = true;
                continue;
            }
            nextChildren.set(childNibble, childOid);
        }
        if (nextChildren.size === 0) {
            if (path.length === 0) {
                const emptyRoot = new TrieBranch(new Map(), this.#geometry);
                this.#markBranchDirty([], emptyRoot);
                return { kind: "branch", changed: true, branch: emptyRoot };
            }
            this.#clearBranchAt(path);
            return { kind: "absent", changed: true };
        }
        if (nextChildren.size === 1 && path.length > 0) {
            const [onlyNibble] = nextChildren.keys();
            if (onlyNibble === undefined) {
                return { kind: "branch", changed, branch };
            }
            const state = childStates.get(onlyNibble);
            if (state?.kind === "leaf") {
                this.#clearLeafAt([...path, onlyNibble]);
                this.#clearBranchAt(path);
                const collapsedLeaf = collapseLeafUp(state.leaf, onlyNibble, path.length + 1, this.#nibbleBits, this.#geometry);
                this.#markLeafDirty(path, collapsedLeaf);
                return { kind: "leaf", changed: true, leaf: collapsedLeaf };
            }
        }
        if (!changed) {
            return { kind: "branch", changed: false, branch };
        }
        const nextBranch = new TrieBranch(nextChildren, this.#geometry);
        this.#markBranchDirty(path, nextBranch);
        return { kind: "branch", changed: true, branch: nextBranch };
    }
}
function compactLeafEntries(entries, includedVV) {
    let changed = false;
    const compacted = [];
    for (const entry of entries) {
        const nextLive = new Set();
        const nextTombstoned = new Set();
        let entryChanged = false;
        for (const dot of entry.dots) {
            if (entry.tombstonedDots.has(dot) && includedVV.contains(Dot.decode(dot))) {
                entryChanged = true;
                continue;
            }
            nextLive.add(dot);
        }
        for (const dot of entry.tombstonedDots) {
            if (includedVV.contains(Dot.decode(dot))) {
                entryChanged = true;
                continue;
            }
            nextTombstoned.add(dot);
        }
        if (nextLive.size === 0 && nextTombstoned.size === 0) {
            changed = true;
            continue;
        }
        changed = changed || entryChanged;
        compacted.push({
            routeKeySuffix: entry.routeKeySuffix,
            element: entry.element,
            dots: nextLive,
            tombstonedDots: nextTombstoned,
        });
    }
    return { changed, entries: compacted };
}
function mergeLeaves(left, right, geometry) {
    const merged = [...left.entries(), ...right.entries()];
    merged.sort((a, b) => compareBytes(a.routeKeySuffix, b.routeKeySuffix));
    return new TrieLeaf(merged, geometry);
}
function collapseLeafUp(leaf, childNibble, childDepth, nibbleBits, geometry) {
    const lengthened = leaf.entries().map((entry) => ({
        routeKeySuffix: prependNibbleToSuffix({
            suffix: entry.routeKeySuffix,
            nibble: childNibble,
            depth: childDepth,
            nibbleBits,
        }),
        element: entry.element,
        dots: entry.dots,
        tombstonedDots: entry.tombstonedDots,
    }));
    lengthened.sort((a, b) => compareBytes(a.routeKeySuffix, b.routeKeySuffix));
    return new TrieLeaf(lengthened, geometry);
}
function prependNibbleToSuffix(args) {
    const oldSlots = Math.floor((256 - args.depth * args.nibbleBits) / args.nibbleBits);
    const totalSlots = oldSlots + 1;
    const totalBits = totalSlots * args.nibbleBits;
    const out = new Uint8Array(Math.ceil(totalBits / 8));
    writeSlot(out, 0, args.nibbleBits, args.nibble);
    for (let i = 0; i < oldSlots; i += 1) {
        writeSlot(out, i + 1, args.nibbleBits, readSlot(args.suffix, i, args.nibbleBits));
    }
    return out;
}
function readSlot(suffix, slot, nibbleBits) {
    const bitOffset = slot * nibbleBits;
    let value = 0;
    for (let bitIndex = 0; bitIndex < nibbleBits; bitIndex += 1) {
        value = (value << 1) | readBitMsbFirst(suffix, bitOffset + bitIndex);
    }
    return value;
}
function writeSlot(out, slot, nibbleBits, value) {
    const bitOffset = slot * nibbleBits;
    for (let bitIndex = 0; bitIndex < nibbleBits; bitIndex += 1) {
        const bit = (value >>> (nibbleBits - bitIndex - 1)) & 1;
        writeBitMsbFirst(out, bitOffset + bitIndex, bit);
    }
}
function writeBitMsbFirst(out, bitOffset, bit) {
    const byteIndex = Math.floor(bitOffset / 8);
    const byte = out[byteIndex] ?? 0;
    const bitInByte = bitOffset % 8;
    const mask = 1 << (7 - bitInByte);
    out[byteIndex] = bit === 1 ? byte | mask : byte & ~mask;
}
function readBitMsbFirst(bytes, bitOffset) {
    const byteIndex = Math.floor(bitOffset / 8);
    const byte = bytes[byteIndex] ?? 0;
    const bitInByte = bitOffset % 8;
    return (byte >>> (7 - bitInByte)) & 1;
}
