import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { ThinkingLevel } from "@earendil-works/pi-agent-core";
import type { ApiRetryRecord, UsageStats, WorkerHistoryMessage, WorkerResult, WorkerToolCall } from "./types.ts";
import { usageFromHistory } from "./usage.ts";

export const API_RETRY_DELAYS_MS = [10_000, 30_000, 60_000] as const;

export interface ProcessPauseControl {
	isPaused(): boolean;
	onChange(listener: (paused: boolean) => void): () => void;
}

export interface PiRunEvent {
	type: string;
	output: string;
	stderr: string;
	tool?: WorkerToolCall;
	toolCalls: WorkerToolCall[];
	history: WorkerHistoryMessage[];
	usage: UsageStats;
	model?: string;
	stopReason?: string;
	errorMessage?: string;
	apiRetries: ApiRetryRecord[];
	apiRetry?: ApiRetryRecord;
}

export interface PiRunOptions {
	cwd: string;
	model: string;
	thinkingLevel?: ThinkingLevel;
	tools?: string[];
	systemPrompt: string;
	prompt: string;
	timeoutMs: number;
	signal?: AbortSignal;
	pauseControl?: ProcessPauseControl;
	onEvent?: (event: PiRunEvent) => void;
	apiRetryDelaysMs?: readonly number[];
	classifyRetryableError?: (message: WorkerHistoryMessage) => boolean | Promise<boolean>;
}

export interface PiRunResult extends WorkerResult {}

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

function extractText(message: WorkerHistoryMessage): string {
	if (typeof message.content === "string") return message.content;
	if (!Array.isArray(message.content)) return "";
	return message.content
		.filter((part): part is { type: "text"; text: string } => isRecord(part) && part.type === "text" && typeof part.text === "string")
		.map((part) => part.text)
		.join("\n");
}

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

function byteLimited(text: string, maxBytes: number): { text: string; truncated: boolean } {
	if (Buffer.byteLength(text, "utf8") <= maxBytes) return { text, truncated: false };
	let result = text.slice(0, maxBytes);
	while (Buffer.byteLength(result, "utf8") > maxBytes) result = result.slice(0, -1);
	return { text: `${result}\n\n[输出已截断]`, truncated: true };
}

async function writeSystemPrompt(prompt: string, name: string): Promise<{ dir: string; file: string }> {
	const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-cluster-"));
	const safeName = name.replace(/[^\w.-]+/g, "_");
	const file = path.join(dir, `${safeName}.md`);
	await fs.promises.writeFile(file, prompt, { encoding: "utf8", mode: 0o600 });
	return { dir, file };
}

async function findSessionFile(directory: string): Promise<string | undefined> {
	let entries: fs.Dirent[];
	try {
		entries = await fs.promises.readdir(directory, { withFileTypes: true });
	} catch {
		return undefined;
	}
	for (const entry of entries) {
		const entryPath = path.join(directory, entry.name);
		if (entry.isDirectory()) {
			const nested = await findSessionFile(entryPath);
			if (nested) return nested;
		} else if (entry.isFile() && entry.name.endsWith(".jsonl")) {
			return entryPath;
		}
	}
	return undefined;
}

async function waitUntilResumed(pauseControl?: ProcessPauseControl, signal?: AbortSignal): Promise<boolean> {
	if (signal?.aborted) return false;
	if (!pauseControl?.isPaused()) return true;
	return new Promise<boolean>((resolve) => {
		let settled = false;
		let removePauseListener = () => {};
		const onAbort = () => finish(false);
		const finish = (resumed: boolean) => {
			if (settled) return;
			settled = true;
			removePauseListener();
			signal?.removeEventListener("abort", onAbort);
			resolve(resumed);
		};
		removePauseListener = pauseControl.onChange((paused) => {
			if (!paused) finish(true);
		});
		signal?.addEventListener("abort", onAbort, { once: true });
		if (!pauseControl.isPaused()) finish(true);
	});
}

