import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import { ReplKernelManager } from "../src/core/kernel/index.js";
import {
	kernelCwdPathIn,
	manifestPathIn,
	readKernelCwd,
	snapshotPathIn,
	writeKernelCwd,
} from "../src/core/kernel/state-snapshot.js";

describe("kernel state snapshot paths", () => {
	it("places snapshot + manifest inside the session artifact directory", () => {
		const artifactDir = "/home/u/.prime/agent/session-artifacts/abc-123";
		expect(snapshotPathIn(artifactDir)).toBe(join(artifactDir, "kernel-state.dill"));
		expect(manifestPathIn(artifactDir)).toBe(join(artifactDir, "kernel-state.json"));
		expect(kernelCwdPathIn(artifactDir)).toBe(join(artifactDir, "kernel-cwd"));
	});
});

describe("recorded kernel working directory", () => {
	const dirs: string[] = [];

	function artifactDir(): string {
		const dir = mkdtempSync(join(tmpdir(), `pi-kernel-cwd-${process.pid}-`));
		dirs.push(dir);
		return dir;
	}

	afterEach(() => {
		while (dirs.length > 0) rmSync(dirs.pop() as string, { recursive: true, force: true });
	});

	it("round-trips a directory that exists", () => {
		const dir = artifactDir();
		const work = join(dir, "work");
		mkdirSync(work);

		writeKernelCwd(dir, work);

		expect(readKernelCwd(dir)).toEqual({ cwd: work });
	});

	it("reports a recorded directory that is gone instead of using it", () => {
		const dir = artifactDir();
		const gone = join(dir, "gone");
		mkdirSync(gone);
		writeKernelCwd(dir, gone);
		rmSync(gone, { recursive: true, force: true });

		expect(readKernelCwd(dir)).toEqual({ missing: gone });
	});

	it("ignores a relative path, an empty record, and a missing record", () => {
		const dir = artifactDir();
		writeKernelCwd(dir, "relative/path");
		expect(readKernelCwd(dir)).toEqual({});

		writeFileSync(kernelCwdPathIn(dir), "\n");
		expect(readKernelCwd(dir)).toEqual({});

		rmSync(kernelCwdPathIn(dir));
		expect(readKernelCwd(dir)).toEqual({});
	});
});

type ReplInternals = {
	state: string;
	writeLine: (request: Record<string, unknown>) => Promise<void>;
	start: () => Promise<void>;
	handleEvent: (event: Record<string, unknown>) => void;
	activeExecution?: { requestId: string };
	scheduleSnapshot: () => void;
};

/** A manager whose kernel is a stub that answers state requests by hand. */
function stubManager(): {
	manager: ReplKernelManager;
	internals: ReplInternals;
	requests: Record<string, unknown>[];
	requestTypes: () => (string | undefined)[];
} {
	const manager = new ReplKernelManager({
		cwd: process.cwd(),
		snapshot: { path: "/tmp/test-state.dill", manifestPath: "/tmp/test-state.json" },
	});
	const requests: Record<string, unknown>[] = [];
	const internals = manager as unknown as ReplInternals;
	Object.assign(internals, {
		state: "running",
		start: async () => {},
		writeLine: async (request: Record<string, unknown>) => {
			requests.push(request);
			if (request.type === "restore" && typeof request.id === "string") {
				internals.handleEvent({ event: "done", id: request.id, status: "ok", restored: ["alpha"], failed: [] });
			}
		},
	});
	return { manager, internals, requests, requestTypes: () => requests.map((request) => String(request.type)) };
}

async function settleExecution(internals: ReplInternals): Promise<void> {
	for (let i = 0; i < 20 && !internals.activeExecution; i += 1) {
		await Promise.resolve();
	}
}

describe("kernel state snapshots across a cell deadline", () => {
	afterEach(() => {
		vi.useRealTimers();
	});

	it("never snapshots the timed-out cell's namespace", async () => {
		vi.useFakeTimers();
		const { manager, internals, requestTypes } = stubManager();
		const scheduleSnapshot = vi.fn();
		Object.assign(internals, { scheduleSnapshot });

		const executePromise = manager.execute("while True: pass", { timeoutMs: 1000 });
		await settleExecution(internals);
		const execution = internals.activeExecution;
		if (!execution) {
			throw new Error("Expected an active execution");
		}
		internals.handleEvent({ event: "stdout", id: execution.requestId, text: "before\n" });

		// Deadline at 1000ms, grace at 2000ms: the cell is reported as aborted.
		await vi.advanceTimersByTimeAsync(2000);
		const result = await executePromise;
		expect(result.status).toBe("aborted");
		expect(result.timedOut).toBe(true);
		expect(result.stdout).toBe("before\n");

		// execute() schedules the post-cell snapshot only on status "ok": the on-disk
		// snapshot stays the last successful cell's, and the next start restores it.
		expect(scheduleSnapshot).not.toHaveBeenCalled();
		await vi.advanceTimersByTimeAsync(60_000);
		expect(requestTypes().filter((type) => type === "snapshot")).toEqual([]);

		// The next kernel over the same snapshot files restores the last successful
		// cell's namespace; the timed-out cell never contributed one.
		internals.handleEvent({ event: "done", id: execution.requestId, status: "error" });
		await vi.advanceTimersByTimeAsync(0);
		expect((await manager.restoreState())?.restored).toEqual(["alpha"]);
		manager.disposeSync();
	});

	it("schedules the snapshot for a cell that succeeded", async () => {
		const { manager, internals } = stubManager();
		const scheduleSnapshot = vi.fn();
		Object.assign(internals, { scheduleSnapshot });

		const executePromise = manager.execute("x = 41");
		await settleExecution(internals);
		const execution = internals.activeExecution;
		if (!execution) {
			throw new Error("Expected an active execution");
		}
		internals.handleEvent({ event: "done", id: execution.requestId, status: "ok" });

		await expect(executePromise).resolves.toMatchObject({ status: "ok" });
		expect(scheduleSnapshot).toHaveBeenCalledTimes(1);
		manager.disposeSync();
	});
});