import { EventEmitter } from "node:events";
import { readFileSync } from "node:fs";
import { resolve } from "node:path";
import { describe, expect, it, vi } from "vitest";
import { type HostRequestHandlers, ReplKernelManager } from "../src/core/kernel/index.js";

type ShutdownInternals = {
	state: "running";
	writeLine: (request: Record<string, unknown>) => Promise<void>;
	handleEvent: (event: Record<string, unknown>) => void;
	wireChild: (child: ShutdownInternals["child"]) => void;
	pendingDoneWaiters: Map<string, () => void>;
	inFlightHostRequests: Set<Promise<void>>;
	kernelStderr: string;
	child: EventEmitter & {
		exitCode: number | null;
		signalCode: NodeJS.Signals | null;
		kill: (signal?: NodeJS.Signals | number) => boolean;
		pid?: number;
		stdin: {
			destroyed: boolean;
			destroy: () => void;
			on: (event: string, listener: (...args: unknown[]) => void) => void;
		};
		stdout?: { destroy: () => void; on: (event: string, listener: (...args: unknown[]) => void) => void };
		stderr?: {
			destroy: () => void;
			on: (event: string, listener: (...args: unknown[]) => void) => void;
			once: (event: string, listener: (...args: unknown[]) => void) => void;
		};
	};
};

function configuredManager(
	onSend: (request: Record<string, unknown>, internals: ShutdownInternals) => void | Promise<void>,
	hostHandlers?: HostRequestHandlers,
	withSnapshot = false,
): {
	manager: ReplKernelManager;
	internals: ShutdownInternals;
} {
	const manager = new ReplKernelManager({
		cwd: process.cwd(),
		hostHandlers,
		snapshot: withSnapshot ? { path: "/tmp/shutdown-test.dill", manifestPath: "/tmp/shutdown-test.json" } : undefined,
	});
	const internals = manager as unknown as ShutdownInternals;
	const child = Object.assign(new EventEmitter(), {
		exitCode: null,
		signalCode: null,
		kill: vi.fn(() => true),
		pid: undefined,
		stdin: { destroyed: false, destroy: vi.fn(), on: vi.fn() },
		stdout: { destroy: vi.fn(), on: vi.fn() },
		stderr: { destroy: vi.fn(), on: vi.fn(), once: vi.fn() },
	});
	Object.assign(internals, {
		state: "running",
		writeLine: vi.fn(async (request: Record<string, unknown>) => onSend(request, internals)),
		child,
	});
	// Attach the real exit/error handlers so teardown-ownership races are exercised.
	internals.wireChild(child);
	return { manager, internals };
}

