/**
 * Child process plumbing: spawn `pi --mode json -p` and parse its JSONL
 * event stream. Used by runChildRlm (recursive sub-agents, which reload
 * this extension) and by the refinement harness (a plain tool-less child).
 *
 * Token usage is accumulated from assistant message_end events so callers
 * can report nested cost.
 */

import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as path from "node:path";
import type { Usage } from "@earendil-works/pi-ai";

const KILL_GRACE_MS = 5_000;
const STDERR_CAP = 8_192;
const MAX_PROMPT_BYTES = 1024 * 1024;
const MAX_OUTPUT_CHARS = 64_000;

export interface PiJsonRunOptions {
	cwd: string;
	env?: NodeJS.ProcessEnv;
	signal?: AbortSignal;
}

export interface PiJsonRunResult {
	ok: boolean;
	output: string;
	error?: string;
	usage: Usage;
	turns: number;
	model?: string;
	stopReason?: string;
}

export interface ChildRunResult {
	ok: boolean;
	output: string;
	error?: string;
	usage: Usage;
	turns: number;
	model?: string;
	stopReason?: string;
}

export interface ChildRunOptions {
	prompt: string;
	/** depth of the child agent (its kernel's PI_RLM_DEPTH) */
	depth: number;
	maxDepth: number;
	/** absolute path to this extension's index.ts (passed as -e) */
	extensionPath: string;
	/** "provider/model" — defaults to the user's configured model */
	model?: string;
	cwd: string;
	signal?: AbortSignal;
}

export function emptyUsage(): Usage {
	return {
		input: 0,
		output: 0,
		cacheRead: 0,
		cacheWrite: 0,
		totalTokens: 0,
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
	};
}

export function addUsage(total: Usage, u: Partial<Usage> | undefined): void {
	if (!u) return;
	total.input += u.input ?? 0;
	total.output += u.output ?? 0;
	total.cacheRead += u.cacheRead ?? 0;
	total.cacheWrite += u.cacheWrite ?? 0;
	total.totalTokens += u.totalTokens ?? 0;
	if (u.cost) {
		total.cost.input += u.cost.input ?? 0;
		total.cost.output += u.cost.output ?? 0;
		total.cost.cacheRead += u.cost.cacheRead ?? 0;
		total.cost.cacheWrite += u.cost.cacheWrite ?? 0;
		total.cost.total += u.cost.total ?? 0;
	}
}

export function getPiInvocation(args: string[]): { command: string; args: string[] } {
	const currentScript = process.argv[1];
	const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
	if (currentScript && !isBunVirtualScript && fs.existsSync(currentScript)) {
		return { command: process.execPath, args: [currentScript, ...args] };
	}
	const execName = path.basename(process.execPath).toLowerCase();
	if (!/^(node|bun)(\.exe)?$/.test(execName)) {
		return { command: process.execPath, args };
	}
	return { command: "pi", args };
}

/** Spawn `pi --mode json -p` with pre-built args and collect the final assistant text. */
export async function runPiJson(args: string[], opts: PiJsonRunOptions): Promise<PiJsonRunResult> {
	const usage = emptyUsage();

	return new Promise<PiJsonRunResult>((resolve) => {
		let settled = false;
		let buffer = "";
		let stderr = "";
		let output = "";
		let turns = 0;
		let model: string | undefined;
		let stopReason: string | undefined;
		let errorMessage: string | undefined;

		const finish = (result: Omit<PiJsonRunResult, "usage" | "turns" | "model" | "stopReason">) => {
			if (settled) return;
			settled = true;
			resolve({ ...result, usage, turns, model, stopReason });
		};

		const processLine = (line: string) => {
			if (!line.trim()) return;
			let event: { type?: string; message?: Record<string, never> };
			try {
				event = JSON.parse(line);
			} catch {
				return;
			}
			if (event.type !== "message_end" || !event.message) return;
			const msg = event.message as {
				role?: string;
				usage?: Partial<Usage>;
				model?: string;
				stopReason?: string;
				errorMessage?: string;
				content?: Array<{ type?: string; text?: string }>;
			};
			if (msg.role !== "assistant") return;
			turns++;
			addUsage(usage, msg.usage);
			if (!model && msg.model) model = msg.model;
			if (msg.stopReason) stopReason = msg.stopReason;
			if (msg.errorMessage) errorMessage = String(msg.errorMessage);
			for (const part of msg.content ?? []) {
				if (part?.type === "text" && typeof part.text === "string") {
					output = part.text; // last assistant text wins
				}
			}
		};

		const invocation = getPiInvocation(args);
		const proc = spawn(invocation.command, invocation.args, {
			cwd: opts.cwd,
			env: opts.env ?? process.env,
			stdio: ["ignore", "pipe", "pipe"],
		});

		proc.stdout.on("data", (d) => {
			buffer += d.toString();
			const lines = buffer.split("\n");
			buffer = lines.pop() ?? "";
			for (const line of lines) processLine(line);
		});
		proc.stderr.on("data", (d) => {
			stderr += d.toString();
			if (stderr.length > STDERR_CAP) stderr = stderr.slice(-STDERR_CAP);
		});

		proc.on("error", (err) => {
			finish({ ok: false, output: "", error: `failed to spawn child pi: ${err.message}` });
		});

		proc.on("close", (code) => {
			if (buffer.trim()) processLine(buffer);
			if (errorMessage) {
				finish({ ok: false, output, error: errorMessage });
			} else if (code !== 0) {
				const tail = stderr.trim().slice(-500);
				finish({
					ok: false,
					output,
					error: tail ? `child pi exited ${code}: ${tail}` : `child pi exited with code ${code}`,
				});
			} else {
				finish({ ok: true, output });
			}
		});

		if (opts.signal) {
			const kill = () => {
				proc.kill("SIGTERM");
				setTimeout(() => {
					if (!proc.killed) proc.kill("SIGKILL");
				}, KILL_GRACE_MS);
			};
			if (opts.signal.aborted) kill();
			else opts.signal.addEventListener("abort", kill, { once: true });
		}
	});
}

export async function runChildRlm(opts: ChildRunOptions): Promise<ChildRunResult> {
	if (Buffer.byteLength(opts.prompt, "utf8") > MAX_PROMPT_BYTES) {
		return {
			ok: false,
			output: "",
			error: `rlm() prompt is ${Buffer.byteLength(opts.prompt, "utf8")} bytes; the limit is ${MAX_PROMPT_BYTES}. Pass smaller chunks.`,
			usage: emptyUsage(),
			turns: 0,
		};
	}

	const args = [
		"--mode",
		"json",
		"-p",
		"--no-session",
		// Don't inherit user/global extensions — a clean RLM child…
		"--no-extensions",
		// …except this one, which re-enables rlm() for deeper recursion.
		"-e",
		opts.extensionPath,
		"--rlm",
		"--rlm-max-depth",
		String(opts.maxDepth),
	];
	if (opts.model) args.push("--model", opts.model);
	args.push(opts.prompt);

	const env = {
		...process.env,
		PI_RLM_DEPTH: String(opts.depth),
		PI_RLM_MAX_DEPTH: String(opts.maxDepth),
	};

	const result = await runPiJson(args, { cwd: opts.cwd, env, signal: opts.signal });
	if (result.output.length > MAX_OUTPUT_CHARS) {
		result.output = `${result.output.slice(0, MAX_OUTPUT_CHARS)}\n... [rlm output truncated at ${MAX_OUTPUT_CHARS} chars]`;
	}
	return result;
}