import type { ChildProcess, StdioOptions } from "node:child_process";
import { randomUUID } from "node:crypto";
import { chmodSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { AgentSession } from "../core/agent-session.js";
import type { AgentSessionRuntime } from "../core/agent-session-runtime.js";
import {
	clearOrphanProcessJournal,
	killOrphanProcess,
	ORPHAN_PROCESS_JOURNAL_ENV,
	readActiveOrphanProcesses,
	shouldReapOrphanProcess,
} from "../core/orphan-process-journal.js";
import { SESSION_LEASE_OWNER_ID_ENV, SESSION_LEASES_ENABLED_ENV } from "../core/session-lease.js";
import { attachJsonlLineReader, serializeJsonLine } from "../modes/rpc/jsonl.js";
import { spawnHidden } from "../utils/child-process.js";
import { isHelpCommandRequest, PUBLIC_COMMAND_NAMES, REMOVED_COMMAND_NAMES } from "./command-registry.js";
import { type CliSubprocessLaunchSpec, createCliSubprocessLaunchSpec } from "./subprocess-launch.js";

const OWNED_WORKER_ENV = "PRIME_AGENT_INTERNAL_OWNED_WORKER";
const OWNED_RECOVERY_DESCRIPTOR_ENV = "PRIME_AGENT_INTERNAL_OWNED_RECOVERY_DESCRIPTOR";
const OWNED_PROFILE_ENV = "PRIME_AGENT_INTERNAL_OWNED_PROFILE";

let closeOwnerWatch: (() => void) | undefined;

export type OwnedSessionWorkerProfile = "print" | "json" | "rpc" | "interactive-ephemeral";

export function isOwnedSessionWorkerProcess(environment: NodeJS.ProcessEnv = process.env): boolean {
	return environment[OWNED_WORKER_ENV] === "1";
}

interface OwnedSessionRecoveryDescriptor {
	version: 1;
	profile: OwnedSessionWorkerProfile;
	sessionId: string;
	sessionFile?: string;
	cwd: string;
	updatedAt: string;
}

const NON_SESSION_FLAGS = new Set(["--help", "-h", "--version", "-v", "--list-models", "--export"]);

const NON_SESSION_COMMANDS = new Set([...PUBLIC_COMMAND_NAMES, ...REMOVED_COMMAND_NAMES]);

function valueAfter(args: readonly string[], flag: string): string | undefined {
	const index = args.indexOf(flag);
	return index === -1 ? undefined : args[index + 1];
}

function hasNonSessionOperation(args: readonly string[]): boolean {
	if (args.some((arg) => NON_SESSION_FLAGS.has(arg) || arg.startsWith("--export="))) {
		return true;
	}
	const first = args[0];
	if (first === "help") {
		return isHelpCommandRequest(args.slice(1));
	}
	return first !== undefined && NON_SESSION_COMMANDS.has(first);
}

function isStartupBenchmark(environment: NodeJS.ProcessEnv): boolean {
	const value = environment.PI_STARTUP_BENCHMARK?.toLowerCase();
	return value === "1" || value === "true" || value === "yes";
}

export function classifyOwnedSessionWorkerInvocation(
	args: readonly string[],
	stdinIsTTY: boolean | undefined,
	environment: NodeJS.ProcessEnv = process.env,
): OwnedSessionWorkerProfile | undefined {
	if (isOwnedSessionWorkerProcess(environment) || hasNonSessionOperation(args)) {
		return undefined;
	}

	const mode = valueAfter(args, "--mode");
	if (mode === "daemon") {
		return undefined;
	}
	if (mode === "rpc") {
		return "rpc";
	}
	if (mode === "json") {
		return "json";
	}
	if (args.includes("--print") || args.includes("-p") || stdinIsTTY === false) {
		return "print";
	}
	if (args.includes("--no-session") || isStartupBenchmark(environment)) {
		return "interactive-ephemeral";
	}
	return undefined;
}

export type OwnedWorkerLaunchSpec = CliSubprocessLaunchSpec;

export function createOwnedWorkerLaunchSpec(
	args: readonly string[],
	executable = process.execPath,
	execArgs: readonly string[] = process.execArgv,
	entrypoint = process.argv[1],
): OwnedWorkerLaunchSpec {
	return createCliSubprocessLaunchSpec(args, executable, execArgs, entrypoint);
}

function readOwnedRecoveryDescriptor(path: string): OwnedSessionRecoveryDescriptor | undefined {
	try {
		const value = JSON.parse(readFileSync(path, "utf8")) as Partial<OwnedSessionRecoveryDescriptor>;
		if (
			value.version === 1 &&
			typeof value.sessionId === "string" &&
			typeof value.cwd === "string" &&
			(value.sessionFile === undefined || typeof value.sessionFile === "string")
		) {
			return value as OwnedSessionRecoveryDescriptor;
		}
	} catch {
		// The worker may have stopped before creating a recoverable session.
	}
	return undefined;
}

function writeOwnedRecoveryDescriptor(path: string, profile: OwnedSessionWorkerProfile, session: AgentSession): void {
	const descriptor: OwnedSessionRecoveryDescriptor = {
		version: 1,
		profile,
		sessionId: session.sessionId,
		...(session.sessionFile ? { sessionFile: session.sessionFile } : {}),
		cwd: session.sessionManager.getCwd(),
		updatedAt: new Date().toISOString(),
	};
	const tempPath = `${path}.${process.pid}.tmp`;
	writeFileSync(tempPath, `${JSON.stringify(descriptor)}\n`, { mode: 0o600 });
	chmodSync(tempPath, 0o600);
	renameSync(tempPath, path);
}

export function installOwnedSessionRecoveryTracking(runtime: AgentSessionRuntime): void {
	const path = process.env[OWNED_RECOVERY_DESCRIPTOR_ENV];
	const profile = process.env[OWNED_PROFILE_ENV] as OwnedSessionWorkerProfile | undefined;
	if (!path || !profile) {
		return;
	}
	let unsubscribeSession: (() => void) | undefined;
	const bind = (session: AgentSession) => {
		unsubscribeSession?.();
		writeOwnedRecoveryDescriptor(path, profile, session);
		let lastSessionFile = session.sessionFile;
		unsubscribeSession = session.subscribe((event) => {
			if (event.type !== "message_start" && event.type !== "session_info_changed") {
				return;
			}
			if (session.sessionFile !== lastSessionFile) {
				lastSessionFile = session.sessionFile;
				writeOwnedRecoveryDescriptor(path, profile, session);
			}
		});
	};
	bind(runtime.session);
	runtime.onSessionReplaced(bind);
}

export function createRpcRecoveryArgs(args: readonly string[], sessionPath: string): string[] {
	const recovered: string[] = [];
	for (let index = 0; index < args.length; index++) {
		const arg = args[index]!;
		if (arg === "--resume" || arg === "-r" || arg === "--fork") {
			index++;
			continue;
		}
		if (arg.startsWith("--resume=")) {
			continue;
		}
		if (arg === "--continue" || arg === "-c") {
			continue;
		}
		recovered.push(arg);
	}
	return [...recovered, "--resume", sessionPath];
}

function exitCodeForSignal(signal: NodeJS.Signals | null): number {
	if (signal === "SIGHUP") {
		return 129;
	}
	if (signal === "SIGINT") {
		return 130;
	}
	if (signal === "SIGTERM") {
		return 143;
	}
	return 1;
}

function forwardSignal(child: ChildProcess, signal: NodeJS.Signals): void {
	if (child.exitCode === null && child.signalCode === null) {
		child.kill(signal);
	}
}

export async function runOwnedSessionWorkerFrontend(
	args: readonly string[],
	profile: OwnedSessionWorkerProfile,
): Promise<number> {
	const interactive = profile === "interactive-ephemeral";
	const recoveryDescriptorPath = join(tmpdir(), `prime-agent-owned-${process.pid}-${randomUUID().slice(0, 12)}.json`);
	const orphanProcessJournalPath = `${recoveryDescriptorPath}.orphans.jsonl`;
	let currentChild: ChildProcess | undefined;
	let terminating = false;
	let terminationSignal: NodeJS.Signals | undefined;
	let stdinEnded = false;
	let currentRpcInput: NodeJS.WritableStream | undefined;
	let currentRpcOutput: NodeJS.ReadableStream | undefined;
	let rpcStdoutPaused = false;
	let detachRpcInput: (() => void) | undefined;
	let detachRpcOutput: (() => void) | undefined;
	const bufferedRpcInput: string[] = [];
	const pendingRpcCommands = new Map<string, { publicId?: string; command: string }>();
	const anonymousRpcIdPrefix = `prime-agent-owned-${randomUUID()}`;
	let anonymousRpcCommandId = 0;

	const prepareRpcInput = (line: string): string => {
		try {
			const command = JSON.parse(line) as { id?: unknown; type?: unknown } | null;
			if (
				!command ||
				Array.isArray(command) ||
				typeof command.type !== "string" ||
				command.type === "extension_ui_response" ||
				command.type === "ack_result"
			) {
				return `${line}\n`;
			}
			const publicId = typeof command.id === "string" ? command.id : undefined;
			const internalId = publicId ?? `${anonymousRpcIdPrefix}-${++anonymousRpcCommandId}`;
			pendingRpcCommands.set(internalId, { publicId, command: command.type });
			return publicId !== undefined ? `${line}\n` : serializeJsonLine({ ...command, id: internalId });
		} catch {
			// The worker preserves the existing parse-error response contract.
			return `${line}\n`;
		}
	};
	const observeRpcOutput = (line: string) => {
		let parsed: unknown;
		try {
			parsed = JSON.parse(line) as unknown;
		} catch {
			return;
		}
		if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
			return;
		}
		const response = parsed as { id?: unknown; type?: unknown; command?: unknown };
		const id = typeof response.id === "string" ? response.id : undefined;
		const pending = id !== undefined ? pendingRpcCommands.get(id) : undefined;
		let framed = `${line}\n`;
		if (response.type === "response" && pending?.publicId === undefined) {
			const { id: _internalId, ...publicResponse } = response;
			framed = serializeJsonLine(publicResponse);
		}
		if (!process.stdout.write(framed) && !rpcStdoutPaused) {
			rpcStdoutPaused = true;
			currentRpcOutput?.pause();
			process.stdout.once("drain", () => {
				rpcStdoutPaused = false;
				currentRpcOutput?.resume();
			});
		}
		if (response.type !== "response" || typeof response.command !== "string") {
			return;
		}
		if (id !== undefined && pending?.command === response.command) {
			pendingRpcCommands.delete(id);
		}
	};
	const failPendingRpcCommands = () => {
		for (const pending of pendingRpcCommands.values()) {
			process.stdout.write(
				serializeJsonLine({
					...(pending.publicId !== undefined ? { id: pending.publicId } : {}),
					type: "response",
					command: pending.command,
					success: false,
					error: "The isolated session worker stopped during this command; its result is uncertain and was not replayed",
				}),
			);
		}
		pendingRpcCommands.clear();
	};
	const reapWorkerResources = (workerPid: number | undefined) => {
		if (!workerPid) {
			return;
		}
		if (process.platform !== "win32") {
			try {
				process.kill(-workerPid, "SIGKILL");
			} catch {
				// The worker process group may already be fully reaped.
			}
		}
		for (const orphan of readActiveOrphanProcesses(orphanProcessJournalPath, workerPid)) {
			if (!shouldReapOrphanProcess(orphan)) {
				continue;
			}
			killOrphanProcess(orphan.pid);
		}
		clearOrphanProcessJournal(orphanProcessJournalPath);
	};

	if (profile === "rpc") {
		detachRpcInput = attachJsonlLineReader(process.stdin, (line) => {
			const framed = prepareRpcInput(line);
			const input = currentRpcInput;
			if (input?.writable) {
				if (!input.write(framed)) {
					process.stdin.pause();
					input.once("drain", () => {
						if (currentRpcInput === input && !stdinEnded) {
							process.stdin.resume();
						}
					});
				}
			} else {
				bufferedRpcInput.push(framed);
			}
		});
		process.stdin.once("end", () => {
			stdinEnded = true;
			currentRpcInput?.end();
		});
	}

	const spawnWorker = (workerArgs: readonly string[]): ChildProcess => {
		const launch = createOwnedWorkerLaunchSpec(workerArgs);
		const bridgeStdin = profile === "rpc" || process.stdin.isTTY !== true;
		const stdio: StdioOptions = interactive
			? ["inherit", "inherit", "inherit", "ipc"]
			: [bridgeStdin ? "pipe" : "inherit", "pipe", "pipe", "ipc"];
		const child = spawnHidden(launch.command, launch.args, {
			cwd: process.cwd(),
			detached: process.platform !== "win32",
			env: {
				...process.env,
				[OWNED_WORKER_ENV]: "1",
				[OWNED_RECOVERY_DESCRIPTOR_ENV]: recoveryDescriptorPath,
				[OWNED_PROFILE_ENV]: profile,
				[ORPHAN_PROCESS_JOURNAL_ENV]: orphanProcessJournalPath,
				[SESSION_LEASES_ENABLED_ENV]: "1",
				[SESSION_LEASE_OWNER_ID_ENV]: `owned-${randomUUID()}`,
			},
			stdio,
		});
		currentChild = child;
		if (!interactive) {
			const childInput = child.stdin ?? undefined;
			const childOutput = child.stdout ?? undefined;
			const childError = child.stderr ?? undefined;
			if ((bridgeStdin && !childInput) || !childOutput || !childError) {
				child.kill("SIGTERM");
				throw new Error("Owned session worker did not expose bridged stdio");
			}
			childError.pipe(process.stderr, { end: false });
			if (profile === "rpc") {
				if (!childInput) {
					throw new Error("Owned RPC worker did not expose stdin");
				}
				currentRpcInput = childInput;
				currentRpcOutput = childOutput;
				if (rpcStdoutPaused) {
					childOutput.pause();
				}
				detachRpcOutput = attachJsonlLineReader(childOutput, observeRpcOutput);
				for (const buffered of bufferedRpcInput.splice(0)) {
					childInput.write(buffered);
				}
				if (stdinEnded) {
					childInput.end();
				}
			} else {
				if (childInput) {
					process.stdin.pipe(childInput);
				}
				childOutput.pipe(process.stdout, { end: false });
			}
		}
		return child;
	};

	const signals: NodeJS.Signals[] = ["SIGINT", "SIGTERM"];
	if (process.platform !== "win32") {
		signals.push("SIGHUP");
	}
	const signalHandlers = signals.map((signal) => {
		const handler = () => {
			terminating = true;
			terminationSignal ??= signal;
			if (currentChild) {
				forwardSignal(currentChild, signal);
			}
		};
		process.on(signal, handler);
		return { signal, handler };
	});

	try {
		let workerArgs = [...args];
		let recoveryAttempt = 0;
		while (true) {
			if (terminating) {
				return exitCodeForSignal(terminationSignal ?? null);
			}
			const workerStartedAt = Date.now();
			const child = spawnWorker(workerArgs);
			const workerPid = child.pid;
			const exit = await new Promise<{ code: number; signal: NodeJS.Signals | null }>((resolveExit, reject) => {
				child.once("error", reject);
				child.once("close", (code, signal) => resolveExit({ code: code ?? exitCodeForSignal(signal), signal }));
			});
			currentChild = undefined;
			currentRpcInput = undefined;
			currentRpcOutput = undefined;
			if (profile === "rpc" && !stdinEnded) {
				process.stdin.resume();
			}
			detachRpcOutput?.();
			detachRpcOutput = undefined;
			if (!interactive && profile !== "rpc" && child.stdin) {
				process.stdin.unpipe(child.stdin);
			}
			if (child.connected) {
				child.disconnect();
			}
			reapWorkerResources(workerPid);
			const rpcCrashed =
				profile === "rpc" &&
				!terminating &&
				(exit.code !== 0 || exit.signal !== null || pendingRpcCommands.size > 0);
			const workerExitCode = rpcCrashed && exit.code === 0 ? 1 : exit.code;
			if (Date.now() - workerStartedAt >= 60_000) {
				recoveryAttempt = 0;
			}
			if (rpcCrashed) {
				failPendingRpcCommands();
			}
			const shouldRecover = rpcCrashed && !stdinEnded && recoveryAttempt < 3;
			if (!shouldRecover) {
				return terminationSignal ? exitCodeForSignal(terminationSignal) : workerExitCode;
			}
			const descriptor = readOwnedRecoveryDescriptor(recoveryDescriptorPath);
			if (!descriptor?.sessionFile) {
				return terminationSignal ? exitCodeForSignal(terminationSignal) : workerExitCode;
			}
			workerArgs = createRpcRecoveryArgs(args, descriptor.sessionFile);
			const retryDelay = [250, 1000, 5000][recoveryAttempt] ?? 5000;
			recoveryAttempt++;
			await new Promise((resolveDelay) => setTimeout(resolveDelay, retryDelay));
			if (terminating) {
				return exitCodeForSignal(terminationSignal ?? null);
			}
		}
	} finally {
		for (const { signal, handler } of signalHandlers) {
			process.off(signal, handler);
		}
		detachRpcInput?.();
		detachRpcOutput?.();
		rmSync(recoveryDescriptorPath, { force: true });
		clearOrphanProcessJournal(orphanProcessJournalPath);
	}
}