describe("ReplKernelManager graceful shutdown", () => {
	it.each([
		["optional policy", {}],
		["snapshot-and-drain policy", { snapshot: true, drainHostRequests: true }],
	] as const)("uses the same diagnostic deadline for the %s", async (_label, options) => {
		vi.useFakeTimers();
		try {
			const { manager, internals } = configuredManager(() => new Promise<void>(() => {}), undefined, true);
			// The final flush snapshots through the private captureSnapshot (bounded, prune-free).
			vi.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot").mockResolvedValue(
				null,
			);
			const shutdown = manager.shutdown(options);
			await vi.advanceTimersByTimeAsync(5_000);
			await expect(shutdown).resolves.toBe(true);
			expect(internals.kernelStderr).toContain("Kernel did not shut down within 5000ms");
			expect(internals.child).toBeUndefined();
		} finally {
			vi.useRealTimers();
		}
	});

	it("does not finish shutdown before the stdin write settles", async () => {
		let finishSend: (() => void) | undefined;
		const sendBlocked = new Promise<void>((resolve) => {
			finishSend = resolve;
		});
		const { manager, internals } = configuredManager(async (request, state) => {
			state.handleEvent({ event: "done", id: request.id, status: "ok" });
			await sendBlocked;
			state.child.exitCode = 0;
			state.child.emit("exit", 0, null);
		});

		let finished = false;
		const shutdown = manager.shutdown().then(() => {
			finished = true;
		});
		await new Promise((resolve) => globalThis.setTimeout(resolve, 0));
		expect(finished).toBe(false);
		finishSend?.();
		await shutdown;
		expect(internals.pendingDoneWaiters.size).toBe(0);
	});

	it("finishes promptly when the runtime exits without a shutdown done", async () => {
		const { manager, internals } = configuredManager((_request, state) => {
			state.child.exitCode = 0;
			state.child.emit("exit", 0, null);
		});
		vi.useFakeTimers();
		try {
			const shutdown = manager.shutdown();
			await vi.advanceTimersByTimeAsync(100);
			// True = this call performed the cleanup: startup-failure recovery relies on it to resurrect to idle.
			await expect(shutdown).resolves.toBe(true);
			expect(internals.child).toBeUndefined();
		} finally {
			vi.useRealTimers();
		}
	});

	it("flushes a snapshot and drains host requests when session teardown selects both policies", async () => {
		let releaseHandler: (() => void) | undefined;
		const handlerBlocked = new Promise<void>((resolve) => {
			releaseHandler = resolve;
		});
		const sentReplies: Record<string, unknown>[] = [];
		const { manager, internals } = configuredManager(
			async (request, state) => {
				if (request.type === "host_reply") sentReplies.push(request);
				if (request.type === "shutdown") {
					state.handleEvent({ event: "done", id: request.id, status: "ok" });
					state.child.exitCode = 0;
					state.child.emit("exit", 0, null);
				}
			},
			{
				"test.slow": async () => {
					await handlerBlocked;
					return { answer: 42, status: "weird" };
				},
			},
			true,
		);
		const snapshot = vi
			.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot")
			.mockResolvedValue(null);

		internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } });
		const shutdown = manager.shutdown({ snapshot: true, drainHostRequests: true });
		await new Promise((resolve) => globalThis.setTimeout(resolve, 0));

		expect(snapshot).toHaveBeenCalledOnce();
		expect(internals.writeLine).not.toHaveBeenCalled();
		releaseHandler?.();
		await shutdown;

		expect(internals.inFlightHostRequests.size).toBe(0);
		expect(sentReplies).toEqual([
			{
				type: "host_reply",
				id: "hr-1",
				data: { status: "ok", result: { answer: 42, status: "weird" } },
			},
		]);
		expect(internals.child).toBeUndefined();
	});

	it("shutdown during an in-flight host request leaves no dangling task once the handler settles", async () => {
		let releaseHandler: (() => void) | undefined;
		const handlerBlocked = new Promise<void>((resolve) => {
			releaseHandler = resolve;
		});
		const { manager, internals } = configuredManager(
			(request, state) => {
				if (request.type !== "shutdown") return;
				state.handleEvent({ event: "done", id: request.id, status: "ok" });
				state.child.exitCode = 0;
				state.child.emit("exit", 0, null);
			},
			{
				"test.slow": async () => {
					await handlerBlocked;
					return { answer: 42 };
				},
			},
			true,
		);
		const snapshot = vi
			.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot")
			.mockResolvedValue(null);

		internals.handleEvent({ event: "host_request", id: "hr-1", data: { type: "test.slow" } });
		const tracked = [...internals.inFlightHostRequests];
		expect(tracked).toHaveLength(1);

		await manager.shutdown();
		expect(snapshot).not.toHaveBeenCalled();
		expect(internals.child).toBeUndefined();

		// The reply write fails against the torn-down child; the task must still settle.
		releaseHandler?.();
		await Promise.all(tracked);
		expect(internals.inFlightHostRequests.size).toBe(0);
	});

	it("shares one wire exchange across concurrent graceful shutdown calls", async () => {
		let exitKernel: (() => void) | undefined;
		const { manager, internals } = configuredManager((request, state) => {
			state.handleEvent({ event: "done", id: request.id, status: "ok" });
			exitKernel = () => {
				state.child.exitCode = 0;
				state.child.emit("exit", 0, null);
			};
		});

		const owner = manager.shutdown();
		await new Promise((resolve) => globalThis.setTimeout(resolve, 0));
		const concurrent = manager.shutdown();
		exitKernel?.();

		await expect(owner).resolves.toBe(true);
		await expect(concurrent).resolves.toBe(false);
		expect(internals.writeLine).toHaveBeenCalledOnce();
	});

	it("stands down when a hard teardown supersedes the snapshot flush", async () => {
		let releaseSnapshot: (() => void) | undefined;
		const snapshotBlocked = new Promise<null>((resolve) => {
			releaseSnapshot = () => resolve(null);
		});
		const { manager, internals } = configuredManager(() => undefined, undefined, true);
		vi.spyOn(manager as unknown as { captureSnapshot: () => Promise<null> }, "captureSnapshot").mockReturnValue(
			snapshotBlocked,
		);
		const child = internals.child;

		const shutdown = manager.shutdown({ snapshot: true });
		await new Promise((resolve) => globalThis.setTimeout(resolve, 0));
		await manager.kill();
		releaseSnapshot?.();

		await expect(shutdown).resolves.toBe(false);
		expect(child.kill).toHaveBeenCalledOnce();
	});

	it("restart after a graceful shutdown starts the kernel again", async () => {
		const { manager, internals } = configuredManager((request, state) => {
			if (request.type !== "shutdown") return;
			state.handleEvent({ event: "done", id: request.id, status: "ok" });
			// The runtime exits after acking shutdown; the exit handler must not
			// steal teardown ownership from the in-flight graceful shutdown().
			state.child.exitCode = 0;
			state.child.emit("exit", 0, null);
		});
		const start = vi.spyOn(manager, "start").mockResolvedValue(undefined);

		await manager.restart();

		expect(internals.child).toBeUndefined();
		expect(start).toHaveBeenCalledTimes(1);
	});

	it("restart does not resurrect a kernel a concurrent kill superseded", async () => {
		const { manager } = configuredManager((_request, state) => {
			// A concurrent kill supersedes the in-flight graceful shutdown.
			void manager.kill();
			state.child.emit("exit", null, "SIGKILL");
		});
		const start = vi.spyOn(manager, "start").mockResolvedValue(undefined);

		await manager.restart();

		expect(start).not.toHaveBeenCalled();
	});

	it("kill() destroys stderr only while the child has not exited", async () => {
		// Alive child: it may ignore the signal and never emit 'exit', so the
		// post-exit drain would never bound the pipe — teardown must.
		const alive = configuredManager(() => {});
		const aliveChild = alive.internals.child;
		await alive.manager.kill();
		expect(aliveChild.stderr?.destroy).toHaveBeenCalledTimes(1);

		// Exited child: the post-exit drain owns the stream; destroying it in
		// teardown would drop the kernel's buffered last words.
		const exited = configuredManager(() => {});
		const exitedChild = exited.internals.child;
		exitedChild.exitCode = 42;
		await exited.manager.kill();
		expect(exitedChild.stderr?.destroy).not.toHaveBeenCalled();
	});

	it("waits for the matching shutdown done and removes its waiter", async () => {
		const { manager, internals } = configuredManager(async (request, state) => {
			expect(request.type).toBe("shutdown");
			expect(request.id).toBeTypeOf("string");
			// An unrelated done must not release the shutdown waiter.
			state.handleEvent({ event: "done", id: "unrelated", status: "ok" });
			expect(state.pendingDoneWaiters.size).toBe(1);
			state.handleEvent({ event: "done", id: request.id, status: "ok" });
			queueMicrotask(() => {
				state.child.exitCode = 0;
				state.child.emit("exit", 0, null);
			});
		});

		await manager.shutdown();

		expect(internals.pendingDoneWaiters.size).toBe(0);
		expect(internals.child).toBeUndefined();
	});

	it("keeps the kernel MCP close budget strictly inside the host shutdown deadline", () => {
		const source = readFileSync(resolve(__dirname, "../../../prime-agent-runtime/src/rlm/mcp.py"), "utf8");
		const match = source.match(/^_SHUTDOWN_TIMEOUT = ([\d.]+)$/m);
		expect(match).not.toBeNull();
		// +1s dispatch slack in mcp.py close(); the sum must undercut the host's 5s kill deadline.
		expect((Number(match![1]) + 1) * 1000).toBeLessThan(5000);
	});
});