import test from "node:test";
import assert from "node:assert/strict";
import {createServer} from "node:http";
import {createApp} from "../src/app.js";
import Database from "better-sqlite3";
import {runMigrations} from "../src/db/schema.js";

const deferred = () => { let resolve; const promise = new Promise(done => { resolve = done; }); return {promise, resolve}; };

async function fixture(t) {
    let mode = "success", sequence = 0, searches = 0;
    let started = deferred(), release = deferred();
    const calls = [];
    const search = createServer((_req, res) => {
        searches++;
        if (mode === "search-fail") { res.writeHead(503); res.end("Unavailable"); return; }
        res.setHeader("Content-Type", "application/json");
        res.end(JSON.stringify({results: [{title: "Satellite launch news", url: `https://example.org/news/${searches}`, content: "Latest satellite launch news today: a research satellite launched successfully."}]}));
    });
    await new Promise(resolve => search.listen(0, "127.0.0.1", resolve));
    const modelService = {
        async generateReply() { return `Original ${++sequence}`; },
        async streamReply(model, messages, chunk, options) {
            calls.push({model, messages, options});
            if (options.reasoning) options.onThinking("New reasoning", "New reasoning");
            chunk("Partial", "Partial");
            started.resolve();
            if (mode === "hold") await new Promise((resolve, reject) => {
                release.promise.then(resolve);
                options.signal.addEventListener("abort", () => reject(new Error("Generation stopped")), {once: true});
            });
            if (mode === "fail") throw new Error("Model failed after partial output");
            if (mode === "empty") return "";
            chunk(" answer", "Partial answer");
            return "Partial answer";
        },
        mapError(error) { return {status: 500, body: {error: error.message}}; }
    };
    const app = createApp({modelService, config: {dbPath: ":memory:", testMode: true, sessionSecret: "regeneration-test",
        chatSummaryUpdateEveryMessages: 0, eventLoopLagMonitorEnabled: false, structuredRequestLoggingEnabled: false,
        slowRequestLoggingEnabled: false, searxngUrl: `http://127.0.0.1:${search.address().port}`}});
    const server = app.listen(0, "127.0.0.1");
    await new Promise(resolve => server.once("listening", resolve));
    t.after(async () => {
        release.resolve();
        await new Promise(resolve => { server.closeAllConnections(); server.close(resolve); });
        await new Promise(resolve => { search.closeAllConnections(); search.close(resolve); });
        await app.locals.close();
    });
    let cookie;
    const api = async (path, body, method = body ? "POST" : "GET") => {
        const response = await fetch(`http://127.0.0.1:${server.address().port}${path}`, {method,
            headers: {"Content-Type": "application/json", ...(cookie ? {Cookie: cookie} : {})}, body: body ? JSON.stringify(body) : undefined});
        if (response.headers.get("set-cookie")) cookie = response.headers.get("set-cookie").split(";")[0];
        return response;
    };
    await api("/register", {username: "regenerator", password: "test-password"});
    const chatId = (await (await api("/chats", {title: "Versions"})).json()).chat.id;
    await api("/chat", {chatId, message: "Search the web for the latest satellite launch news today"});
    await api("/chat", {chatId, message: "Later private question"});
    const messages = (await (await api(`/chats/${chatId}/messages`)).json()).messages;
    const bot = messages[1];
    app.locals.repositories.updateChatMessageThinking(chatId, bot.id, "Original reasoning");
    app.locals.repositories.updateChatMessageSources(chatId, bot.id, [{title: "Old", url: "https://example.org/old", snippet: "Old evidence"}]);
    return {app, api, chatId, bot, messages, calls,
        get searches() { return searches; },
        path: `/chats/${chatId}/messages/${bot.id}`,
        hold() { mode = "hold"; started = deferred(); release = deferred(); return started.promise; },
        setMode(value) { mode = value; }, release: () => release.resolve(),
        history: async () => (await (await api(`/chats/${chatId}/messages`)).json()).messages};
}

