import { describe, expect, test } from "bun:test";
import { collectCitationSources } from "../app/utils/citations";
import { createMarkdownParser, getNodeChildren, getNodeTag } from "../app/utils/markdown";
import type { MarkdownAstNode } from "../app/utils/markdown";
import { readSearchResponse } from "../server/utils/search";

const citations = (nodes: MarkdownAstNode[]): MarkdownAstNode[] => nodes.flatMap(node =>
    getNodeTag(node) === "source-cite" ? [node] : citations(getNodeChildren(node)),
);

describe("citation parsing", () => {
    test("keeps citation markers as text unless explicitly enabled", async () => {
        const content = "Claim [[cite:1234]] and [[cite:s7,s9]].";
        for (const options of [{}, { citations: false }]) {
            for (const streaming of [false, true]) {
                const tree = await createMarkdownParser(options)(content, { streaming });
                expect(citations(tree.nodes)).toEqual([]);
                expect(getNodeChildren(tree.nodes[0] ?? "")).toEqual([content]);
            }
        }
    });

    test("renders single and grouped references as dedicated nodes", async () => {
        const tree = await createMarkdownParser({ citations: true })("Claim [[cite:s7]]. Another [[cite:s7,s9]].");
        expect(citations(tree.nodes)).toEqual([
            ["source-cite", { ids: "s7" }],
            ["source-cite", { ids: "s7,s9" }],
        ]);
    });

    test("does not interpret code, escaped references, or link labels", async () => {
        for (const content of [
            "`[[cite:s7]]`",
            "```text\n[[cite:s7]]\n```",
            "\\[[cite:s7]]",
            "[Claim [[cite:s7]]](https://example.com)",
        ]) {
            const tree = await createMarkdownParser({ citations: true })(content);
            expect(citations(tree.nodes)).toEqual([]);
        }
    });

    test("streaming never auto-completes a citation", async () => {
        const parse = createMarkdownParser({ citations: true });
        const content = "Claim [[cite:s7,s9]]";
        for (let length = 1; length < content.length; length++) {
            expect(citations((await parse(content.slice(0, length), { streaming: true })).nodes)).toEqual([]);
        }
        expect(citations((await parse(content, { streaming: true })).nodes)).toHaveLength(1);
        expect(citations((await parse(content)).nodes)).toHaveLength(1);
    });
});

describe("citation sources", () => {
    test("IDs survive storage, including duplicate URLs and multiple searches", () => {
        const first = readSearchResponse({ results: [
            { title: "One", url: "https://example.com" },
            { title: "Two", url: "https://example.com" },
        ] });
        const second = readSearchResponse({ results: [{ title: "Three", url: "https://example.com" }] });
        expect(first.ok && second.ok).toBe(true);
        if (!first.ok || !second.ok) {
            return;
        }
        const parts = [first.data, second.data].map(output => ({
            toolCall: { toolName: "search", output: { type: "json", value: JSON.stringify(output) } },
        }));
        const sources = collectCitationSources(parts);
        expect(sources.size).toBe(3);
        expect(sources.get(first.data.results[0]?.sourceId ?? "")?.title).toBe("One");
        expect(collectCitationSources(JSON.parse(JSON.stringify(parts)))).toEqual(sources);
        expect(sources.has("invented")).toBe(false);
        expect(collectCitationSources([]).size).toBe(0);
    });

    test("ignores legacy results, non-search tools, and unsafe URLs", () => {
        const output = { type: "json", value: JSON.stringify({ results: [
            { title: "Legacy", link: "https://example.com" },
            { sourceId: "bad", title: "Bad", link: "javascript:alert(1)" },
        ] }) };
        expect(collectCitationSources([{ toolCall: { toolName: "search", output } }]).size).toBe(0);
        expect(collectCitationSources([{ toolCall: { toolName: "python", output } }]).size).toBe(0);
    });
});