import { expect, test } from "bun:test";
import {
    NODE_BYTES,
    Queue,
    VIEW_HIGH,
    VIEW_LOW,
    before,
    blocks,
    bytes,
    clean,
    compressTask,
    contains,
    grow,
    key,
    merge,
    mergeTask,
    pieces,
    render,
    tooLong,
    viewBytes,
    type Line,
    type View,
} from "../src/optchat-core.js";

interface State {
    keep: number;
    state: number;
    older: State | null;
}

// Taelin's rollback push without `life`, which stays 0 under push alone.
function push(next: number, states: State | null): State {
    if (states === null) return { keep: 0, state: next, older: null };
    if (states.keep === 0) return { ...states, keep: 1 };
    return { keep: 0, state: next, older: push(states.state, states.older) };
}

// Reads the newest-first list as view lines, oldest first.
function lines(states: State | null, t: number) {
    const starts: number[] = [];
    for (let s = states; s; s = s.older) starts.push(s.state);
    starts.reverse();
    return starts.map((first, k): Line => {
        const size = (starts[k + 1] ?? t + 1) - first;
        const l = Math.log2(size);
        expect(Number.isInteger(l) && first % size === 0).toBe(true);
        return [l, first / size];
    });
}

test("due ordering reproduces Taelin's push at its list length for t=0..2000", () => {
    // Every parent counts as built. A target one byte under the current size merges exactly the
    // most due pair, so the budget is push's list length in lines.
    const text = () => "";
    let states: State | null = null;
    let view: Line[] = [];
    for (let t = 0; t <= 2000; t++) {
        states = push(t, states);
        const expected = lines(states, t);
        view.push([0, t]);
        while (view.length > expected.length) {
            const next = merge(view, t + 1, viewBytes(view, text) - 1, text);
            expect(next.length).toBe(view.length - 1);
            view = next;
        }
        expect(view).toEqual(expected);
    }
});

test("merging from the first message diverges from push (the recipe's T=10 example)", () => {
    const text = () => "";
    // View 0+4, 4+4, 8+1, 9+1 at T=10: push merges 8-9, not 0-7.
    const view: Line[] = [
        [2, 0],
        [2, 1],
        [0, 8],
        [0, 9],
    ];
    expect(merge(view, 10, viewBytes(view, text) - 1, text)).toEqual([
        [2, 0],
        [2, 1],
        [1, 4],
    ]);
});

test("merge skips pairs whose parent is unbuilt", () => {
    const built = new Set(["1:0"]);
    const text = (line: Line) => (line[0] === 0 || built.has(key(line)) ? "" : undefined);
    const view: Line[] = [
        [0, 0],
        [0, 1],
        [0, 2],
        [0, 3],
    ];
    expect(merge(view, 4, 0, text)).toEqual([
        [1, 0],
        [0, 2],
        [0, 3],
    ]);
});

test("view is a sawtooth: grows past 128 KB, then one batch drops it to 64 KB", () => {
    const line = "y".repeat(500);
    const text = () => line;
    const view: View = { lines: [], batching: false };
    let peak = 0;
    let batches = 0;
    let total = 0;
    for (let i = 0; i < 2000; i++) {
        const size = viewBytes(view.lines, text);
        expect(total).toBe(size);
        peak = Math.max(peak, size);
        const grown = grow(view, i, text, total, true);
        total = grown.total;
        if (grown.merged) {
            batches++;
            expect(size).toBeLessThanOrEqual(VIEW_HIGH);
            expect(viewBytes(view.lines, text)).toBeLessThanOrEqual(VIEW_LOW);
        } else expect(view.lines.at(-1)).toEqual([0, i]);
    }
    expect(peak).toBeGreaterThan(VIEW_LOW);
    expect(peak).toBeLessThanOrEqual(VIEW_HIGH + 600);
    expect(batches).toBeGreaterThan(1);
    // Lines stay contiguous and cover every message.
    let next = 0;
    for (const [l, i] of view.lines) {
        expect(i * 2 ** l).toBe(next);
        next += 2 ** l;
    }
    expect(next).toBe(2000);
});