async function waitActiveDelay(delayMs: number, pauseControl?: ProcessPauseControl, signal?: AbortSignal): Promise<boolean> {
	let remaining = Math.max(0, delayMs);
	while (remaining > 0) {
		if (!await waitUntilResumed(pauseControl, signal)) return false;
		const startedAt = Date.now();
		const outcome = await new Promise<"elapsed" | "paused" | "aborted">((resolve) => {
			let settled = false;
			let removePauseListener = () => {};
			const timeout = setTimeout(() => finish("elapsed"), remaining);
			const onAbort = () => finish("aborted");
			const finish = (value: "elapsed" | "paused" | "aborted") => {
				if (settled) return;
				settled = true;
				clearTimeout(timeout);
				removePauseListener();
				signal?.removeEventListener("abort", onAbort);
				resolve(value);
			};
			removePauseListener = pauseControl?.onChange((paused) => {
				if (paused) finish("paused");
			}) ?? (() => {});
			signal?.addEventListener("abort", onAbort, { once: true });
			if (signal?.aborted) finish("aborted");
			else if (pauseControl?.isPaused()) finish("paused");
		});
		remaining = Math.max(0, remaining - (Date.now() - startedAt));
		if (outcome === "aborted") return false;
	}
	return !signal?.aborted;
}

function setContentBlock(message: WorkerHistoryMessage, index: number, block: Record<string, unknown>): void {
	const content = Array.isArray(message.content) ? [...message.content] : [];
	content[index] = block;
	message.content = content;
}

function updateStreamingMessage(message: WorkerHistoryMessage, event: Record<string, any>): WorkerHistoryMessage {
	const update = event.assistantMessageEvent;
	if (!isRecord(update)) return message;
	if (event.usage && isRecord(event.usage)) message.usage = event.usage;
	if (update.type === "text_start") setContentBlock(message, update.contentIndex, { type: "text", text: "" });
	if (update.type === "text_delta") {
		const content = Array.isArray(message.content) ? message.content : [];
		const current = isRecord(content[update.contentIndex]) ? content[update.contentIndex] : { type: "text", text: "" };
		setContentBlock(message, update.contentIndex, { ...current, type: "text", text: `${typeof current.text === "string" ? current.text : ""}${update.delta ?? ""}` });
	}
	if (update.type === "text_end") setContentBlock(message, update.contentIndex, { type: "text", text: update.content ?? "" });
	if (update.type === "thinking_start") setContentBlock(message, update.contentIndex, { type: "thinking", thinking: "" });
	if (update.type === "thinking_delta") {
		const content = Array.isArray(message.content) ? message.content : [];
		const current = isRecord(content[update.contentIndex]) ? content[update.contentIndex] : { type: "thinking", thinking: "" };
		setContentBlock(message, update.contentIndex, { ...current, type: "thinking", thinking: `${typeof current.thinking === "string" ? current.thinking : ""}${update.delta ?? ""}` });
	}
	if (update.type === "thinking_end") setContentBlock(message, update.contentIndex, { type: "thinking", thinking: update.content ?? "" });
	if (update.type === "toolcall_start") setContentBlock(message, update.contentIndex, { type: "toolCall", id: update.id, name: update.toolName, arguments: "" });
	if (update.type === "toolcall_delta") {
		const content = Array.isArray(message.content) ? message.content : [];
		const current = isRecord(content[update.contentIndex]) ? content[update.contentIndex] : { type: "toolCall", arguments: "" };
		setContentBlock(message, update.contentIndex, { ...current, arguments: `${typeof current.arguments === "string" ? current.arguments : ""}${update.delta ?? ""}` });
	}
	if (update.type === "toolcall_end" && isRecord(update.toolCall)) setContentBlock(message, update.contentIndex, update.toolCall);
	if (update.type === "done" && isRecord(update.message)) return update.message as WorkerHistoryMessage;
	if (update.type === "error" && isRecord(update.error)) return update.error as WorkerHistoryMessage;
	return message;
}

function summarizeError(value: unknown): string {
	const text = typeof value === "string" ? value : String(value ?? "未知瞬时 API 错误");
	return text.replace(/\s+/g, " ").trim().slice(0, 1_000) || "未知瞬时 API 错误";
}

function fallbackIsContextOverflow(message: WorkerHistoryMessage): boolean {
	const text = `${message.errorMessage ?? ""} ${extractText(message)}`;
	return /(context.?length|context.?window|maximum context|token limit|too many tokens|prompt is too long|上下文.{0,8}(过长|溢出|超限))/i.test(text);
}

function fallbackIsRetryableAssistantError(message: WorkerHistoryMessage): boolean {
	if (message.stopReason !== "error" || !message.errorMessage) return false;
	const text = message.errorMessage;
	if (/(insufficient[_ ]quota|out of budget|quota exceeded|billing|usage limit|monthly usage|available balance)/i.test(text)) return false;
	return /(overloaded|rate.?limit|too many requests|\b429\b|\b500\b|\b502\b|\b503\b|\b504\b|\b524\b|service.?unavailable|server.?error|internal.?error|provider.?returned.?error|network.?error|connection.?error|connection.?refused|connection.?lost|fetch failed|getaddrinfo|EAI_AGAIN|upstream.?connect|reset before headers|socket hang up|timed? out|timeout|terminated|websocket.?closed|websocket.?error|retry)/i.test(text);
}

