/**
 * KernelProcess — owns the Python REPL subprocess and the unix-socket
 * JSONL protocol spoken with kernel/rlm_kernel.py.
 *
 * One cell executes at a time (executions are mutex-queued). While a cell
 * runs, the kernel may send `stream` messages (stdout/stderr mirroring) and
 * blocking `rlm_request` messages; both are dispatched from the same reader.
 */

import { type ChildProcess, spawn } from "node:child_process";
import { rmSync } from "node:fs";
import { createServer, type Server, type Socket } from "node:net";
import { tmpdir } from "node:os";
import { join } from "node:path";

export interface VarInfo {
	name: string;
	type: string;
	repr: string;
}

export interface ExecResult {
	ok: boolean;
	stdout: string;
	stderr: string;
	/** repr() of a trailing expression, if any */
	result: string | null;
	/** the cell was interrupted via the abort signal */
	cancelled: boolean;
	/** the kernel was (re)booted for this cell — previous variables are gone */
	restarted: boolean;
	/** variables restored from the session snapshot (non-empty only on the first cell after a boot that restored) */
	restoredVars: string[];
}

export interface ExecOptions {
	signal?: AbortSignal;
	onStream?: (stream: "stdout" | "stderr", data: string) => void;
}

export interface RlmRequest {
	id: number;
	prompt: string;
	depth: number;
}

export interface RlmReply {
	ok: boolean;
	result?: string;
	error?: string;
}

export interface RefineRequest {
	id: number;
	instructions?: string | null;
	global: boolean;
}

export interface RefineReply {
	scheduled: boolean;
	reason?: string;
}

export interface SnapshotReply {
	ok: boolean;
	saved?: string[];
	skipped?: string[];
	error?: string;
}

export interface KernelOptions {
	pythonPath: string;
	kernelPath: string;
	depth: number;
	/** namespace pickle path; kernel restores it at boot and writes it on snapshot() */
	snapshotPath?: string;
	onRlmRequest: (req: RlmRequest, signal?: AbortSignal) => Promise<RlmReply>;
	onRefineRequest: (req: RefineRequest) => Promise<RefineReply>;
}

const BOOT_TIMEOUT_MS = 30_000;
const INTERRUPT_GRACE_MS = 4_000;
const PYTHON_LOG_CAP = 4_000;

export class KernelProcess {
	private opts: KernelOptions;
	private server: Server | null = null;
	private socket: Socket | null = null;
	private proc: ChildProcess | null = null;
	private socketPath = "";
	private buffer = "";
	private ready = false;
	private starting: Promise<void> | null = null;
	private queue: Promise<unknown> = Promise.resolve();
	private nextId = 1;
	private pending = new Map<number, { resolve: (v: unknown) => void; reject: (e: Error) => void }>();
	private readySettle: { resolve: () => void; reject: (e: Error) => void } | null = null;
	private pythonLog = "";
	private pythonVersion = "";
	private activeSignal: AbortSignal | undefined;
	private onStream: ((stream: "stdout" | "stderr", data: string) => void) | null = null;
	private generation = 0;
	private lastExecGeneration = -1;
	private execCount = 0;
	private rlmCount = 0;
	private restarts = 0;
	private restoredVars: string[] = [];

	constructor(opts: KernelOptions) {
		this.opts = opts;
	}

	get isReady(): boolean {
		return this.ready;
	}

	get stats() {
		return {
			ready: this.ready,
			python: this.pythonVersion,
			execs: this.execCount,
			rlmCalls: this.rlmCount,
			restarts: this.restarts,
			restored: this.restoredVars,
			pid: this.proc?.pid,
		};
	}

	async ensureStarted(): Promise<void> {
		if (this.ready && this.proc) return;
		if (this.starting) return this.starting;
		this.starting = this.boot();
		try {
			await this.starting;
		} finally {
			this.starting = null;
		}
	}