test("regeneration streams before commit in the same message; versions retain sources and reasoning and branch from selection", async t => {
    const s = await fixture(t);
    const started = s.hold();
    const response = await s.api(`${s.path}/retry/stream`, {reasoning: true, webSearch: true});
    await started;
    const reader = response.body.getReader();
    let wire = new TextDecoder().decode((await reader.read()).value);
    while (!wire.includes("Partial")) wire += new TextDecoder().decode((await reader.read()).value);
    assert.match(wire, /event: thinking/);
    assert.match(wire, /Partial/);
    assert.doesNotMatch(wire, /event: done/);
    let history = await s.history();
    assert.equal(history[1].content, s.bot.content, "database retains original while streaming");
    assert.equal(history[1].versionCount, 1);
    assert.equal(history.length, 4);
    assert.equal((await s.api(`${s.path}/retry`, {})).status, 409);
    assert.equal((await s.api(s.path, {content: "unsafe edit"}, "PUT")).status, 409);
    assert.equal((await s.api(s.path, null, "DELETE")).status, 409);
    assert.equal((await s.api(`/chats/${s.chatId}/messages/${s.messages[3].id}`, {content: "Other answer edited"}, "PUT")).status, 200);
    s.release();
    while (true) { const result = await reader.read(); if (result.done) break; wire += new TextDecoder().decode(result.value); }
    assert.match(wire, /event: done/);
    assert.ok(wire.indexOf('"status":"searching"') < wire.indexOf("event: chunk"));
    history = await s.history();
    assert.deepEqual(history.map(m => m.id), s.messages.map(m => m.id));
    assert.equal(history[1].content, "Partial answer");
    assert.equal(history[1].versionCount, 2);
    assert.equal(history[1].versionIndex, 1);
    assert.equal(history[1].thinking, "New reasoning");
    assert.equal(history[1].sources[0].url, "https://example.org/news/1");
    assert.equal(s.searches, 1);
    assert.ok(s.calls[0].messages.every(m => !m.content.includes("Later private question") && !m.content.includes("Original 2")));
    assert.equal(s.calls[0].messages.at(-1).content, s.messages[0].content);
    await s.api(`${s.path}/versions/0`, {});
    history = await s.history();
    assert.equal(history[1].content, s.bot.content);
    assert.equal(history[1].thinking, "Original reasoning");
    assert.equal(history[1].sources[0].url, "https://example.org/old");
    const branch = await (await s.api(`${s.path}/branch`, {confirmed: true})).json();
    assert.equal(branch.messages.at(-1).content, s.bot.content);
    assert.equal(branch.messages.at(-1).thinking, "Original reasoning");
    assert.deepEqual(branch.messages.at(-1).sources, history[1].sources);
    await s.api(s.path, {content: "Edited original"}, "PUT");
    await s.api(`${s.path}/versions/1`, {});
    assert.equal((await s.history())[1].content, "Partial answer");
    await s.api(`${s.path}/versions/0`, {});
    assert.equal((await s.history())[1].content, "Edited original");
    assert.equal((await s.api(`${s.path}/versions/99`, {})).status, 400);
    s.setMode("success");
    const third = await (await s.api(`${s.path}/retry/stream`, {})).text();
    assert.match(third, /event: done/);
    const thirdMessage = (await s.history())[1];
    assert.equal(thirdMessage.versionCount, 3);
    assert.equal(thirdMessage.versionIndex, 2);
    assert.equal(thirdMessage.thinking, undefined);
    assert.deepEqual(thirdMessage.sources, []);
    const backup = (await (await s.api("/exports/workspace")).json()).workspace;
    assert.equal((await s.api("/imports/workspace", {workspace: backup, allowDuplicates: true})).status, 200);
    const copiedVersions = s.app.locals.db.prepare("SELECT * FROM chat_message_versions WHERE message_id != ?").all(s.bot.id);
    assert.equal(copiedVersions.length, 3);
    assert.equal(copiedVersions[0].thinking, "Original reasoning");
    assert.equal(copiedVersions[1].thinking, "New reasoning");
    const reopened = new Database(s.app.locals.db.serialize());
    try {
        runMigrations(reopened, {dropLegacyChats: false});
        assert.equal(reopened.prepare("SELECT count(*) n FROM chat_message_versions WHERE message_id = ?").get(s.bot.id).n, 3);
        assert.equal(reopened.prepare("SELECT retry_active_index FROM chat_messages WHERE id = ?").get(s.bot.id).retry_active_index, 2);
    } finally { reopened.close(); }
    await s.api(s.path, null, "DELETE");
    assert.equal(s.app.locals.repositories.listMessageVersions(s.bot.id).length, 0);
});

test("failed, empty, stopped and disconnected regeneration preserve the saved answer and versions", async t => {
    const s = await fixture(t);
    const original = (await s.history())[1];
    for (const mode of ["fail", "empty", "search-fail"]) {
        s.setMode(mode);
        const wire = await (await s.api(`${s.path}/retry/stream`, {reasoning: true, webSearch: mode === "search-fail"})).text();
        assert.match(wire, /event: error/);
        assert.doesNotMatch(wire, /event: done/);
        assert.deepEqual((await s.history())[1], original);
        assert.equal(s.app.locals.repositories.listMessageVersions(s.bot.id).length, 0);
        if (mode === "search-fail") assert.match(wire, /"searchState":"failed"/);
    }
    const started = s.hold();
    const response = await s.api(`${s.path}/retry/stream`, {});
    await started;
    assert.equal((await (await s.api(`/chats/${s.chatId}/stop`, {})).json()).stopped, true);
    assert.match(await response.text(), /event: error/);
    assert.deepEqual((await s.history())[1], original);
    const disconnected = s.hold();
    const disconnectResponse = await s.api(`${s.path}/retry/stream`, {});
    await disconnected;
    await disconnectResponse.body.cancel();
    await new Promise(resolve => setTimeout(resolve, 20));
    assert.deepEqual((await s.history())[1], original);
    assert.equal(s.app.locals.repositories.listMessageVersions(s.bot.id).length, 0);
});