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

describe("RLM bootstrap", () => {
	it("pre-imports asyncio so the prompt's subagent patterns work without a manual import", () => {
		expect(buildRlmBootstrapCode()).toMatch(/^import asyncio$/m);
	});

	it("gives subagent registry operations the actionable missing-runtime fallback", () => {
		const code = buildRlmBootstrapCode();
		// Role routing was replaced by `rlm.models()` discovery, so the stub set is the spawn
		// and registry surface the prompt actually calls.
		expect(code).toContain("async def spawn(self, prompt, **kwargs)");
		expect(code).toContain("async def models(self)");
		expect(code).toContain("async def list_subagents(self)");
		expect(code).toContain("async def delete_subagent(self, target)");
		expect(code).toContain("self._rebind()");
	});

	it("disables colored output for subprocesses launched by the kernel", () => {
		expect(buildRlmBootstrapCode()).toContain('_prime_agent_os.environ["NO_COLOR"] = "1"');
	});

	it("binds bash from the runtime with a retrying missing-runtime stub fallback", () => {
		const code = buildRlmBootstrapCode();
		expect(code).toContain("bash=_prime_agent_rlm_module.bash");
		// The stub accepts the real call shape, so a cell that passes a timeout gets the
		// runtime's own signature check rather than the stub's arity.
		expect(code).toContain("def bash(command, *args, **kwargs):");
		expect(code).toContain('_prime_agent_require_runtime()["bash"](command, *args, **kwargs)');
	});

	it("does not install a kernel shutdown hook (the runtime closes MCP on shutdown itself)", () => {
		expect(buildRlmBootstrapCode()).not.toContain("install_shutdown_hook");
	});

	it("guards Python skill imports so a broken skill does not abort bootstrap", () => {
		const code = buildRlmBootstrapCode([
			{
				name: "broken-skill",
				importName: "broken_skill",
				packagePath: "/tmp/broken-skill",
				pyprojectPath: "/tmp/broken-skill/pyproject.toml",
			},
		]);

		expect(code).toContain("except Exception as _prime_agent_skill_error");
		expect(code).toContain("_PrimeAgentUnavailableSkill");
		expect(code).toContain("_PRIME_AGENT_SKILL_IMPORT_ERRORS");
		expect(code).toContain("globals()[_prime_agent_skill_name] = _PrimeAgentUnavailableSkill");
	});
});

const FAKE_RUNTIME_SOURCE = `
rlm = object()

def bash(command, *args, **kwargs):
    return f"real:{command}:{sorted(kwargs)}"
`;

describe("RLM bootstrap recovery", () => {
	const python = process.env.PRIME_AGENT_KERNEL_PYTHON ?? "python3";
	const dir = mkdtempSync(join(tmpdir(), "prime-agent-rlm-recovery-"));

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

	it("waits out a build lock, then rebinds the runtime on the next call", async () => {
		const site = join(dir, "site");
		const venv = join(dir, "kernel-venv");
		const lock = `${venv}.bootstrap.lock`;
		const missing = join(site, "rlm", "mcp.py");
		mkdirSync(join(site, "rlm"), { recursive: true });
		writeFileSync(join(site, "rlm", "__init__.py"), FAKE_RUNTIME_SOURCE);
		writeFileSync(join(site, "remote.py"), "MARK = 'remote'\n");
		writeFileSync(lock, `${process.pid}\n`);
		const script = [
			buildRlmBootstrapCode([{ name: "flaky", importName: "flaky", packagePath: site, pyprojectPath: lock }]),
			'print("lock_held_after_wait:", _prime_agent_build_in_progress(), flush=True)',
			'print("ready:", _PRIME_AGENT_RUNTIME_READY, flush=True)',
			"try:",
			'    bash("too-early")',
			"except RuntimeError as error:",
			'    print("terminal:", "prime-agent-runtime is not installed" in str(error), flush=True)',
			'print("awaiting-go", flush=True)',
			"input()",
			'print("bash:", bash("echo", timeout=5), flush=True)',
			'print("mcp:", mcp.MARK, remote.MARK, flush=True)',
			'print("skill:", asyncio.run(flaky.run()), flush=True)',
		].join("\n");
		const child = spawn(python, ["-c", script], {
			env: {
				...process.env,
				PYTHONPATH: site,
				PRIME_AGENT_KERNEL_VENV: venv,
				PRIME_AGENT_KERNEL_BOOTSTRAP_WAIT_SECONDS: "10",
			},
		});
		let stdout = "";
		child.stdout.setEncoding("utf8");
		child.stdout.on("data", (chunk: string) => {
			stdout += chunk;
			// Act only on the kernel's own progress: release the lock it is waiting on,
			// then publish the module whose absence failed its first import.
			if (stdout.includes("waiting up to")) rmSync(lock, { force: true });
			if (stdout.includes("awaiting-go")) {
				writeFileSync(missing, "MARK = 'mcp'\n");
				// The skill's editable install landed with the same rebuild.
				writeFileSync(join(site, "flaky.py"), "async def run():\n    return 'skill-ok'\n");
				child.stdin.write("go\n");
			}
		});
		const exitCode = await new Promise<number | null>((settle) => child.on("close", settle));

		expect(exitCode).toBe(0);
		expect(stdout).toContain("lock_held_after_wait: False");
		expect(stdout).toContain("waiting up to 10s for the kernel venv build");
		expect(stdout).toContain("ready: False");
		expect(stdout).toContain("terminal: True");
		expect(stdout).toContain("bash: real:echo:['timeout']");
		expect(stdout).toContain("mcp: mcp remote");
		expect(stdout).toContain("skill: skill-ok");
	});

	it("treats a lock with an alive builder pid as in-progress beyond stale age and clears dead locks", () => {
		const venv = join(dir, "pid-lock-venv");
		const lock = `${venv}.bootstrap.lock`;
		const script = `${buildRlmBootstrapCode()}\nimport os, time\nopen('${lock}', 'w').write(f'{os.getpid()}:tok\\n')\nos.utime('${lock}', (time.time() - 60, time.time() - 60))\nassert _prime_agent_build_in_progress() is True\nopen('${lock}', 'w').write('999999999:dead\\n')\nassert _prime_agent_build_in_progress() is False and not os.path.exists('${lock}')`;
		expect(spawnSync(python, ["-c", script], { env: { ...process.env, PRIME_AGENT_KERNEL_VENV: venv } }).status).toBe(
			0,
		);
	});

	it("isolates optional remote import failure so rlm, bash, and mcp remain available", () => {
		const site = join(dir, "isolated-site");
		mkdirSync(join(site, "rlm"), { recursive: true });
		writeFileSync(join(site, "rlm", "__init__.py"), FAKE_RUNTIME_SOURCE);
		writeFileSync(join(site, "rlm", "mcp.py"), "MARK = 'mcp'\n");
		const script = `${buildRlmBootstrapCode()}\nassert _PRIME_AGENT_RUNTIME_READY and callable(bash) and hasattr(mcp, 'MARK') and isinstance(globals().get('remote'), _PrimeAgentUnavailableSkill)`;
		expect(spawnSync(python, ["-c", script], { env: { ...process.env, PYTHONPATH: site } }).status).toBe(0);
	});

	it("placeholder objects self-heal on attribute access and report teaching error on failure", () => {
		const venv = join(dir, "placeholder-venv");
		const script = `${buildRlmBootstrapCode([{ name: "git", importName: "git", packagePath: "/tmp/git", pyprojectPath: "/tmp/git/pyproject.toml" }])}\ntry: rlm.progress_note('t')\nexcept RuntimeError as e: assert 'kernel-venv' in str(e) or '${venv}' in str(e)\ntry: git.status()\nexcept RuntimeError as e: assert 'is unavailable in this kernel' in str(e) and '${venv}' in str(e)`;
		expect(spawnSync(python, ["-c", script], { env: { ...process.env, PRIME_AGENT_KERNEL_VENV: venv } }).status).toBe(
			0,
		);
	});
});