test("a batch that cannot reach 64 KB continues at the next message", () => {
    let parents = false;
    const text = (line: Line) => (line[0] === 0 || parents ? "z".repeat(500) : undefined);
    const view: View = { lines: [], batching: false };
    let total = 0;
    for (let i = 0; i < 260; i++) {
        const grown = grow(view, i, text, total, true);
        expect(grown.merged).toBe(false);
        total = grown.total;
    }
    expect(view.batching).toBe(true);
    parents = true;
    // Nothing was built since the last try: the batch waits.
    const waiting = grow(view, 260, text, total, false);
    expect(waiting).toEqual({ merged: false, tried: false, total: viewBytes(view.lines, text) });
    expect(grow(view, 261, text, waiting.total, true).merged).toBe(true);
    expect(viewBytes(view.lines, text)).toBeLessThanOrEqual(VIEW_LOW);
    expect(view.batching).toBe(false);
});

test("compaction context stops before the target and at the first unbuilt line", () => {
    const text = (line: Line) => (key(line) === "0:5" ? undefined : "t");
    const view: Line[] = [
        [2, 0],
        [0, 4],
        [0, 5],
        [0, 6],
    ];
    expect(before(view, 4, text)).toEqual([[2, 0]]);
    expect(before(view, 7, text)).toEqual([
        [2, 0],
        [0, 4],
    ]);
});

test("view blocks keep whole 4-line blocks stable as lines are appended", () => {
    const a = blocks(["1", "2", "3", "4", "5", "6"]);
    expect(a.whole).toEqual(["<chat>\n1\n2\n3\n4\n"]);
    expect(a.tail).toBe("5\n6\n</chat>");
    const b = blocks(["1", "2", "3", "4", "5", "6", "7", "8", "9"]);
    expect(b.whole[0]).toBe(a.whole[0]!);
    expect(b.tail).toBe("9\n</chat>");
    expect(blocks([])).toEqual({ whole: [], tail: "<chat>\n</chat>" });
});

test("lines render with labels and without newlines", () => {
    expect(render([3, 5], "a\nb")).toBe("40+8|a b");
    expect(render([0, 7], undefined)).toBe("7+1|(not summarized yet: zoom it)");
});

test("long tool outputs are clipped; other long texts are split", () => {
    const long = "a".repeat(20_000) + "b".repeat(20_000);
    const [echo] = pieces("echo", long);
    expect(echo!.startsWith("a".repeat(15_000))).toBe(true);
    expect(echo!.endsWith("b".repeat(15_000))).toBe(true);
    expect(echo!).toContain("10000 characters clipped");
    expect(pieces("agent", long).join("")).toBe(long);
    expect(pieces("agent", long)).toHaveLength(2);
    expect(pieces("user", "short")).toEqual(["short"]);
});

test("tasks are verbatim and Too long cuts at 512 bytes", () => {
    const ruler = "-".repeat(512);
    expect(compressTask(7, "user", "hi")).toBe(
        `Compaction: compress message 7 into one line of at most 512 bytes\n(about 70 words), the length of this ruler:\n${ruler}\n<input>\nuser: hi\n</input>`,
    );
    expect(mergeTask([3, 5], "A", "B")).toBe(
        `Compaction: merge lines 40+4 and 44+4, adjacent, into one line of at most\n512 bytes (about 70 words), the length of this ruler:\n${ruler}\n<chat> may hold their messages, 40 to 47, in more detail: take details\nof them from there too.\n<input>\nA\nB\n</input>`,
    );
    const reply = "é".repeat(300);
    const message = tooLong(reply);
    expect(message).toStartWith("Too long: your line is 600 bytes, over the 512-byte limit.");
    expect(message).toEndWith(`\n${"é".repeat(256)}| ← LIMIT`);
    expect(bytes("é".repeat(256))).toBe(NODE_BYTES);
    expect(clean(" 40+8| the line \n")).toBe("the line");
});

