import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";
import type { JsonValue } from "@earendil-works/chord";
import {
	type AssistantMessage,
	fauxAssistantMessage,
	fauxText,
	fauxToolCall,
	type ToolResultMessage,
	Type,
} from "@earendil-works/pi-ai";
import {
	defineTool,
	type EntryRecord,
	GenerationTask,
	type Harness,
	LiveDoc,
	type TaskId,
	type ToolRegistration,
	ToolResultEntry,
	ToolTask,
} from "@earendil-works/pi-durable";
import { afterEach, describe, expect, it } from "vitest";
import type { ExecutionEnv } from "../src/env/index.ts";
import { NodeExecutionEnv } from "../src/env/node.ts";
import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts";
import { createBashTool } from "../src/tools/index.ts";
import { allEntries, type ChatSetup, chatSetup, openChat, waitFor } from "./chat-support.ts";
import { addHooks, addTool, type Installed } from "./harness-support.ts";
import { context } from "./session-support.ts";
import { aborted, deferred } from "./task-support.ts";

const directories = new Set<string>();

afterEach(async () => {
	for (const directory of directories) await rm(directory, { recursive: true, force: true });
	directories.clear();
});

async function sqlitePath(): Promise<string> {
	const directory = await mkdtemp(join(tmpdir(), "pi-durable-tools-"));
	directories.add(directory);
	return join(directory, "session.sqlite");
}

async function open(path: string, setup: ChatSetup, env?: ExecutionEnv) {
	const opened = await openChat(await openNodeSqliteStorage(path), setup, env === undefined ? {} : { env });
	opened.harness.resume();
	return opened;
}

function tool(
	name: string,
	execute: ToolRegistration["execute"],
	extra: Partial<ToolRegistration> = {},
): ToolRegistration {
	return defineTool({ name, description: name, parameters: Type.Object({}), execute, ...extra });
}

function call(name: string, id = "c1"): AssistantMessage {
	return fauxAssistantMessage([fauxToolCall(name, {}, { id })], { stopReason: "toolUse" });
}

const DONE = fauxAssistantMessage([fauxText("done")]);

function results(entries: readonly EntryRecord[]): ToolResultMessage[] {
	return entries.filter((entry) => ToolResultEntry.is(entry)).map((entry) => entry.model![0] as ToolResultMessage);
}

function text(message: ToolResultMessage | undefined): string {
	return (message?.content ?? []).map((item) => (item.type === "text" ? item.text : "")).join("|");
}

async function toolTaskId(harness: Harness): Promise<TaskId> {
	let id: TaskId | undefined;
	await waitFor(async () => {
		id = (await harness.snapshot(LiveDoc, 1 as never, context))?.tools?.[0]?.taskId;
		return id !== undefined;
	});
	return id!;
}

/**
 * A tool that writes output, then blocks until its invocation is cancelled the first `blocking` times it runs.
 * `started` resolves once the output is durable.
 */
function blockingTool(name: string, extra: Partial<ToolRegistration> = {}) {
	const started = deferred();
	const state = { runs: 0, blocking: 1 };
	const registration = tool(
		name,
		async (_args, api, callContext) => {
			state.runs++;
			api.output(`run ${state.runs}\n`);
			await api.details({ run: state.runs }, callContext);
			if (state.runs <= state.blocking) {
				started.resolve();
				await aborted(callContext.abortSignal!);
			}
			return {};
		},
		extra,
	);
	return { registration, started, state };
}