	private async boot(): Promise<void> {
		await this.shutdown();
		this.socketPath = join(
			tmpdir(),
			`pi-rlm-${process.pid}-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 6)}.sock`,
		);
		this.pythonLog = "";
		if (this.generation > 0) this.restarts++;
		this.generation++;

		await new Promise<void>((resolve, reject) => {
			const timeout = setTimeout(() => {
				this.settleBoot(
					new Error(
						`Timed out waiting for the pi-rlm Python kernel (${BOOT_TIMEOUT_MS / 1000}s).\n` +
							(this.pythonLog ? `Python output:\n${this.pythonLog}\n` : "") +
							`Check that "${this.opts.pythonPath}" is Python 3 (or set PI_RLM_PYTHON).`,
					),
				);
			}, BOOT_TIMEOUT_MS);

			this.readySettle = {
				resolve: () => {
					clearTimeout(timeout);
					resolve();
				},
				reject: (err: Error) => {
					clearTimeout(timeout);
					reject(err);
				},
			};

			this.server = createServer((sock) => {
				this.socket = sock;
				sock.setNoDelay(true);
				sock.on("data", (chunk) => this.onData(chunk));
				sock.on("close", () => this.onSocketClose(sock));
				sock.on("error", () => {});
			});
			this.server.on("error", (err) => this.settleBoot(err));
			this.server.listen(this.socketPath, () => {
				const env: NodeJS.ProcessEnv = {
					...process.env,
					PI_RLM_SOCKET: this.socketPath,
					PI_RLM_DEPTH: String(this.opts.depth),
				};
				if (this.opts.snapshotPath) env.PI_RLM_SNAPSHOT = this.opts.snapshotPath;
				const proc = spawn(this.opts.pythonPath, [this.opts.kernelPath], {
					env,
					stdio: ["ignore", "pipe", "pipe"],
				});
				this.proc = proc;
				proc.stdout?.on("data", (d) => this.logPython(d));
				proc.stderr?.on("data", (d) => this.logPython(d));
				proc.on("error", (err) =>
					this.settleBoot(
						new Error(
							`Failed to start "${this.opts.pythonPath}": ${err.message}\n` +
								"Install Python 3 or set PI_RLM_PYTHON to a Python 3 interpreter.",
						),
					),
				);
				proc.on("exit", () => {
					if (this.proc === proc) {
						this.proc = null;
						this.onKernelDeath();
					}
				});
			});
		});
	}

	private settleBoot(err?: Error) {
		const settle = this.readySettle;
		this.readySettle = null;
		if (!settle) return;
		if (err) settle.reject(err);
		else settle.resolve();
	}

	private logPython(chunk: Buffer | string) {
		this.pythonLog += chunk.toString();
		if (this.pythonLog.length > PYTHON_LOG_CAP) {
			this.pythonLog = this.pythonLog.slice(-PYTHON_LOG_CAP);
		}
	}

	private onData(chunk: Buffer | string) {
		this.buffer += chunk.toString("utf8");
		let idx: number;
		while ((idx = this.buffer.indexOf("\n")) >= 0) {
			const line = this.buffer.slice(0, idx);
			this.buffer = this.buffer.slice(idx + 1);
			if (line.trim()) this.dispatch(line);
		}
	}

	private dispatch(line: string) {
		let msg: Record<string, never>;
		try {
			msg = JSON.parse(line);
		} catch {
			return;
		}
		switch (msg.type as string) {
			case "ready": {
				this.ready = true;
				const m = msg as unknown as { python?: string; restored?: string[] };
				this.pythonVersion = String(m.python ?? "");
				this.restoredVars = Array.isArray(m.restored) ? m.restored : [];
				this.settleBoot();
				break;
			}
			case "result":
			case "set_done":
			case "snapshot_done":
				this.settleRequest((msg as unknown as { id: number }).id, msg);
				break;
			case "vars": {
				const m = msg as unknown as { id: number; vars?: VarInfo[] };
				this.settleRequest(m.id, m.vars ?? []);
				break;
			}
			case "stream": {
				const m = msg as unknown as { stream: "stdout" | "stderr"; data: string };
				this.onStream?.(m.stream, m.data);
				break;
			}
			case "rlm_request": {
				const m = msg as unknown as RlmRequest;
				this.rlmCount++;
				this.opts
					.onRlmRequest(
						{ id: m.id, prompt: String(m.prompt ?? ""), depth: m.depth ?? this.opts.depth + 1 },
						this.activeSignal,
					)
					.then((reply) => this.send({ type: "rlm_response", id: m.id, ...reply }))
					.catch((err: unknown) =>
						this.send({
							type: "rlm_response",
							id: m.id,
							ok: false,
							error: err instanceof Error ? err.message : String(err),
						}),
					);
				break;
			}
			case "refine_request": {
				const m = msg as unknown as RefineRequest;
				this.opts
					.onRefineRequest({
						id: m.id,
						instructions: m.instructions ?? undefined,
						global: Boolean(m.global),
					})
					.then((reply) => this.send({ type: "refine_response", id: m.id, ...reply }))
					.catch((err: unknown) =>
						this.send({
							type: "refine_response",
							id: m.id,
							scheduled: false,
							reason: err instanceof Error ? err.message : String(err),
						}),
					);
				break;
			}
		}
	}

	private settleRequest(id: number, value: unknown) {
		const p = this.pending.get(id);
		if (p) {
			this.pending.delete(id);
			p.resolve(value);
		}
	}

	private rejectAll(err: Error) {
		for (const [, p] of this.pending) p.reject(err);
		this.pending.clear();
		this.settleBoot(err);
	}

	private onSocketClose(sock: Socket) {
		if (this.socket === sock) {
			this.socket = null;
			this.ready = false;
			this.rejectAll(new Error("pi-rlm kernel closed the connection"));
		}
	}

