import { spawnSync } from "node:child_process";
import { existsSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { homedir, tmpdir } from "node:os";
import { join, resolve } from "node:path";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { ReplKernelManager } from "../src/core/kernel/index.js";

const runtimePython = resolve("../../prime-agent-runtime/.venv/bin/python");
const fallbackPython = join(homedir(), ".prime", "agent", "kernel-venv", "bin", "python");

function resolveKernelPython(): string | null {
	for (const python of [process.env.PRIME_AGENT_KERNEL_PYTHON, runtimePython, fallbackPython]) {
		if (!python || !existsSync(python)) continue;
		const check = spawnSync(python, ["-c", "import rlm.repl, mcp, rlm"], { encoding: "utf8" });
		if (check.status === 0) return python;
	}
	return null;
}

const python = resolveKernelPython();
const describeIfKernel = python ? describe : describe.skip;

const MCP_SERVER = `import asyncio, json, os, sys
from pathlib import Path
Path(sys.argv[1]).write_text(str(os.getpid()))
async def main():
    while line := await asyncio.get_running_loop().run_in_executor(None, sys.stdin.readline):
        request = json.loads(line)
        if request.get("id") is None:
            continue
        if request["method"] == "initialize":
            result = {"protocolVersion": "2025-06-18", "capabilities": {"tools": {}}, "serverInfo": {"name": "shutdown-fixture", "version": "1"}}
        elif request["method"] == "tools/list":
            result = {"tools": [{"name": "fixture.echo", "description": "echo", "inputSchema": {"type": "object"}}]}
        else:
            result = {"content": [{"type": "text", "text": "ok"}]}
        print(json.dumps({"jsonrpc": "2.0", "id": request["id"], "result": result}), flush=True)
asyncio.run(main())
`;

function pidExists(pid: number): boolean {
	try {
		process.kill(pid, 0);
		return true;
	} catch {
		return false;
	}
}

async function waitForExit(pid: number, timeoutMs: number): Promise<boolean> {
	const deadline = Date.now() + timeoutMs;
	while (Date.now() < deadline) {
		if (!pidExists(pid)) return true;
		await new Promise((resolveWait) => globalThis.setTimeout(resolveWait, 25));
	}
	return !pidExists(pid);
}

describeIfKernel("real REPL kernel MCP shutdown", { tags: ["kernel-heavy"] }, () => {
	let dir = "";
	let fixture = "";
	let pidFile = "";

	beforeAll(() => {
		dir = mkdtempSync(join(tmpdir(), "prime-agent-repl-mcp-shutdown-"));
		fixture = join(dir, "stdio_server.py");
		pidFile = join(dir, "stdio.pid");
		writeFileSync(fixture, MCP_SERVER);
	});

	afterAll(() => {
		if (dir) rmSync(dir, { recursive: true, force: true });
	});

	it("closes a stdio MCP child on graceful shutdown", async () => {
		let manager: ReplKernelManager | undefined = new ReplKernelManager({
			python: python as string,
			cwd: resolve("../../prime-agent-runtime"),
			hostHandlers: {
				"mcp.config": async () => ({
					type: "stdio",
					command: python as string,
					args: [fixture, pidFile],
				}),
			},
		});
		try {
			const opened = await manager.execute(
				"import rlm.mcp as mcp; tools = await mcp.list_tools('fixture'); [t['name'] for t in tools]",
			);
			expect(opened.status, opened.stderr || opened.error?.traceback.join("\n")).toBe("ok");
			expect(opened.result).toContain("fixture.echo");
			const childPid = Number(readFileSync(pidFile, "utf8"));
			expect(pidExists(childPid)).toBe(true);

			await manager.shutdown();
			expect(await waitForExit(childPid, 2_000)).toBe(true);
			manager = undefined;
		} finally {
			await manager?.kill();
		}
	});
});