test("queue keeps at most 8 builds running, starts messages within 8 of the first unbuilt, and merges built pairs", async () => {
    const built = new Set<string>();
    const pending = new Map<string, (ok: boolean) => void>();
    let peak = 0;
    const queue = new Queue(
        (line) => built.has(key(line)),
        (line) =>
            new Promise<boolean>((resolve) => {
                pending.set(key(line), (ok) => {
                    if (ok) built.add(key(line));
                    resolve(ok);
                });
                peak = Math.max(peak, pending.size);
            }),
    );
    const finish = async (name: string, ok = true) => {
        const resolve = pending.get(name)!;
        pending.delete(name);
        resolve(ok);
        await Bun.sleep(0);
    };
    for (let i = 0; i < 12; i++) queue.add(i);
    expect([...pending.keys()].sort()).toEqual(
        ["0:0", "0:1", "0:2", "0:3", "0:4", "0:5", "0:6", "0:7"].sort(),
    );
    await finish("0:1");
    // Message 8 starts: 7 earlier messages remain unbuilt.
    expect(pending.has("0:8")).toBe(true);
    await finish("0:0");
    // 0 and 1 are built: their merge competes with message 9 for the free slot; messages come first.
    expect(pending.has("0:9")).toBe(true);
    expect(pending.has("1:0")).toBe(false);
    await finish("0:2", false);
    // The failure keeps its slot as unbuilt; no leaf may start, so the merge runs.
    expect(pending.has("1:0")).toBe(true);
    expect(pending.has("0:10")).toBe(false);
    expect(peak).toBeLessThanOrEqual(8);
    const settled = queue.settle(3, 1_000);
    // settle requeues the failure at once; it starts in the next free slot.
    await finish("1:0");
    expect(pending.has("0:2")).toBe(true);
    await finish("0:2");
    expect(await settled).toBe(true);
});

test("queue resumes unbuilt messages and ready merges from disk", async () => {
    const built = new Set(["0:0", "0:1", "0:2"]);
    const started: string[] = [];
    const queue = new Queue(
        (line) => built.has(key(line)),
        async (line) => {
            started.push(key(line));
            built.add(key(line));
            return true;
        },
    );
    queue.resume(4, [
        [0, 0],
        [0, 1],
        [0, 2],
    ]);
    await Bun.sleep(0);
    expect(started).toEqual(["0:3", "1:0", "1:1", "2:0"]);
});

test("a long unbuilt view grows in linear time and pending batches wait for builds", () => {
    // Regression: enabling optchat in a long session re-measured and re-merged the whole view per message.
    const text = () => undefined;
    const view: View = { lines: [], batching: false };
    let total = 0;
    let built = true;
    const started = performance.now();
    for (let i = 0; i < 20_000; i++) {
        const grown = grow(view, i, text, total, built);
        total = grown.total;
        if (grown.tried) built = false;
    }
    expect(total).toBe(viewBytes(view.lines, text));
    expect(view.batching).toBe(true);
    expect(performance.now() - started).toBeLessThan(2_000);
});

test("contains finds exactly the view's lines", () => {
    const view: Line[] = [
        [2, 0],
        [1, 2],
        [0, 6],
    ];
    expect(contains(view, [1, 2])).toBe(true);
    expect(contains(view, [0, 6])).toBe(true);
    expect(contains(view, [0, 4])).toBe(false);
    expect(contains(view, [0, 7])).toBe(false);
});

test("a merge is queued once when both halves finish in the same tick", async () => {
    // Regression: short leaves build synchronously, so both finishes saw the parent unbuilt.
    const tree = new Set<string>();
    const starts: string[] = [];
    const queue = new Queue(
        (line) => tree.has(key(line)),
        async (line) => {
            starts.push(key(line));
            if (line[0] > 0) await Bun.sleep(5);
            tree.add(key(line));
            return true;
        },
    );
    queue.add(0);
    queue.add(1);
    await Bun.sleep(30);
    expect(starts).toEqual(["0:0", "0:1", "1:0"]);
});

test("settle stops waiting once failures leave nothing running", async () => {
    // Regression: it waited the full timeout though failures retry only at the next message.
    const queue = new Queue(
        () => false,
        async () => {
            await Bun.sleep(5);
            return false;
        },
    );
    queue.add(0);
    const started = Date.now();
    expect(await queue.settle(1, 2_000)).toBe(false);
    expect(Date.now() - started).toBeLessThan(500);
    expect(queue.running).toBe(0);
});