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 { afterAll, beforeAll, describe, expect, it } from "vitest";
import { 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 describeIfKernel = python ? describe : describe.skip;

describeIfKernel("repl kernel state snapshot round-trip (real runtime)", { tags: ["kernel-heavy"] }, () => {
	let dir = "";
	let snapshotPath = "";
	let manifestPath = "";

	beforeAll(() => {
		dir = mkdtempSync(join(tmpdir(), "prime-agent-repl-roundtrip-"));
		snapshotPath = join(dir, "session.dill");
		manifestPath = join(dir, "session.json");
	});

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

	function newManager(): ReplKernelManager {
		return new ReplKernelManager({
			python: python as string,
			cwd: dir,
			snapshot: { path: snapshotPath, manifestPath },
		});
	}

	it("saves picklable names, reports unpicklable ones, then revives them in a fresh runtime", async () => {
		const writer = newManager();
		try {
			await writer.execute("x = 42");
			await writer.execute("df = [1, 2, 3]");
			await writer.execute("def double(n):\n    return n * 2");
			await writer.execute("gen = (n for n in range(3))");

			const snap = await writer.snapshotState();
			expect(snap).not.toBeNull();
			expect(snap?.saved).toEqual(expect.arrayContaining(["x", "df", "double"]));
			expect(snap?.skipped.map((s) => s.name)).toContain("gen");
			expect(existsSync(snapshotPath)).toBe(true);
			expect(existsSync(manifestPath)).toBe(true);
		} finally {
			await writer.shutdown({ snapshot: true, drainHostRequests: true });
		}

		const reader = newManager();
		try {
			const restore = await reader.restoreState();
			expect(restore?.restored).toEqual(expect.arrayContaining(["df", "double", "x"]));
			expect(restore?.failed.map((f) => f.name) ?? []).not.toContain("x");

			const echo = await reader.execute("print(x, double(x), sum(df))");
			expect(echo.stdout.trim()).toBe("42 84 6");
		} finally {
			await reader.shutdown({ snapshot: true, drainHostRequests: true });
		}
	}, 60_000);

	it("treats a missing snapshot as an empty restore (clean start)", async () => {
		const freshDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-state-empty-"));
		const manager = new ReplKernelManager({
			python: python as string,
			cwd: freshDir,
			snapshot: { path: join(freshDir, "missing.dill"), manifestPath: join(freshDir, "missing.json") },
		});
		try {
			const restore = await manager.restoreState();
			expect(restore).toEqual({ restored: [], failed: [], path: join(freshDir, "missing.dill") });
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
			rmSync(freshDir, { recursive: true, force: true });
		}
	}, 60_000);

	it("restores a snapshot artifact containing IPython-injected blobs, skipping those names", async () => {
		// Synthesize the artifact shape an IPython-kernel snapshot writes: a dict of
		// dill blobs including In/Out/get_ipython entries.
		const ipythonArtifactDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-ipython-artifact-"));
		const ipythonArtifactPath = join(ipythonArtifactDir, "kernel-state.dill");
		const buildScript = [
			"import dill",
			"dill.settings['recurse'] = True",
			"payload = {",
			"    'kept_number': dill.dumps(41),",
			"    'kept_text': dill.dumps('hello'),",
			"    'In': dill.dumps(['print(1)']),",
			"    'Out': dill.dumps({1: 'x'}),",
			"    'get_ipython': dill.dumps(None),",
			"}",
			`with open(${JSON.stringify(ipythonArtifactPath)}, "wb") as fh:`,
			"    dill.dump(payload, fh)",
		].join("\n");
		const build = spawnSync(python as string, ["-c", buildScript], { encoding: "utf8" });
		expect(build.status).toBe(0);

		const manager = new ReplKernelManager({
			python: python as string,
			cwd: ipythonArtifactDir,
			snapshot: { path: ipythonArtifactPath, manifestPath: join(ipythonArtifactDir, "kernel-state.json") },
		});
		try {
			const restore = await manager.restoreState();
			expect(restore?.restored).toEqual(["kept_number", "kept_text"]);
			expect(restore?.failed).toEqual([]);
			const echo = await manager.execute("print(kept_number + 1, kept_text, 'In' in dir(), 'Out' in dir())");
			expect(echo.stdout.trim()).toBe("42 hello False False");
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
			rmSync(ipythonArtifactDir, { recursive: true, force: true });
		}
	}, 60_000);

	it("lists live user-defined names, filtering internals and live handles", async () => {
		const listDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-state-list-"));
		const manager = new ReplKernelManager({ python: python as string, cwd: listDir });
		try {
			expect(await manager.listNamespaceNames()).toBeNull();
			await manager.execute("alpha = 1\ndef helper(n):\n    return n\n_hidden = 2\nrlm = object()");
			const names = await manager.listNamespaceNames();
			expect(names).toEqual(expect.arrayContaining(["alpha", "helper"]));
			expect(names).not.toContain("_hidden");
			expect(names).not.toContain("rlm");
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
			rmSync(listDir, { recursive: true, force: true });
		}
	}, 60_000);

	it("prunes oversized variables via a compaction snapshot", async () => {
		const boundedDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-state-bounded-"));
		const manager = new ReplKernelManager({
			python: python as string,
			cwd: boundedDir,
			snapshot: {
				path: join(boundedDir, "bounded.dill"),
				manifestPath: join(boundedDir, "bounded.json"),
				maxBytes: 10 * 1024,
				maxVariableBytes: 8 * 1024,
			},
		});
		try {
			await manager.execute('small_text = "a" * 100\nlarge_text = "x" * 16_384');
			const snapshot = await manager.snapshotState();
			expect(snapshot?.skipped.map(({ name }) => name)).toContain("large_text");
			expect(snapshot?.saved).toContain("small_text");

			const compacted = await manager.pruneOversizedVariables();
			expect(compacted?.pruned).toEqual(["large_text"]);
			const remaining = await manager.listNamespaceNames();
			expect(remaining).toContain("small_text");
			expect(remaining).not.toContain("large_text");
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
			rmSync(boundedDir, { recursive: true, force: true });
		}
	}, 60_000);

	it("auto-snapshots after a successful execution (debounced)", async () => {
		const autoDir = mkdtempSync(join(tmpdir(), "prime-agent-repl-state-auto-"));
		const autoPath = join(autoDir, "auto.dill");
		const manager = new ReplKernelManager({
			python: python as string,
			cwd: autoDir,
			snapshot: { path: autoPath, manifestPath: join(autoDir, "auto.json"), debounceMs: 50 },
		});
		try {
			await manager.execute("auto_var = 'persisted'");
			await expect.poll(() => existsSync(autoPath), { timeout: 10_000 }).toBe(true);
		} finally {
			await manager.shutdown({ snapshot: true, drainHostRequests: true });
			rmSync(autoDir, { recursive: true, force: true });
		}
	}, 60_000);
});