	private onKernelDeath() {
		this.ready = false;
		this.rejectAll(
			new Error(
				`pi-rlm Python kernel exited unexpectedly.` +
					(this.pythonLog ? `\nPython output:\n${this.pythonLog}` : ""),
			),
		);
		this.cleanupSocketFile();
	}

	private send(msg: Record<string, unknown>) {
		if (this.socket && !this.socket.destroyed) {
			this.socket.write(`${JSON.stringify(msg)}\n`);
		}
	}

	/**
	 * Execute a cell. Executions are serialized; pass `signal` to interrupt
	 * (SIGINT first, SIGKILL after a grace period — which kills the kernel
	 * and its namespace).
	 */
	exec(code: string, options: ExecOptions = {}): Promise<ExecResult> {
		const run = this.queue.then(() => this.execInner(code, options));
		this.queue = run.catch(() => {});
		return run;
	}

	private async execInner(code: string, options: ExecOptions): Promise<ExecResult> {
		await this.ensureStarted();
		const restarted = this.lastExecGeneration !== -1 && this.lastExecGeneration !== this.generation;
		const restoredVars = this.generation !== this.lastExecGeneration ? this.restoredVars : [];

		const id = this.nextId++;
		this.onStream = options.onStream ?? null;
		this.activeSignal = options.signal;
		let cancelled = false;
		let killTimer: NodeJS.Timeout | undefined;
		const onAbort = () => {
			cancelled = true;
			this.proc?.kill("SIGINT");
			killTimer = setTimeout(() => {
				if (this.pending.has(id)) this.proc?.kill("SIGKILL");
			}, INTERRUPT_GRACE_MS);
		};
		if (options.signal) {
			if (options.signal.aborted) onAbort();
			else options.signal.addEventListener("abort", onAbort, { once: true });
		}

		try {
			const result = await new Promise<{
				ok: boolean;
				stdout?: string;
				stderr?: string;
				result?: string | null;
			}>((resolve, reject) => {
				this.pending.set(id, { resolve: resolve as (v: unknown) => void, reject });
				this.send({ type: "exec", id, code });
			});
			this.execCount++;
			this.lastExecGeneration = this.generation;
			return {
				ok: Boolean(result.ok),
				stdout: result.stdout ?? "",
				stderr: result.stderr ?? "",
				result: result.result ?? null,
				cancelled,
				restarted,
				restoredVars,
			};
		} finally {
			if (killTimer) clearTimeout(killTimer);
			options.signal?.removeEventListener("abort", onAbort);
			this.activeSignal = undefined;
			this.onStream = null;
		}
	}

	/** Best-effort variable injection. Ordered before any later exec on the socket. */
	async setVar(name: string, value: unknown): Promise<void> {
		if (!this.ready) return;
		const id = this.nextId++;
		await new Promise<void>((resolve, reject) => {
			this.pending.set(id, { resolve: resolve as (v: unknown) => void, reject });
			this.send({ type: "set", id, name, value });
		});
	}

	async listVars(): Promise<VarInfo[]> {
		if (!this.ready) return [];
		const id = this.nextId++;
		return new Promise<VarInfo[]>((resolve, reject) => {
			this.pending.set(id, { resolve: resolve as (v: unknown) => void, reject });
			this.send({ type: "vars_request", id });
		});
	}

	/** Persist picklable namespace values to the snapshot file. Best-effort. */
	async snapshot(): Promise<SnapshotReply> {
		if (!this.ready || !this.opts.snapshotPath) {
			return { ok: false, error: "kernel not running or no snapshot path" };
		}
		const id = this.nextId++;
		return new Promise<SnapshotReply>((resolve) => {
			const timer = setTimeout(() => {
				if (this.pending.delete(id)) resolve({ ok: false, error: "snapshot timed out" });
			}, 15_000);
			this.pending.set(id, {
				resolve: (v: unknown) => {
					clearTimeout(timer);
					resolve(v as SnapshotReply);
				},
				reject: () => {
					clearTimeout(timer);
					resolve({ ok: false, error: "kernel went away" });
				},
			});
			this.send({ type: "snapshot_request", id });
		});
	}

	async shutdown(): Promise<void> {
		const proc = this.proc;
		this.proc = null;
		this.ready = false;
		this.socket = null;
		this.readySettle = null;
		this.rejectAll(new Error("pi-rlm kernel shut down"));
		if (this.server) {
			try {
				this.server.close();
			} catch {
				/* ignore */
			}
			this.server = null;
		}
		if (proc) {
			try {
				proc.kill("SIGTERM");
			} catch {
				/* ignore */
			}
			const timer = setTimeout(() => {
				try {
					proc.kill("SIGKILL");
				} catch {
					/* ignore */
				}
			}, 2_000);
			proc.once("exit", () => clearTimeout(timer));
		}
		this.cleanupSocketFile();
	}

	private cleanupSocketFile() {
		if (this.socketPath) {
			try {
				rmSync(this.socketPath, { force: true });
			} catch {
				/* ignore */
			}
		}
	}
}