import { spawnSync } from "node:child_process";
import { existsSync, mkdtempSync, readFileSync, rmSync } from "node:fs";
import { homedir, tmpdir } from "node:os";
import { join, resolve } from "node:path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import {
	AGENT_MESSAGE_DISPLAY_MIME,
	ATTACHMENT_DISPLAY_MIME,
	CONTENT_DISPLAY_MIME,
	DIFF_DISPLAY_MIME,
	type HostRequestHandlers,
	ReplKernelManager,
} from "../src/core/kernel/index.js";

function resolveReplPython(): string | null {
	const candidates = [
		process.env.PRIME_AGENT_KERNEL_PYTHON,
		resolve(__dirname, "..", "..", "..", "prime-agent-runtime", ".venv", "bin", "python"),
		join(homedir(), ".prime", "agent", "kernel-venv", "bin", "python"),
	].filter((p): p is string => Boolean(p));
	for (const python of candidates) {
		if (!existsSync(python)) continue;
		const check = spawnSync(python, ["-c", "import rlm.repl, dill"], { encoding: "utf8" });
		if (check.status === 0) return python;
	}
	return null;
}

const python = resolveReplPython();
const describeIf = python ? describe : describe.skip;

describeIf("ReplKernelManager execute (real runtime)", () => {
	let dir = "";
	let manager: ReplKernelManager | undefined;

	beforeEach(() => {
		dir = mkdtempSync(join(tmpdir(), "prime-agent-repl-execute-"));
	});

	afterEach(async () => {
		await manager?.shutdown({ snapshot: true, drainHostRequests: true });
		manager = undefined;
		if (dir) {
			rmSync(dir, { recursive: true, force: true });
			dir = "";
		}
	});

	it("streams stdout/stderr, returns results, and persists state across cells", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const chunks: { name: string; text: string }[] = [];
		const first = await manager.execute("import sys\nx = 21\nprint('to-out')\nsys.stderr.write('to-err\\n')", {
			onStream: (chunk, name) => chunks.push({ name, text: chunk }),
		});
		expect(first.status).toBe("ok");
		expect(first.stdout).toContain("to-out");
		expect(first.stderr).toContain("to-err");
		expect(chunks.some((c) => c.name === "stdout" && c.text.includes("to-out"))).toBe(true);

		const second = await manager.execute("x * 2");
		expect(second.status).toBe("ok");
		expect(second.result).toBe("42");
	}, 30_000);

	it("returns the generated skill-surface digest the prompt publishes", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		await manager.start();

		const digest = await manager.skillDigest(["rlm"]);

		// Generated from the live modules, so it names the runtime's own calls and the
		// settled/pending convention rather than any hand-written copy.
		expect(digest).toContain("# skill surface");
		expect(digest).toContain("bash(");
	}, 30_000);

	it("reports cell errors with a clean traceback", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const r = await manager.execute("def boom():\n    raise ValueError('nope')\nboom()");
		expect(r.status).toBe("error");
		expect(r.error?.ename).toBe("ValueError");
		expect(r.error?.evalue).toBe("nope");
		expect(r.error?.traceback.join("")).toContain("raise ValueError('nope')");
	}, 30_000);

	it("parses emitted display payloads into diffs, attachments, and sent messages", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const code = [
			"from rlm import emit",
			`emit({${JSON.stringify(DIFF_DISPLAY_MIME)}: {"path": "/tmp/f.py", "old_str": "a", "new_str": "b", "start_line": 3}})`,
			`emit({${JSON.stringify(CONTENT_DISPLAY_MIME)}: {"content": "diff --git a/f.py b/f.py", "title": "my reading diff", "format": "diff"}})`,
			`emit({${JSON.stringify(ATTACHMENT_DISPLAY_MIME)}: {"mime_type": "image/png", "data": "aGVsbG8=", "path": "/tmp/i.png"}})`,
			`emit({${JSON.stringify(AGENT_MESSAGE_DISPLAY_MIME)}: {"id": "m1", "message": "hi", "deliveryStatus": "delivered", "receiverRole": "parent", "target": {"activeSessionId": "a", "sessionId": "s"}}})`,
		].join("\n");
		const r = await manager.execute(code);
		expect(r.status).toBe("ok");
		expect(r.diffs).toEqual([{ path: "/tmp/f.py", oldStr: "a", newStr: "b", startLine: 3 }]);
		expect(r.displays).toEqual([{ content: "diff --git a/f.py b/f.py", title: "my reading diff", format: "diff" }]);
		expect(r.attachments).toEqual([{ mimeType: "image/png", data: "aGVsbG8=", path: "/tmp/i.png" }]);
		expect(r.sentAgentMessages).toEqual([
			{
				id: "m1",
				message: "hi",
				deliveryStatus: "delivered",
				receiverRole: "parent",
				target: { activeSessionId: "a", sessionId: "s" },
			},
		]);
	}, 30_000);

	it("round-trips host requests through hostHandlers, including error replies", async () => {
		const hostHandlers: HostRequestHandlers = {
			"test.echo": async (payload) => ({ echoed: payload.value, cell: payload.cellSourceCode }),
			"test.fail": async () => {
				throw new Error("handler exploded");
			},
		};
		manager = new ReplKernelManager({ python: python as string, cwd: dir, hostHandlers });

		const ok = await manager.execute(
			"import rlm\nreply = await rlm.host_request('test.echo', {'value': 7})\nreply['echoed']",
		);
		expect(ok.status).toBe("ok");
		expect(ok.result).toBe("7");

		const cellSource = await manager.execute("reply['cell']");
		expect(cellSource.status).toBe("ok");
		expect(cellSource.result).toContain("test.echo");

		const failed = await manager.execute("import rlm\nawait rlm.host_request('test.fail')");
		expect(failed.status).toBe("error");
		expect(failed.error?.ename).toBe("RuntimeError");
		expect(failed.error?.evalue).toBe("handler exploded");

		const unknown = await manager.execute("import rlm\nawait rlm.host_request('test.unknown')");
		expect(unknown.status).toBe("error");
		expect(unknown.error?.evalue).toContain('host request type "test.unknown" is not available');
	}, 30_000);

	it("spawns through rlm.spawn over the unchanged rlm.run wire type, requires a child name, and refuses a direct rlm call", async () => {
		const requestTypes: string[] = [];
		const hostHandlers: HostRequestHandlers = {
			"rlm.run": async (payload) => {
				requestTypes.push("rlm.run");
				expect(payload.prompt).toMatch(/^You are child subagent 'worker' at depth \d+\.[\s\S]*\nTASK:\nchild work/);
				expect(payload.kwargs).toEqual({ name: "worker", model: "anthropic/claude-opus-4-7" });
				return {
					rlm_child_id: "sub-a1b2c3d4",
					name: "worker",
					session_dir: dir,
					model: "anthropic/claude-opus-4-7",
				};
			},
		};
		manager = new ReplKernelManager({ python: python as string, cwd: dir, hostHandlers });

		const spawned = await manager.execute(
			"import rlm\nhandle = await rlm.spawn('child work', name='worker', model='anthropic/claude-opus-4-7')\nhandle.name",
		);
		expect(spawned.status).toBe("ok");
		expect(spawned.result).toBe("'worker'");
		expect(requestTypes).toEqual(["rlm.run"]);

		const nameless = await manager.execute("await rlm.spawn('child work')");
		expect(nameless.status).toBe("error");
		expect(nameless.error?.ename).toBe("TypeError");
		expect(nameless.error?.evalue).toContain("required keyword-only arguments");
		expect(nameless.error?.evalue).toContain("'name'");
		expect(nameless.error?.evalue).toContain("'model'");
		expect(requestTypes).toEqual(["rlm.run"]);

		const called = await manager.execute("await rlm('child work')");
		expect(called.status).toBe("error");
		expect(called.error?.ename).toBe("TypeError");
		expect(called.error?.evalue).toContain("await rlm.spawn(");
	}, 30_000);

	it("dispose sends the protocol shutdown so live bash children die with the kernel", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const r = await manager.execute("from rlm import bash\nh = bash('sleep 600')\nh.pid");
		expect(r.status).toBe("ok");
		const pid = Number(r.result);
		expect(Number.isInteger(pid)).toBe(true);
		await manager.shutdown({ snapshot: true, drainHostRequests: true });
		let alive = true;
		for (let i = 0; i < 100; i++) {
			try {
				process.kill(pid, 0);
			} catch {
				alive = false;
				break;
			}
			await new Promise((resolve) => setTimeout(resolve, 50));
		}
		expect(alive).toBe(false);
	}, 30_000);

	it("stops a cell at its deadline, keeps partial output, and keeps serving cells", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const r = await manager.execute("import time\nprint('before-the-deadline', flush=True)\ntime.sleep(5)", {
			timeoutMs: 1500,
		});
		expect(r.status).toBe("aborted");
		expect(r.timedOut).toBe(true);
		expect(r.stdout).toContain("before-the-deadline");
		expect(r.durationMs).toBeGreaterThanOrEqual(1500);

		// The kernel survived the interrupt, and so did its namespace.
		const again = await manager.execute("x = 41\nx + 1");
		expect(again.status).toBe("ok");
		expect(again.result).toBe("42");
	}, 30_000);

	it("ends a cell blocked on a one-shot await bash() at its deadline", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const pidFile = join(dir, "child.pid");
		// The audited 61-minute blocking await, now bounded by the cell deadline.
		const r = await manager.execute(
			`from rlm import bash\nawait bash(${JSON.stringify(`echo $$ > ${pidFile}; sleep 600`)})\nprint('never reached', flush=True)`,
			{ timeoutMs: 2000 },
		);
		expect(r.status).toBe("aborted");
		expect(r.timedOut).toBe(true);
		expect(r.stdout).not.toContain("never reached");

		// `await bash(cmd)` owns its process: the cancelled await kills the group.
		const pid = Number(readFileSync(pidFile, "utf8").trim());
		expect(Number.isInteger(pid)).toBe(true);
		let alive = true;
		for (let i = 0; i < 60; i++) {
			try {
				process.kill(pid, 0);
			} catch {
				alive = false;
				break;
			}
			await new Promise((resolve) => setTimeout(resolve, 100));
		}
		expect(alive).toBe(false);
	}, 30_000);

	it("runs past any deadline when the caller arms no timer", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		// Deliberately longer than the 1.5s deadline above: without timeoutMs there is no timer.
		const slow = await manager.execute("import time\ntime.sleep(2.0)\n'finished'");
		expect(slow.status).toBe("ok");
		expect(slow.result).toBe("'finished'");
		expect(slow.timedOut).toBeUndefined();

		// State ops keep their own explicit bound instead of inheriting a cell deadline.
		const internal = await manager.execute("'state-op'", { internal: true });
		expect(internal.status).toBe("ok");
		expect(internal.result).toBe("'state-op'");
		expect(internal.timedOut).toBeUndefined();
	}, 30_000);

	it("surfaces unattributed background output separately from cell stdout", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const first = await manager.execute(
			[
				"import threading, time",
				"def late():",
				"    time.sleep(0.5)",
				"    print('SECRET-thread', flush=True)",
				"threading.Thread(target=late, daemon=True).start()",
			].join("\n"),
		);
		expect(first.status).toBe("ok");

		const second = await manager.execute("import time\ntime.sleep(1.0)\nprint('own-output')");
		expect(second.status).toBe("ok");
		expect(second.stdout).toContain("own-output");
		expect(second.stdout).not.toContain("SECRET-thread");
		expect(second.backgroundOutput ?? "").toContain("SECRET-thread");
	}, 30_000);
});