import { afterEach, expect, test } from "bun:test";
import { mkdtemp, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { Plugin } from "@opencode/plugin";
import optchat from "../src/optchat.js";

const directories: string[] = [];
afterEach(async () => {
    await Promise.all(
        directories.splice(0).map((directory) => rm(directory, { recursive: true, force: true })),
    );
});

interface Part {
    type: string;
    text?: string;
    cache?: { type: string; ttlSeconds?: number };
    name?: string;
    input?: object;
    id?: string;
    providerExecuted?: boolean;
    result?: { type: string; value: string };
}
interface Message {
    id?: string;
    role: string;
    content: Part[];
}
// Stands in for the host's Message class, which the plugin reuses through `constructor`.
class Entry implements Message {
    id?: string;
    role = "";
    content: Part[] = [];
    constructor(input: Message) {
        Object.assign(this, input);
    }
}
interface Event {
    sessionID: string;
    agent: string;
    system: Part[];
    messages: Entry[];
    tools: Record<string, string>;
    result?: { summary: string };
}
interface ZoomInput {
    id: number;
    n?: number;
}
interface Tool {
    name: string;
    options: { codemode: boolean };
    execute: (
        input: ZoomInput,
        context: { sessionID: string; agent: string },
    ) => Promise<{ content: string }>;
}

interface Gate {
    promise?: Promise<void>;
}

async function host() {
    const directory = await mkdtemp(join(tmpdir(), "optchat-plugin-"));
    directories.push(directory);
    const hooks = new Map<string, (event: Event) => Promise<void> | void>();
    const storage = new Map<string, string | boolean>();
    const notices: string[] = [];
    const tools = new Map<string, Tool>();
    const generated: { system: Part[]; messages: Entry[] }[] = [];
    const replies: string[] = [];
    // While set, compactions wait here before the host runs the generate hook.
    const gate: Gate = {};
    let command:
        | ((input: { sessionID: string; prompt: { text: string } }) => Promise<void>)
        | undefined;
    const context = {
        options: { dataDir: directory, compactionModel: "acme/cheap#xhigh" },
        storage: {
            get: async (key: string) => storage.get(key),
            set: async (key: string, value: string | boolean) => void storage.set(key, value),
            remove: async (key: string) => void storage.delete(key),
        },
        tool: {
            transform: async (callback: (editor: { add: (tool: Tool) => void }) => void) =>
                callback({ add: (tool) => tools.set(tool.name, tool) }),
        },
        command: {
            transform: async (
                callback: (editor: {
                    add: (definition: { execute: typeof command }) => void;
                }) => void,
            ) =>
                callback({
                    add: (definition) => (command = definition.execute),
                }),
        },
        session: {
            hook: async (name: string, callback: (event: Event) => Promise<void> | void) => {
                hooks.set(name, callback);
                return { dispose: async () => {} };
            },
            get: async ({ sessionID }: { sessionID: string }) =>
                sessionID === "ses_child"
                    ? { id: sessionID, parentID: "ses_root" }
                    : sessionID === "ses_worker"
                      ? { id: sessionID, parentID: "ses_root" }
                      : { id: sessionID, agent: "build" },
            create: async (input: { parentID: string; model: object }) => {
                expect(input.parentID).toBe("ses_root");
                expect(input.model).toEqual({
                    providerID: "acme",
                    id: "cheap",
                    variant: "xhigh",
                });
                return { id: "ses_worker" };
            },
            generate: async ({ sessionID, prompt }: { sessionID: string; prompt: string }) => {
                await gate.promise;
                const event: Event = {
                    sessionID,
                    agent: "build",
                    system: [{ type: "text", text: "host system" }],
                    messages: [
                        new Entry({
                            role: "user",
                            content: [{ type: "text", text: prompt }],
                        }),
                    ],
                    tools: {},
                };
                await hooks.get("generate")!(event);
                generated.push(event);
                return { text: replies.shift() ?? "short summary" };
            },
            synthetic: async ({ text }: { text: string }) => void notices.push(text),
        },
    };
    // SAFETY: optchat uses only these host services; a missing one fails the test.
    const setup = optchat.setup as (
        context: Pick<Plugin.Context, "options">,
    ) => ReturnType<typeof optchat.setup>;
    const cleanup = await setup(context);
    const turn = async (sessionID: string, messages: Message[]) => {
        const event: Event = {
            sessionID,
            agent: "build",
            system: [
                {
                    type: "text",
                    text: "host one",
                    cache: { type: "ephemeral" },
                },
                { type: "text", text: "host two" },
            ],
            messages: messages.map((message) => new Entry(message)),
            tools: { read: "", zoom: "", date: "" },
        };
        await hooks.get("context")!(event);
        return event;
    };
    return {
        directory,
        hooks,
        notices,
        tools,
        generated,
        replies,
        gate,
        turn,
        command: () => command!,
        cleanup,
    };
}

const hints = (event: Event) =>
    [...event.system, ...event.messages.flatMap((message) => message.content)].filter(
        (part) => part.cache,
    ).length;
const settle = () => Bun.sleep(20);

test("optchat is off by default, refuses child sessions, and hides its tools when off", async () => {
    const h = await host();
    expect(h.tools.get("zoom")!.options.codemode).toBe(false);
    expect(h.tools.get("date")!.options.codemode).toBe(false);
    const off = await h.turn("ses_root", [
        { id: "m1", role: "user", content: [{ type: "text", text: "hi" }] },
    ]);
    expect(Object.keys(off.tools)).toEqual(["read"]);
    expect(off.messages).toHaveLength(1);
    await h.command()({ sessionID: "ses_child", prompt: { text: "on" } });
    expect(h.notices.at(-1)).toContain("only to root sessions");
    await h.command()({ sessionID: "ses_root", prompt: { text: "status" } });
    expect(h.notices.at(-1)).toBe("[optchat] off.");
    await expect(
        h.tools.get("zoom")!.execute({ id: 0, n: 1 }, { sessionID: "ses_root", agent: "build" }),
    ).rejects.toThrow("optchat is off");
    if (h.cleanup) await h.cleanup();
});

test("turns see the view and the new message; compactions retry Too long and build the tree", async () => {
    const h = await host();
    await h.command()({
        sessionID: "ses_root",
        prompt: { text: "/optchat on" },
    });
    expect(h.notices.at(-1)).toStartWith("[optchat] on. 0 messages");
    const long = "please remember: " + "detail ".repeat(100);
    h.replies.push("x".repeat(700), "y".repeat(600), "z".repeat(400));
    const first = await h.turn("ses_root", [
        { id: "m1", role: "user", content: [{ type: "text", text: long }] },
    ]);
    expect(first.system[0]!.text).toStartWith("You are build, an AI agent");
    expect(first.system.map((part) => part.text)).toEqual([
        first.system[0]!.text,
        "host one",
        "host two",
    ]);
    expect(first.messages[0]!.content.map((part) => part.text)).toEqual(["<chat>\n</chat>"]);
    expect(first.messages[1]!.content[0]!.text).toBe(long);
    // System start and end, and the request's end; no whole view block yet.
    expect(hints(first)).toBe(3);
    expect(first.messages[1]!.content[0]!.cache).toEqual({
        type: "ephemeral",
        ttlSeconds: 300,
    });
    await settle();

    // The compaction ran in the worker with the same system prompt, retried twice, and kept the shortest line.
    expect(h.generated).toHaveLength(3);
    const task = h.generated[0]!.messages[1]!.content[0]!.text!;
    expect(task).toStartWith("Compaction: compress message 0 into one line of at most 512 bytes");
    expect(task).toContain(`<input>\nuser: ${long}\n</input>`);
    expect(h.generated[0]!.system[0]!.text).toBe(first.system[0]!.text);
    const retry = h.generated[2]!.messages.map((message) => message.role);
    expect(retry).toEqual(["user", "user", "assistant", "user", "assistant", "user"]);
    expect(h.generated[2]!.messages.at(-1)!.content[0]!.text).toStartWith(
        "Too long: your line is 600 bytes",
    );

    const second = await h.turn("ses_root", [
        { id: "m1", role: "user", content: [{ type: "text", text: long }] },
        {
            id: "m2",
            role: "assistant",
            content: [
                { type: "reasoning", text: "never logged" },
                {
                    type: "tool-call",
                    id: "c1",
                    name: "read",
                    input: { path: "a" },
                },
            ],
        },
        {
            id: "m3",
            role: "tool",
            content: [
                {
                    type: "tool-result",
                    id: "c1",
                    name: "read",
                    result: { type: "text", value: "file a" },
                },
            ],
        },
        {
            id: "m4",
            role: "assistant",
            content: [{ type: "text", text: "Read a." }],
        },
        {
            id: "m5",
            role: "user",
            content: [{ type: "text", text: "and now?" }],
        },
    ]);
    // Messages before the new one are summarized before the turn; short ones are their own lines.
    expect(second.messages.map((message) => message.role)).toEqual(["user", "user"]);
    expect(second.messages[0]!.content.map((part) => part.text)).toEqual([
        `<chat>\n0+1|${"z".repeat(400)}\n1+1|tool: read {"path":"a"}\n2+1|echo: read: file a\n3+1|agent: Read a.\n`,
        "</chat>",
    ]);
    // Last whole block, system start and end, and the end.
    expect(hints(second)).toBe(4);
    await settle();

    const zoom = h.tools.get("zoom")!;
    const at = { sessionID: "ses_root", agent: "build" };
    expect((await zoom.execute({ id: 0, n: 1 }, at)).content).toBe(`user: ${long}`);
    expect((await zoom.execute({ id: 0, n: 4 }, at)).content).toBe(
        `0+2|${"z".repeat(400)} tool: read {"path":"a"}\n2+2|echo: read: file a agent: Read a.`,
    );
    await expect(zoom.execute({ id: 1, n: 2 }, at)).rejects.toThrow("No line 1+2");
    expect((await h.tools.get("date")!.execute({ id: 4 }, at)).content).toMatch(
        /^\d{4}-\d{2}-\d{2}T/,
    );

    const compaction: Event = {
        sessionID: "ses_root",
        agent: "build",
        system: [],
        messages: [],
        tools: {},
    };
    await h.hooks.get("compaction")!(compaction);
    expect(compaction.result!.summary).toStartWith("[optchat]");

    if (h.cleanup) await h.cleanup();
    const view = JSON.parse(await readFile(join(h.directory, "ses_root", "view.json"), "utf8"));
    expect(view.lines).toHaveLength(5);
    await expect(readFile(join(h.directory, "ses_root", "lock"))).rejects.toThrow();
});

async function logged(directory: string) {
    const files = await Array.fromAsync(new Bun.Glob("ses_root/main/*.jsonl").scan(directory));
    const lines: string[] = [];
    for (const file of files)
        lines.push(...(await readFile(join(directory, file), "utf8")).trim().split("\n"));
    return lines.map((line) => {
        const message = JSON.parse(line);
        return `${message.kind}: ${message.text}`;
    });
}

test("/optchat off stops queued and running compactions without writing nodes", async () => {
    // Regression: the generate hook found no chat after off, and the reply was saved as a node.
    const h = await host();
    await h.command()({ sessionID: "ses_root", prompt: { text: "on" } });
    let release = () => {};
    h.gate.promise = new Promise((resolve) => (release = resolve));
    await h.turn("ses_root", [
        {
            id: "m1",
            role: "user",
            content: [{ type: "text", text: "remember " + "word ".repeat(200) }],
        },
    ]);
    await settle();
    await h.command()({ sessionID: "ses_root", prompt: { text: "off" } });
    release();
    await settle();
    // The in-flight request still got its own chat's view and task.
    expect(h.generated).toHaveLength(1);
    expect(h.generated[0]!.messages[1]!.content[0]!.text).toStartWith(
        "Compaction: compress message 0",
    );
    const tree = await Array.fromAsync(new Bun.Glob("ses_root/tree/*.jsonl").scan(h.directory));
    expect(tree).toEqual([]);
    if (h.cleanup) await h.cleanup();
});

test("dropped reasoning parts and tool messages without ids do not re-log or skip parts", async () => {
    // Regression: sources were positions, which shift when a model switch drops reasoning parts.
    const h = await host();
    await h.command()({ sessionID: "ses_root", prompt: { text: "on" } });
    const before: Message[] = [
        { id: "m1", role: "user", content: [{ type: "text", text: "hi" }] },
        {
            id: "m2",
            role: "assistant",
            content: [
                { type: "reasoning", text: "" },
                {
                    type: "tool-call",
                    id: "c1",
                    name: "read",
                    input: { path: "a" },
                },
            ],
        },
        {
            role: "tool",
            content: [
                {
                    type: "tool-result",
                    id: "c1",
                    name: "read",
                    result: { type: "text", value: "A" },
                },
            ],
        },
        {
            id: "m3",
            role: "assistant",
            content: [
                { type: "reasoning", text: "" },
                { type: "text", text: "done" },
            ],
        },
        { id: "m4", role: "user", content: [{ type: "text", text: "next" }] },
    ];
    await h.turn("ses_root", before);
    const after = before.map((message) => ({
        ...message,
        content: message.content.filter((part) => part.type !== "reasoning"),
    }));
    await h.turn("ses_root", [
        {
            id: "m0",
            role: "system",
            content: [{ type: "text", text: "Instructions updated" }],
        },
        ...after,
        {
            id: "m5",
            role: "assistant",
            content: [{ type: "text", text: "ok" }],
        },
        { id: "m6", role: "user", content: [{ type: "text", text: "again" }] },
    ]);
    await settle();
    if (h.cleanup) await h.cleanup();
    expect(await logged(h.directory)).toEqual([
        "user: hi",
        'tool: read {"path":"a"}',
        "echo: read: A",
        "agent: done",
        "user: next",
        "agent: ok",
        "user: again",
    ]);
});

test("without compactionModel, optchat registers nothing", async () => {
    const options = { compactionModel: "", dataDir: "" };
    const registered: string[] = [];
    const record = (name: string) => async () => {
        registered.push(name);
        return { dispose: async () => {} };
    };
    const context = {
        options,
        tool: { transform: record("tool") },
        command: { transform: record("command") },
        session: { hook: record("hook") },
    };
    // SAFETY: setup must return before touching any host service; a call is recorded and fails the test.
    const setup = optchat.setup as (
        context: Pick<Plugin.Context, "options">,
    ) => ReturnType<typeof optchat.setup>;
    expect(await setup(context)).toBeUndefined();
    expect(registered).toEqual([]);
});