describe("tool recovery", () => {
	it("answers an unsafe tool interrupted after intent with its durable partial output", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		const { registration, started, state } = blockingTool("work");
		addTool(setup.registry, registration);
		setup.faux.setResponses([call("work"), DONE]);
		let opened = await open(path, setup);
		const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
		await started.promise;
		const taskId = await toolTaskId(opened.harness);
		await opened.harness.close(context);

		opened = await open(path, setup);
		expect((await opened.harness.getTask(taskId, context))?.state).toMatchObject({
			checkpoint: { phase: "execute", arguments: {}, replay: "unsafe" },
		});
		expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
		expect(state.runs).toBe(1);
		const [result] = results(await allEntries(opened.root));
		expect(result).toMatchObject({ isError: true, details: { run: 1 } });
		expect(text(result)).toBe(
			"run 1\n|<harness>\n[error] Tool work was interrupted and may have partially run\n</harness>",
		);
		expect(await opened.harness.snapshot(LiveDoc, opened.root.id, context)).toEqual({});
		await opened.harness.close(context);
	});

	it("reruns a tool only when both the stored and the current replay policy are safe", async () => {
		const cases: [stored: "safe" | "unsafe", current: "safe" | "unsafe", reruns: boolean][] = [
			["safe", "safe", true],
			["safe", "unsafe", false],
			["unsafe", "safe", false],
		];
		for (const [stored, current, reruns] of cases) {
			const path = await sqlitePath();
			const setup = chatSetup();
			const blocking = blockingTool("work", { replay: stored });
			const registration: Installed = addTool(setup.registry, blocking.registration);
			setup.faux.setResponses([call("work"), DONE]);
			let opened = await open(path, setup);
			const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
			await blocking.started.promise;
			await opened.harness.close(context);

			registration.dispose();
			addTool(setup.registry, { ...blocking.registration, replay: current });
			opened = await open(path, setup);
			expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
			const [result] = results(await allEntries(opened.root));
			expect(blocking.state.runs).toBe(reruns ? 2 : 1);
			expect(result!.isError).toBe(!reruns);
			if (reruns) expect(text(result)).toBe("run 2\n");
			await opened.harness.close(context);
		}
	});

	it("reruns a safe tool with the environment of the conversation's cwd at rerun, and not once it is deselected", async () => {
		for (const change of ["cwd", "deselect"] as const) {
			const path = await sqlitePath();
			const setup = chatSetup();
			const cwds: (string | undefined)[] = [];
			const started = deferred();
			const work = tool(
				"work",
				async (_args, api, callContext) => {
					cwds.push(api.env?.cwd);
					if (cwds.length === 1) {
						started.resolve();
						await aborted(callContext.abortSignal!);
					}
					return {};
				},
				{ replay: "safe" },
			);
			addTool(setup.registry, work);
			setup.faux.setResponses([call("work"), DONE]);
			const env = ({ cwd = "/" }: { readonly cwd?: string }) => new NodeExecutionEnv({ cwd });
			const reopen = async () => {
				const opened = await openChat(await openNodeSqliteStorage(path), setup, { env });
				opened.harness.resume();
				return opened;
			};
			let opened = await reopen();
			await opened.root.configure({ cwd: "/one" }, context);
			const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
			await started.promise;
			await opened.root.configure(change === "cwd" ? { cwd: "/two" } : { extensions: [] }, context);
			await opened.harness.close(context);

			opened = await reopen();
			expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
			const [result] = results(await allEntries(opened.root));
			if (change === "cwd") {
				expect(cwds).toEqual(["/one", "/two"]);
				expect(result!.isError).toBe(false);
			} else {
				// A tool that no longer resolves is treated as unsafe: interrupted, not rerun.
				expect(cwds).toEqual(["/one"]);
				expect(text(result)).toContain("was interrupted");
			}
			await opened.harness.close(context);
		}
	});

	it("reruns beforeTool when interrupted before intent, and executes once", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		let runs = 0;
		addTool(
			setup.registry,
			tool("work", async () => {
				runs++;
				return { content: [] };
			}),
		);
		const reached = deferred();
		let asked = 0;
		const decisions: string[] = [];
		addHooks(setup.registry, ToolTask, {
			beforeTool: async (_call, api, callContext) => {
				asked++;
				// A durable first-writer-wins decision survives the rerun.
				decisions.push(await api.memo("test:decision", `attempt ${asked}`, callContext));
				if (asked === 1) {
					reached.resolve();
					await aborted(callContext.abortSignal!);
				}
				return undefined;
			},
		});
		setup.faux.setResponses([call("work"), DONE]);
		let opened = await open(path, setup);
		const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
		await reached.promise;
		await opened.harness.close(context);

		opened = await open(path, setup);
		expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
		expect([asked, runs]).toEqual([2, 1]);
		expect(decisions).toEqual(["attempt 1", "attempt 1"]);
		await opened.harness.close(context);
	});

	it("reruns the generation tools phase interrupted before its commit", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		addTool(
			setup.registry,
			tool("work", async () => ({ content: [] })),
		);
		const reached = deferred();
		let observed = 0;
		addHooks(setup.registry, GenerationTask, {
			afterTools: async (_assistant, _results, _api, callContext) => {
				observed++;
				if (observed === 1) {
					reached.resolve();
					await aborted(callContext.abortSignal!);
				}
			},
		});
		setup.faux.setResponses([call("work"), DONE]);
		let opened = await open(path, setup);
		const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
		await reached.promise;
		await opened.harness.close(context);

		opened = await open(path, setup);
		expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
		expect(observed).toBe(2);
		expect((await allEntries(opened.root)).map((entry) => entry.kind)).toEqual([
			"pi.user",
			"pi.system",
			"pi.assistant",
			"pi.tool-result",
			"pi.assistant",
		]);
		await opened.harness.close(context);
	});

	it("answers an aborted tool with its partial output and continues the run", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		const { registration, started } = blockingTool("work");
		addTool(setup.registry, registration);
		setup.faux.setResponses([call("work"), DONE]);
		const opened = await open(path, setup);
		const submission = await opened.root.submit({ type: "input", content: "go" }, context);
		await started.promise;
		const taskId = await toolTaskId(opened.harness);
		expect(await opened.harness.abortTask(taskId, context)).toBe("marked");
		expect((await opened.harness.waitForTask(taskId, context)).state.outcome).toMatchObject({
			status: "aborted",
			result: { entryId: expect.any(Number) },
		});
		expect((await submission.wait(context)).status).toBe("done");
		const [result] = results(await allEntries(opened.root));
		expect(text(result)).toBe("run 1\n|<harness>\n[error] Tool work was aborted\n</harness>");
		await opened.harness.close(context);
	});

	it("lets context derivation answer a faulted tool and continues the run", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		// A result that is not strict JSON makes the result commit throw, so the scheduler faults the task.
		addTool(
			setup.registry,
			tool("bad", async () => ({ content: [], details: { fn: (() => 1) as unknown as JsonValue } })),
		);
		const requests: string[] = [];
		setup.faux.setResponses([
			call("bad"),
			(request) => {
				const result = request.messages.find((message) => message.role === "toolResult");
				requests.push(result?.role === "toolResult" ? text(result) : "none");
				return DONE;
			},
		]);
		const opened = await open(path, setup);
		const settled = await (await opened.root.submit({ type: "input", content: "go" }, context)).wait(context);
		expect(settled.status).toBe("done");
		expect(results(await allEntries(opened.root))).toEqual([]);
		expect(requests).toEqual(["Tool result unavailable: history ends before this call completed."]);
		const tasks = await opened.harness.commit((tx) => tx.scanTasks({ conversationId: opened.root.id }, 20), context);
		const faulted = tasks.items.find((task) => task.kind === "pi.tool")!;
		expect(faulted.state).toMatchObject({ status: "terminal", outcome: { status: "faulted" } });
		await opened.harness.close(context);
	});

	it("answers a real bash command interrupted by close and reopen, then finishes the run", async () => {
		const path = await sqlitePath();
		const env = new NodeExecutionEnv({ cwd: dirname(path) });
		const setup = chatSetup();
		addTool(setup.registry, createBashTool());
		setup.faux.setResponses([
			fauxAssistantMessage([fauxToolCall("bash", { command: "echo started; sleep 30" }, { id: "b" })], {
				stopReason: "toolUse",
			}),
			DONE,
		]);
		let opened = await open(path, setup, env);
		const id = (await opened.root.submit({ type: "input", content: "go" }, context)).id;
		await waitFor(async () => {
			const slot = (await opened.harness.snapshot(LiveDoc, opened.root.id, context))?.tools?.[0];
			return slot?.output === "started\n";
		});
		await opened.harness.close(context);

		opened = await open(path, setup, env);
		expect((await (await opened.harness.submission(id, context))!.wait(context)).status).toBe("done");
		const [result] = results(await allEntries(opened.root));
		expect(text(result)).toBe(
			"started\n|<harness>\n[error] Tool bash was interrupted and may have partially run\n</harness>",
		);
		await opened.harness.close(context);
	});

	it("clears the interrupted attempt's progress before a safe rerun", async () => {
		const path = await sqlitePath();
		const setup = chatSetup();
		const started = [deferred(), deferred()];
		let runs = 0;
		addTool(
			setup.registry,
			tool(
				"work",
				async (_args, api, callContext) => {
					const run = runs++;
					if (run === 0) {
						api.diagnostic({ severity: "info", message: "first a" });
						api.diagnostic({ severity: "info", message: "first b" });
						await api.details({ run: 1, extra: true }, callContext);
					} else {
						api.diagnostic({ severity: "info", message: "second" });
						await api.details({ run: 2 }, callContext);
					}
					started[run]!.resolve();
					await aborted(callContext.abortSignal!);
					return {};
				},
				{ replay: "safe" },
			),
		);
		setup.faux.setResponses([call("work"), DONE]);
		let opened = await open(path, setup);
		await opened.root.submit({ type: "input", content: "go" }, context);
		await started[0]!.promise;
		await opened.harness.close(context);

		opened = await open(path, setup);
		await started[1]!.promise;
		const slot = (await opened.harness.snapshot(LiveDoc, opened.root.id, context))?.tools?.[0];
		expect(slot?.diagnostics).toEqual([{ severity: "info", message: "second" }]);
		expect(slot?.details).toEqual({ run: 2 });
		await opened.harness.close(context);
	});
});