import { icons as LucideIcons } from "lucide-svelte";
import { untrack } from "svelte";
import { SvelteMap, SvelteSet } from "svelte/reactivity";
import type { FoldKey, TocBlockItem } from "../../../types";

export interface ToolbarState {
    foldsByKey: SvelteMap<FoldKey, boolean>;
    totalFolds: number;
    readonly tocCollapsed: boolean;
    readonly collapsedCount: number;
    readonly hasAnyFolds: boolean;
    readonly allFoldsCollapsed: boolean;
}

class SvMap<K, V> extends SvelteMap<K, V> {
    get(key: K): V | undefined;
    get(key: K, fallback: V): V;
    override get(key: K, fallback?: V): V | undefined {
        let value = super.get(key);
        if (value === undefined && fallback !== undefined) value = fallback;
        return value;
    }
}

function collectFoldEntries(items: TocBlockItem[]): Array<[FoldKey, boolean]> {
    const entries: Array<[FoldKey, boolean]> = [];

    for (const item of items) {
        if (item.children.length > 0 && item.foldKey) {
            entries.push([ item.foldKey, item.initialCollapsed ]);
        }

        entries.push(...collectFoldEntries(item.children));
    }

    return entries;
}

type PersistedFoldStateCb = (key: FoldKey, collapsed: boolean) => void;
type PersistTocCollapsedStateCb = (value: boolean) => void;

export class TocToolbarState implements ToolbarState {
    public foldsByKey = new SvMap<FoldKey, boolean>();
    public tocCollapsed = $state<boolean>(false);

    public readonly totalFolds = $derived<number>(this.foldsByKey.size);
    public readonly collapsedCount = $derived<number>([ ...this.foldsByKey.values() ].filter(Boolean).length);
    public readonly hasAnyFolds = $derived<boolean>(this.totalFolds > 0);
    public readonly allFoldsCollapsed = $derived<boolean>(
        this.totalFolds > 0 && this.collapsedCount === this.totalFolds
    );
    public readonly icons: typeof LucideIcons = LucideIcons;

    constructor(
        private persistedFoldState: PersistedFoldStateCb,
        private persistTocCollapsed: PersistTocCollapsedStateCb
    ) {}

    public isCollapsed(key: FoldKey): boolean {
        return this.foldsByKey.get(key, false);
    }

    public setCollapsed(key: FoldKey, value: boolean): void {
        if (this.foldsByKey.get(key, false) === value) return;

        this.foldsByKey.set(key, value);
        this.persistedFoldState(key, value);
    }

    public toggleFold(key: FoldKey): void {
        this.setCollapsed(key, !this.isCollapsed(key));
    }

    public setAllFolds(value: boolean): void {
        for (const key of this.foldsByKey.keys()) {
            this.setCollapsed(key, value);
        }
    }

    public setTocCollapsed(value: boolean): void {
        if (this.tocCollapsed === value) return;

        this.tocCollapsed = value;
        this.persistTocCollapsed(value);
    }

    public syncTocCollapsed(value: boolean): void {
        if (this.tocCollapsed === value) return;

        this.tocCollapsed = value;
    }

    public syncFoldsFromItems(items: TocBlockItem[]): void {
        const existing = untrack(() => new SvMap(this.foldsByKey));
        const nextEntries = collectFoldEntries(items);
        const nextKeys = new SvelteSet<FoldKey>();

        for (const [ key, initialCollapsed ] of nextEntries) {
            nextKeys.add(key);

            const nextValue = existing.get(key, initialCollapsed);

            if (!existing.has(key) || existing.get(key) !== nextValue) {
                this.foldsByKey.set(key, nextValue);
            }
        }

        for (const key of existing.keys()) {
            if (!nextKeys.has(key)) {
                this.foldsByKey.delete(key);
            }
        }
    }
}