async function isRetryableAssistantError(message: WorkerHistoryMessage, override?: PiRunOptions["classifyRetryableError"]): Promise<boolean> {
	if (fallbackIsContextOverflow(message)) return false;
	if (override) return override(message);
	try {
		const compat = await import("@earendil-works/pi-ai/compat");
		const assistantMessage = message as Parameters<typeof compat.isRetryableAssistantError>[0];
		if (compat.isContextOverflow(assistantMessage)) return false;
		return compat.isRetryableAssistantError(assistantMessage);
	} catch {
		return fallbackIsRetryableAssistantError(message);
	}
}

export function applyPiRunEvent(target: WorkerResult, event: PiRunEvent): void {
	target.output = event.output;
	target.stderr = event.stderr;
	target.toolCalls = event.toolCalls;
	target.history = event.history;
	target.usage = event.usage;
	target.model = event.model;
	target.stopReason = event.stopReason;
	target.errorMessage = event.errorMessage;
	target.apiRetries = event.apiRetries;
}

export async function runPiJson(options: PiRunOptions): Promise<PiRunResult> {
	const messages: WorkerHistoryMessage[] = [];
	const toolCalls: WorkerToolCall[] = [];
	const liveToolResults = new Map<string, WorkerHistoryMessage>();
	const apiRetries: ApiRetryRecord[] = [];
	let streamingMessage: WorkerHistoryMessage | undefined;
	const currentHistory = (): WorkerHistoryMessage[] => [
		...messages,
		...(streamingMessage ? [streamingMessage] : []),
		...liveToolResults.values(),
	];
	let output = "";
	let stderr = "";
	let model: string | undefined;
	let stopReason: string | undefined;
	let errorMessage: string | undefined;
	let tempDir: string | undefined;
	let tempPrompt: string | undefined;
	let sessionFile: string | undefined;
	let agentEnded = false;
	let remainingTimeout = options.timeoutMs;
	let result = 1;
	let timedOut = false;
	let aborted = false;
	let continuationReason: "pause" | "api_retry" | undefined;
	let stopBuiltInRetry: ((error: string) => void) | undefined;

	const emitUpdate = (type: string, tool?: WorkerToolCall, apiRetry?: ApiRetryRecord) => {
		options.onEvent?.({
			type,
			output,
			stderr,
			tool,
			toolCalls: [...toolCalls],
			history: currentHistory(),
			usage: usageFromHistory(currentHistory()),
			model,
			stopReason,
			errorMessage,
			apiRetries: [...apiRetries],
			apiRetry,
		});
	};

	const baseArgs = ["--mode", "json", "-p", "--no-extensions", "--approve", "--model", options.model];
	if (options.thinkingLevel) baseArgs.push("--thinking", options.thinkingLevel);
	if (options.tools && options.tools.length > 0) baseArgs.push("--tools", options.tools.join(","));

	const processLine = (line: string) => {
		if (!line.trim()) return;
		let event: any;
		try {
			event = JSON.parse(line);
		} catch {
			return;
		}
		if (event.type === "message_start" && event.message?.role === "assistant") {
			streamingMessage = event.message as WorkerHistoryMessage;
			emitUpdate("message_start");
		}
		if (event.type === "message_update") {
			streamingMessage = event.message
				? event.message as WorkerHistoryMessage
				: updateStreamingMessage(streamingMessage ?? { role: "assistant", content: [] }, event);
			const text = extractText(streamingMessage);
			if (text) output = text;
			model = streamingMessage.model ?? model;
			stopReason = streamingMessage.stopReason === "pending" ? stopReason : streamingMessage.stopReason ?? stopReason;
			errorMessage = streamingMessage.errorMessage ?? errorMessage;
			emitUpdate("message_update");
		}
		if (event.type === "message_end" && event.message) {
			const message = event.message as WorkerHistoryMessage;
			streamingMessage = undefined;
			if (message.toolCallId) liveToolResults.delete(message.toolCallId);
			messages.push(message);
			if (message.role === "assistant") {
				const text = extractText(message);
				if (text) output = text;
				model = message.model ?? model;
				stopReason = message.stopReason;
				errorMessage = message.errorMessage;
			}
			emitUpdate("message_end");
		}
		if (event.type === "tool_result_end" && event.message) {
			const message = event.message as WorkerHistoryMessage;
			if (message.toolCallId) liveToolResults.delete(message.toolCallId);
			const duplicate = message.toolCallId && messages.some((item) => item.role === message.role && item.toolCallId === message.toolCallId);
			if (!duplicate) messages.push(message);
			emitUpdate("tool_result_end");
		}
		if (event.type === "tool_execution_update" && event.toolCallId) {
			const partial = event.partialResult;
			liveToolResults.set(event.toolCallId, {
				role: "toolResult",
				toolName: event.toolName,
				toolCallId: event.toolCallId,
				content: partial?.content ?? partial,
			});
			emitUpdate("tool_execution_update");
		}
		if (event.type === "tool_execution_end" && event.toolCallId) {
			const toolResult = event.result;
			liveToolResults.set(event.toolCallId, {
				role: "toolResult",
				toolName: event.toolName,
				toolCallId: event.toolCallId,
				content: toolResult?.content ?? toolResult,
				isError: event.isError,
			});
			emitUpdate("tool_execution_end");
		}
		if (event.type === "tool_execution_start") {
			const tool = event.toolName as string | undefined;
			if (tool) {
				const call = { name: tool, args: (event.args ?? {}) as Record<string, unknown> };
				toolCalls.push(call);
				emitUpdate("tool_execution_start", call);
			}
		}
		if (event.type === "agent_start") emitUpdate("agent_start");
		if (event.type === "agent_end") {
			agentEnded = true;
			emitUpdate("agent_end");
		}
		if (event.type === "auto_retry_start") stopBuiltInRetry?.(summarizeError(event.errorMessage));
	};

	try {
		const promptFile = await writeSystemPrompt(options.systemPrompt, "system");
		tempDir = promptFile.dir;
		tempPrompt = promptFile.file;
		const sessionDir = path.join(tempDir, "sessions");
		await fs.promises.mkdir(sessionDir, { recursive: true });

		while (!timedOut && !aborted) {
			if (!await waitUntilResumed(options.pauseControl, options.signal)) {
				aborted = true;
				break;
			}
			const prompt = continuationReason === "pause"
				? "任务因用户暂停而中断。请基于当前会话和工作区状态继续完成原任务；先检查已有结果，避免重复执行已完成的操作。"
				: continuationReason === "api_retry"
					? "上一次模型请求因瞬时 API 错误失败。请基于当前会话和工作区状态继续完成原任务，避免重复执行已完成的操作。"
					: options.prompt;
			const args = [
				...baseArgs,
				"--session-dir",
				sessionDir,
				...(sessionFile ? ["--session", sessionFile] : []),
				"--append-system-prompt",
				tempPrompt,
				prompt,
			];
			const invocation = getPiInvocation(args);
			agentEnded = false;
			let lastAssistant: WorkerHistoryMessage | undefined;
			const previousMessageCount = messages.length;
			const execution = await new Promise<{ code: number; reason: "exit" | "pause" | "abort" | "timeout" | "api_retry"; retryError?: string }>((resolve) => {
				const child = spawn(invocation.command, invocation.args, {
					cwd: options.cwd,
					shell: false,
					detached: process.platform !== "win32",
					stdio: ["ignore", "pipe", "pipe"],
				});
				let buffer = "";
				let settled = false;
				let terminationReason: "pause" | "abort" | "timeout" | "api_retry" | undefined;
				let interceptedRetryError: string | undefined;
				let timeout: NodeJS.Timeout | undefined;
				let killTimeout: NodeJS.Timeout | undefined;
				let activeStartedAt = Date.now();
				let removePauseListener = () => {};

				const sendSignal = (signal: NodeJS.Signals) => {
					if (process.platform !== "win32" && child.pid) {
						try {
							process.kill(-child.pid, signal);
							return;
						} catch {
						}
					}
					child.kill(signal);
				};

				const consumeActiveTime = () => {
					if (activeStartedAt === 0 || options.timeoutMs <= 0) return;
					remainingTimeout = Math.max(0, remainingTimeout - (Date.now() - activeStartedAt));
					activeStartedAt = 0;
				};

				const finish = (code: number) => {
					if (settled) return;
					settled = true;
					consumeActiveTime();
					if (timeout) clearTimeout(timeout);
					if (killTimeout) clearTimeout(killTimeout);
					removePauseListener();
					options.signal?.removeEventListener("abort", onAbort);
					stopBuiltInRetry = undefined;
					resolve({ code, reason: terminationReason ?? "exit", retryError: interceptedRetryError });
				};

				const terminate = (reason: "pause" | "abort" | "timeout" | "api_retry") => {
					if (settled || terminationReason) return;
					terminationReason = reason;
					consumeActiveTime();
					if (timeout) {
						clearTimeout(timeout);
						timeout = undefined;
					}
					sendSignal("SIGTERM");
					killTimeout = setTimeout(() => {
						if (!settled) sendSignal("SIGKILL");
					}, 5_000);
				};

				const onAbort = () => terminate("abort");
				stopBuiltInRetry = (error) => {
					interceptedRetryError = error;
					terminate("api_retry");
				};
				removePauseListener = options.pauseControl?.onChange((paused) => {
					if (paused && !agentEnded) terminate("pause");
				}) ?? (() => {});
				options.signal?.addEventListener("abort", onAbort, { once: true });

				child.stdout.on("data", (data: Buffer) => {
					buffer += data.toString("utf8");
					const lines = buffer.split("\n");
					buffer = lines.pop() ?? "";
					for (const line of lines) processLine(line);
				});
				child.stderr.on("data", (data: Buffer) => {
					stderr += data.toString("utf8");
				});
				child.once("error", (error) => {
					errorMessage = error.message;
					finish(1);
				});
				child.once("close", (code) => {
					if (buffer.trim()) processLine(buffer);
					finish(code ?? 1);
				});

				if (options.signal?.aborted) terminate("abort");
				else if (options.pauseControl?.isPaused() && !agentEnded) terminate("pause");
				else if (options.timeoutMs > 0) timeout = setTimeout(() => {
					if (!agentEnded) terminate("timeout");
				}, remainingTimeout);
			});

			for (let index = messages.length - 1; index >= previousMessageCount; index--) {
				if (messages[index].role === "assistant") {
					lastAssistant = messages[index];
					break;
				}
			}
			result = execution.code;
			if (execution.reason === "pause") {
				streamingMessage = undefined;
				liveToolResults.clear();
				sessionFile ??= await findSessionFile(sessionDir);
				continuationReason = "pause";
				continue;
			}
			if (execution.reason === "abort") {
				aborted = true;
				break;
			}
			if (execution.reason === "timeout") {
				timedOut = true;
				break;
			}

			const retryError = execution.retryError ?? (lastAssistant?.stopReason === "error" && lastAssistant.errorMessage && await isRetryableAssistantError(lastAssistant, options.classifyRetryableError)
				? summarizeError(lastAssistant.errorMessage)
				: undefined);
			const retryDelay = (options.apiRetryDelaysMs ?? API_RETRY_DELAYS_MS).slice(0, API_RETRY_DELAYS_MS.length)[apiRetries.length];
			if (retryError && retryDelay !== undefined) {
				sessionFile ??= await findSessionFile(sessionDir);
				const retry: ApiRetryRecord = {
					timestamp: Date.now(),
					attempt: apiRetries.length + 1,
					delayMs: retryDelay,
					errorSummary: retryError,
				};
				apiRetries.push(retry);
				emitUpdate("api_retry", undefined, retry);
				if (!await waitActiveDelay(retryDelay, options.pauseControl, options.signal)) {
					aborted = true;
					break;
				}
				continuationReason = "api_retry";
				continue;
			}
			break;
		}

		const history = currentHistory();
		const limited = byteLimited(output, 50 * 1024);
		return {
			exitCode: timedOut ? 124 : aborted ? 130 : result,
			output: limited.text,
			stderr,
			toolCalls,
			model,
			stopReason: aborted ? "aborted" : stopReason,
			errorMessage: timedOut
				? `子进程累计运行超过 ${options.timeoutMs}ms 未完成`
				: aborted
					? "子进程已取消"
					: errorMessage,
			usage: usageFromHistory(history),
			history,
			apiRetries,
			outputTruncated: limited.truncated,
		};
	} finally {
		if (tempPrompt) await fs.promises.rm(tempPrompt, { force: true }).catch(() => undefined);
		if (tempDir) await fs.promises.rm(tempDir, { recursive: true, force: true }).catch(() => undefined);
	}
}

export function isSuccessfulWorker(result: WorkerResult): boolean {
	return result.exitCode === 0 && !result.errorMessage && result.stopReason !== "error" && result.stopReason !== "aborted";
}