import { describe, expect, test } from "bun:test";
import { z } from "zod";
import HarnessError from "../harness-error";
import type { HarnessCandidate, HarnessTool } from "../types";
import type { GenerateTextParams, GenerateTextResult } from "./generate-text";
import runTextLoop, { IncomingMessagesNote, RedraftNote } from "./run-text-loop";

const candidate: HarnessCandidate = {
    key: "1:m",
    provider: { id: 1, kind: "openrouter", base_url: "https://x.test" },
    model: {
        slug: "m",
        input_modalities: ["text"],
        output_modalities: ["text"],
        supports_tools: true,
    },
    source: "auto",
};

const memory: string[] = [];

const remember: HarnessTool<{ note: string }> = {
    name: "remember",
    description: "Keeps a note.",
    schema: z.object({ note: z.string() }),
    run: async ({ args }) => {
        memory.push(args.note);
        return "Saved.";
    },
};

/**
 * A fake model that plays back scripted turns and records what it saw.
 */
function scripted(turns: Partial<GenerateTextResult>[]) {
    const seen: GenerateTextParams[] = [];

    const generate = async (params: GenerateTextParams) => {
        seen.push(structuredClone({ ...params, signal: undefined }));
        const turn = turns[seen.length - 1] || { text: "Done." };

        return {
            text: null,
            tool_calls: [],
            candidate,
            ...turn,
        } as GenerateTextResult;
    };

    return { generate, seen };
}

const resolveCandidates = async () => [candidate];

