import { spawnSync } from "node:child_process";
import { existsSync, mkdtempSync, 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,
	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("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("remains usable after a cell raises SystemExit", async () => {
		manager = new ReplKernelManager({ python: python as string, cwd: dir });
		const exited = await manager.execute("raise SystemExit(3)");
		expect(exited.status).toBe("error");
		expect(exited.error?.ename).toBe("SystemExit");

		await expect(manager.execute("1 + 1")).resolves.toMatchObject({ status: "ok", result: "2" });
	}, 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(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.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("emits one background bash completion and acknowledges it when observed", async () => {
		const events: Array<{ kind: string; handleId: string; pid?: number; exitCode?: number }> = [];
		let resolveCompletion = () => {};
		const completion = new Promise<void>((resolve) => {
			resolveCompletion = resolve;
		});
		manager = new ReplKernelManager({
			python: python as string,
			cwd: dir,
			onBashHandleEvent: (event) => {
				events.push(event);
				if (event.kind === "completed") resolveCompletion();
			},
		});

		const started = await manager.execute("from rlm import bash\nh = bash('sleep 0.2; exit 7')\nh.pid");
		expect(started.status).toBe("ok");
		await completion;
		expect(events).toHaveLength(1);
		expect(events[0]).toMatchObject({ kind: "completed", pid: Number(started.result), exitCode: 7 });

		const observed = await manager.execute("h.poll().exit_code");
		expect(observed.result).toBe("7");
		expect(events).toEqual([
			expect.objectContaining({ kind: "completed" }),
			{ kind: "observed", handleId: events[0].handleId },
		]);
	}, 30_000);

	it("defers a background completion until the active host request settles", async () => {
		const events: Array<{ kind: string; handleId: string }> = [];
		let releaseHandler = () => {};
		const handlerGate = new Promise<void>((resolve) => {
			releaseHandler = resolve;
		});
		let markHandlerStarted = () => {};
		const handlerStarted = new Promise<void>((resolve) => {
			markHandlerStarted = resolve;
		});
		manager = new ReplKernelManager({
			python: python as string,
			cwd: dir,
			hostHandlers: {
				"test.block": async () => {
					markHandlerStarted();
					await handlerGate;
					return { released: true };
				},
			},
			onBashHandleEvent: (event) => events.push(event),
		});

		await manager.execute("from rlm import bash\nh = bash('sleep 0.2; exit 7')\nh.pid");
		const blocked = manager.execute("import rlm\nawait rlm.host_request('test.block')");
		await handlerStarted;
		await new Promise((resolve) => setTimeout(resolve, 500));
		expect(events).toEqual([]);

		releaseHandler();
		await expect(blocked).resolves.toMatchObject({ status: "ok" });
		expect(events).toEqual([expect.objectContaining({ kind: "completed" })]);
	}, 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("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("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);
});