export async function maybeRunOwnedSessionWorkerFrontend(
	args: readonly string[],
	forceLegacyFrontend = false,
): Promise<boolean> {
	if (!forceLegacyFrontend && process.env.PRIME_AGENT_INTERNAL_LEGACY_OWNED_WORKER_FRONTEND !== "1") {
		return false;
	}
	const profile = classifyOwnedSessionWorkerInvocation(args, process.stdin.isTTY);
	if (!profile) {
		return false;
	}
	process.exitCode = await runOwnedSessionWorkerFrontend(args, profile);
	return true;
}

export function installOwnedSessionWorkerOwnerWatch(): void {
	if (!isOwnedSessionWorkerProcess()) {
		return;
	}
	if (!process.channel) {
		throw new Error("Owned session worker is missing its owner channel");
	}

	let ownerGone = false;
	const terminate = () => {
		if (ownerGone) {
			return;
		}
		ownerGone = true;
		closeOwnerWatch = undefined;
		const forceTimer = setTimeout(() => {
			if (process.platform !== "win32") {
				try {
					process.kill(-process.pid, "SIGKILL");
					return;
				} catch {
					// Fall through to terminating only this process.
				}
			}
			process.exit(143);
		}, 5000);
		forceTimer.unref();
		process.kill(process.pid, "SIGTERM");
	};
	process.once("disconnect", terminate);
	process.channel.unref();
	closeOwnerWatch = () => {
		if (ownerGone) {
			return;
		}
		ownerGone = true;
		process.off("disconnect", terminate);
		if (process.connected) {
			process.disconnect();
		}
		closeOwnerWatch = undefined;
	};
}

export function closeOwnedSessionWorkerOwnerWatch(): void {
	closeOwnerWatch?.();
}