describe("runTextLoop", () => {
    test("runs tools and feeds their results back until the model answers", async () => {
        memory.length = 0;
        const { generate, seen } = scripted([
            {
                text: "Noting that.",
                tool_calls: [
                    { id: "c1", name: "remember", arguments: `{"note":"tea"}` },
                ],
                usage: { prompt_tokens: 10, completion_tokens: 3, cost: 0.5 },
            },
            {
                text: "I'll remember you like tea.",
                usage: { prompt_tokens: 20, completion_tokens: 5 },
            },
        ]);
        const steps: number[] = [];

        const result = await runTextLoop({
            system: "You are a bot.",
            messages: [{ role: "user", content: "I like tea." }],
            tools: [remember],
            generate,
            resolveCandidates,
            onStep: ({ index }) => {
                steps.push(index);
            },
        });

        expect(result.text).toBe("I'll remember you like tea.");
        expect(result.steps).toBe(2);
        expect(steps).toEqual([0, 1]);
        expect(memory).toEqual(["tea"]);
        expect(result.usage).toEqual({
            prompt_tokens: 30,
            completion_tokens: 8,
            cost: 0.5,
        });
        expect(seen[0]!.messages[0]).toEqual({
            role: "system",
            content: "You are a bot.",
        });
        expect(seen[0]!.tools?.[0]?.name).toBe("remember");
        expect(seen[1]!.prefer_key).toBe("1:m");
        expect(seen[1]!.messages.slice(-2)).toEqual([
            {
                role: "assistant",
                content: "Noting that.",
                tool_calls: [
                    { id: "c1", name: "remember", arguments: `{"note":"tea"}` },
                ],
            },
            { role: "tool", tool_call_id: "c1", content: "Saved." },
        ]);
    });

    test("asks the model to wrap up, without tools, after the last step", async () => {
        const loop_call = {
            tool_calls: [{ id: "c", name: "remember", arguments: `{"note":"x"}` }],
        };
        const { generate, seen } = scripted([
            loop_call,
            loop_call,
            { text: "Ran out of room; here's where I got to.", ...loop_call },
        ]);

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "go" }],
            tools: [remember],
            max_steps: 2,
            generate,
            resolveCandidates,
        });

        expect(seen).toHaveLength(3);
        expect(seen[2]!.tools).toBe(undefined);
        expect(seen[2]!.messages.at(-1)).toMatchObject({ role: "user" });
        expect(seen[2]!.messages.at(-1)!.content).toContain(
            "The tools you actually called this run: remember x2.",
        );
        expect(result.tool_uses).toEqual([
            { name: "remember", ok: true },
            { name: "remember", ok: true },
        ]);
        expect(result.text).toBe("Ran out of room; here's where I got to.");
        expect(result.messages.at(-1)).toEqual({
            role: "assistant",
            content: "Ran out of room; here's where I got to.",
        });
    });

    test("records failed tool calls, and tells the wrap up so", async () => {
        const { generate, seen } = scripted([
            { tool_calls: [{ id: "c", name: "missing_tool", arguments: "{}" }] },
            { text: "Couldn't do it." },
        ]);

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "go" }],
            tools: [remember],
            max_steps: 1,
            generate,
            resolveCandidates,
        });

        expect(result.tool_uses).toEqual([{ name: "missing_tool", ok: false }]);
        expect(seen[1]!.messages.at(-1)!.content).toContain(
            "The tools you actually called this run: missing_tool (failed).",
        );
    });

    test("stops when cancelled", async () => {
        const controller = new AbortController();
        controller.abort();

        const error = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "go" }],
            signal: controller.signal,
            generate: scripted([]).generate,
            resolveCandidates,
        }).catch((e) => e);

        expect(error).toBeInstanceOf(HarnessError);
        expect(error.code).toBe("aborted");
    });

    test("takes in what the user wrote while tools ran, before the next turn", async () => {
        memory.length = 0;
        const { generate, seen } = scripted([
            { tool_calls: [{ id: "c1", name: "remember", arguments: `{"note":"tea"}` }] },
            { text: "Coffee it is." },
        ]);
        const arrivals = [[{ role: "user" as const, content: "Actually, coffee." }]];

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "I like tea." }],
            tools: [remember],
            generate,
            resolveCandidates,
            incoming: async () => arrivals.shift() || [],
        });

        expect(result.text).toBe("Coffee it is.");
        expect(seen).toHaveLength(2);
        expect(seen[1]!.messages.slice(-2)).toEqual([
            { role: "tool", tool_call_id: "c1", content: "Saved." },
            { role: "user", content: `${IncomingMessagesNote}\n\nActually, coffee.` },
        ]);
    });

    test("writes an answer again when the user wrote past it, keeping the draft and tool turns", async () => {
        const { generate, seen } = scripted([
            { text: "Here's the summary in English." },
            { text: "Here's the summary in French." },
        ]);
        const arrivals = [
            [
                { role: "user" as const, content: "In French, please." },
                { role: "user" as const, content: "And short." },
            ],
        ];

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "Summarise it." }],
            generate,
            resolveCandidates,
            incoming: async () => arrivals.shift() || [],
        });

        expect(result.text).toBe("Here's the summary in French.");
        expect(result.steps).toBe(2);
        expect(seen[1]!.messages.slice(-2)).toEqual([
            { role: "assistant", content: "Here's the summary in English." },
            {
                role: "user",
                content: `${RedraftNote}\n\nIn French, please.\n\nAnd short.`,
            },
        ]);
    });

    test("joins files from new messages onto the same user turn", async () => {
        const { generate, seen } = scripted([{ text: "Draft." }, { text: "With the photo." }]);
        const arrivals = [
            [
                {
                    role: "user" as const,
                    content: [
                        { type: "text" as const, text: "Use this one." },
                        { type: "image" as const, url: "data:image/png;base64,x" },
                    ],
                },
            ],
        ];

        await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "Write the caption." }],
            generate,
            resolveCandidates,
            incoming: async () => arrivals.shift() || [],
        });

        expect(seen[1]!.messages.at(-1)).toEqual({
            role: "user",
            content: [
                { type: "text", text: RedraftNote },
                { type: "text", text: "Use this one." },
                { type: "image", url: "data:image/png;base64,x" },
            ],
        });
    });

    test("returns at once when nothing new came in", async () => {
        const { generate, seen } = scripted([{ text: "Hi." }]);
        let looks = 0;

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "Hello" }],
            generate,
            resolveCandidates,
            incoming: async () => {
                looks++;
                return [];
            },
        });

        expect(result.text).toBe("Hi.");
        expect(seen).toHaveLength(1);
        expect(looks).toBe(1);
    });

    test("doesn't redraft the wrap up: a later reply picks the new messages up", async () => {
        const loop_call = {
            tool_calls: [{ id: "c", name: "remember", arguments: `{"note":"x"}` }],
        };
        const { generate, seen } = scripted([loop_call, { text: "Out of steps." }]);
        let looks = 0;

        const result = await runTextLoop({
            system: "s",
            messages: [{ role: "user", content: "go" }],
            tools: [remember],
            max_steps: 1,
            generate,
            resolveCandidates,
            incoming: async () => {
                looks++;
                return looks == 1 ? [{ role: "user", content: "Also this." }] : [];
            },
        });

        expect(result.text).toBe("Out of steps.");
        expect(seen).toHaveLength(2);
        expect(looks).toBe(1);
        expect(typeof seen[1]!.messages.at(-1)!.content).toBe("string");
        expect(seen[1]!.messages.at(-1)!.content).toContain("Also this.");
        expect(seen[1]!.messages.at(-1)!.content).toContain("You've used every step");
    });
});