import { spawn } from "node:child_process";
import { existsSync, promises as fs } from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { type ProcessRunner, runChecked } from "./process.ts";
import type { AgentOutcome, ModelSelection, ResolvedModel } from "./types.ts";

export interface PiAgentRequest {
	prompt: string;
	task: string;
	cwd: string;
	selection: ModelSelection;
	signal?: AbortSignal;
	runner?: ProcessRunner;
	executable?: string;
	tools?: readonly string[];
}

function invocation(
	override: string | undefined,
	args: string[],
): { command: string; args: string[] } {
	if (override) return { command: override, args };
	const script = process.argv[1];
	if (script && !script.startsWith("/$bunfs/root/") && existsSync(script)) {
		return { command: process.execPath, args: [script, ...args] };
	}
	const runtime = path.basename(process.execPath).toLowerCase();
	return /^(node|bun)(\.exe)?$/.test(runtime)
		? { command: "pi", args }
		: { command: process.execPath, args };
}

const TERMINAL_STOP_REASONS = new Set(["error", "aborted"]);

interface StreamingState {
	finalText: string;
	partialText: string;
	stopReason: string;
	errorMessage: string | undefined;
	resolvedModel: ResolvedModel | undefined;
	piRetryAttempts: number;
}

function newStreamingState(): StreamingState {
	return {
		finalText: "",
		partialText: "",
		stopReason: "end",
		errorMessage: undefined,
		resolvedModel: undefined,
		piRetryAttempts: 0,
	};
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function captureResolvedModel(event: unknown, state: StreamingState): void {
	if (
		!isRecord(event) ||
		event.type !== "model_select" ||
		!isRecord(event.model)
	)
		return;
	const model = event.model;
	const provider = typeof model.provider === "string" ? model.provider : "";
	const id = typeof model.id === "string" ? model.id : "";
	const contextWindow =
		typeof model.contextWindow === "number" ? model.contextWindow : 0;
	const maxTokens = typeof model.maxTokens === "number" ? model.maxTokens : 0;
	if (!provider || !id || contextWindow <= 0 || maxTokens <= 0) return;
	state.resolvedModel = { provider, id, contextWindow, maxTokens };
}

function updateStreamingState(line: string, state: StreamingState): void {
	if (!line.trim()) return;
	let event: unknown;
	try {
		event = JSON.parse(line);
	} catch {
		throw new Error("Pi subprocess returned malformed JSONL");
	}
	captureResolvedModel(event, state);
	if (
		!isRecord(event) ||
		event.type !== "message_end" ||
		!isRecord(event.message)
	) {
		return;
	}
	const message = event.message;
	if (typeof message.stopReason === "string")
		state.stopReason = message.stopReason;
	if (typeof message.errorMessage === "string")
		state.errorMessage = message.errorMessage;
	if (
		typeof message.retryAttempts === "number" &&
		message.retryAttempts > state.piRetryAttempts
	) {
		state.piRetryAttempts = message.retryAttempts;
	}
	if (message.role !== "assistant" || !Array.isArray(message.content)) return;
	for (const part of message.content) {
		if (
			isRecord(part) &&
			part.type === "text" &&
			typeof part.text === "string"
		) {
			state.finalText = part.text;
			if (!state.partialText) state.partialText = part.text;
		}
	}
}

function buildOutcome(state: StreamingState): AgentOutcome {
	const common = {
		partialText: state.partialText,
		resolvedModel: state.resolvedModel,
		piRetryAttempts: state.piRetryAttempts,
	};
	if (TERMINAL_STOP_REASONS.has(state.stopReason) || state.errorMessage) {
		return {
			ok: false,
			stopReason: state.stopReason,
			errorMessage: state.errorMessage ?? `Pi stopped: ${state.stopReason}`,
			...common,
		};
	}
	if (!state.finalText) {
		return {
			ok: false,
			stopReason: state.stopReason,
			errorMessage: "Pi subprocess returned no final assistant output",
			...common,
		};
	}
	return {
		ok: true,
		text: state.finalText,
		stopReason: state.stopReason,
		...common,
	};
}

function runStreamingPi(
	command: string,
	args: string[],
	cwd: string,
	signal?: AbortSignal,
): Promise<AgentOutcome> {
	return new Promise((resolve, reject) => {
		const child = spawn(command, args, {
			cwd,
			shell: false,
			stdio: ["ignore", "pipe", "pipe"],
		});
		let buffer = "";
		let stderr = "";
		const state = newStreamingState();
		let failure: Error | undefined;
		let aborted = false;
		let closed = false;
		let killTimer: NodeJS.Timeout | undefined;

		const fail = (error: Error) => {
			if (failure) return;
			failure = error;
			child.kill("SIGTERM");
		};
		const processLine = (line: string) => {
			try {
				updateStreamingState(line, state);
			} catch (error) {
				fail(
					error instanceof Error ? error : new Error("Invalid Pi JSONL event"),
				);
			}
		};
		child.stdout.on("data", (chunk: Buffer) => {
			buffer += chunk.toString("utf8");
			const lines = buffer.split("\n");
			buffer = lines.pop() ?? "";
			for (const line of lines) processLine(line);
			if (buffer.length > 16 * 1024 * 1024)
				fail(new Error("Pi emitted an oversized JSONL record"));
		});
		child.stderr.on("data", (chunk: Buffer) => {
			if (stderr.length < 1024 * 1024) stderr += chunk.toString("utf8");
		});
		const abort = () => {
			aborted = true;
			child.kill("SIGTERM");
			killTimer = setTimeout(() => {
				if (!closed) child.kill("SIGKILL");
			}, 5000);
		};
		if (signal?.aborted) abort();
		else signal?.addEventListener("abort", abort, { once: true });
		child.on("error", (error) =>
			fail(new Error(`Failed to start ${command}: ${error.message}`)),
		);
		child.on("close", (exitCode) => {
			closed = true;
			if (killTimer) clearTimeout(killTimer);
			signal?.removeEventListener("abort", abort);
			if (buffer.trim() && !failure) processLine(buffer);
			if (failure) reject(failure);
			else if (aborted)
				resolve({
					ok: false,
					stopReason: "aborted",
					errorMessage: `${command} was aborted`,
					partialText: state.partialText,
					resolvedModel: state.resolvedModel,
					piRetryAttempts: state.piRetryAttempts,
				});
			else if (exitCode !== 0)
				resolve({
					ok: false,
					stopReason: "error",
					errorMessage: `${command} exited with ${exitCode}: ${stderr.trim()}`,
					partialText: state.partialText,
					resolvedModel: state.resolvedModel,
					piRetryAttempts: state.piRetryAttempts,
				});
			else resolve(buildOutcome(state));
		});
	});
}

export async function runPiAgent(
	request: PiAgentRequest,
): Promise<AgentOutcome> {
	const directory = await fs.mkdtemp(
		path.join(os.tmpdir(), "pi-code-review-prompt-"),
	);
	const promptPath = path.join(directory, "system.md");
	try {
		await fs.writeFile(promptPath, request.prompt, {
			encoding: "utf8",
			mode: 0o600,
		});
		const toolArgs =
			request.tools?.length === 0
				? ["--no-tools"]
				: [
						"--tools",
						(request.tools ?? ["read", "grep", "find", "ls"]).join(","),
					];
		const args = [
			"--mode",
			"json",
			"--print",
			"--no-session",
			"--no-extensions",
			"--no-skills",
			"--no-prompt-templates",
			"--no-themes",
			"--no-context-files",
			"--model",
			request.selection.model,
			"--thinking",
			request.selection.thinking,
			...toolArgs,
			"--append-system-prompt",
			promptPath,
			request.task,
		];
		const resolved = invocation(request.executable, args);
		if (!request.runner) {
			return await runStreamingPi(
				resolved.command,
				resolved.args,
				request.cwd,
				request.signal,
			);
		}
		const result = await runChecked(request.runner, {
			command: resolved.command,
			args: resolved.args,
			cwd: request.cwd,
			signal: request.signal,
		});
		const state = newStreamingState();
		for (const line of result.stdout.split("\n")) {
			updateStreamingState(line, state);
		}
		return buildOutcome(state);
	} finally {
		await fs.rm(directory, { recursive: true, force: true });
	}
}