/** Find a python with a current rlm runtime, or null to skip. */
function resolveKernelPython(): 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, rlm; assert callable(rlm.emit)"], {
			encoding: "utf8",
		});
		if (check.status === 0) return python;
	}
	return null;
}

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

describeIfKernel("RLM bootstrap (real kernel)", () => {
	const dir = mkdtempSync(join(tmpdir(), "prime-agent-bootstrap-"));

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

	it("binds asyncio in the user namespace", async () => {
		const manager = new ReplKernelManager({ python: python as string, cwd: dir });
		try {
			await manager.start();
			const bootstrap = await manager.execute(buildRlmBootstrapCode());
			expect(bootstrap.status).toBe("ok");

			const result = await manager.execute("_t = asyncio.create_task(asyncio.sleep(0))\nprint(type(_t).__name__)");
			expect(result.status).toBe("ok");
			expect(result.stdout).toContain("Task");

			const envResult = await manager.execute('import os\nprint(os.environ["NO_COLOR"])');
			expect(envResult.status).toBe("ok");
			expect(envResult.stdout.trim()).toBe("1");
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
		}
	}, 60_000);

	it("emits canonical paths for edits after the kernel changes directories", async () => {
		const firstDir = join(dir, "first");
		const secondDir = join(dir, "second");
		mkdirSync(firstDir, { recursive: true });
		mkdirSync(secondDir, { recursive: true });
		writeFileSync(join(firstDir, "same.txt"), "old");
		writeFileSync(join(secondDir, "same.txt"), "old");
		const editSkillRoot = join(process.cwd(), "skills", "edit");
		const manager = new ReplKernelManager({
			python: python as string,
			cwd: dir,
			env: { PYTHONPATH: join(editSkillRoot, "src") },
		});
		try {
			await manager.start();
			const bootstrap = await manager.execute(
				buildRlmBootstrapCode([
					{
						name: "edit",
						importName: "edit",
						packagePath: editSkillRoot,
						pyprojectPath: join(editSkillRoot, "pyproject.toml"),
					},
				]),
			);
			expect(bootstrap.status).toBe("ok");

			const first = await manager.execute(
				'import os\nos.chdir("first")\nawait edit(path="same.txt", old_str="old", new_str="new")',
			);
			const second = await manager.execute(
				'os.chdir("../second")\nawait edit(path="same.txt", old_str="old", new_str="new")',
			);

			expect(first.diffs?.[0]?.path).toBe(realpathSync(join(firstDir, "same.txt")));
			expect(second.diffs?.[0]?.path).toBe(realpathSync(join(secondDir, "same.txt")));
			expect(first.diffs?.[0]?.path).not.toBe(second.diffs?.[0]?.path);
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
		}
	}, 60_000);
});