import { performance } from "node:perf_hooks";

import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
	InputEvent,
	SessionBeforeTreeEvent as PiSessionBeforeTreeEvent,
	SessionTreeEvent as PiSessionTreeEvent,
} from "@earendil-works/pi-coding-agent";

import type { OperationResult, SessionTreeEvent, UndoController } from "../src/controller.ts";
import { browseDiff } from "../src/diff-ui.ts";
import { computeCheckpointDiff, type DiffSource, formatDiffSummary, sanitizeDisplayText } from "../src/diff-view.ts";
import { createPiUndoRuntime } from "../src/pi-runtime.ts";
import { StatusReporter } from "../src/status-reporter.ts";

export interface PiUndoRuntime {
	readonly controller: UndoController;
	readonly reporter: StatusReporter;
	readonly diffSource?: DiffSource;
	readonly recovery?: { readonly reason?: string; readonly files?: number; readonly opId?: string };
	setCommandContext?(context: ExtensionCommandContext | undefined): void;
	isInternalNavigation?(): boolean;
	normalizeTreeEvent?(event: PiSessionTreeEvent): SessionTreeEvent;
	dispose?(): Promise<void>;
}

export type PiUndoRuntimeFactory = (
	context: ExtensionContext,
	pi: ExtensionAPI,
) => Promise<PiUndoRuntime>;

type DeferredImage = NonNullable<InputEvent["images"]>[number];

interface DeferredPrompt {
	readonly text: string;
	readonly images?: readonly DeferredImage[];
}

export function createPiUndoExtension(runtimeFactory: PiUndoRuntimeFactory): (pi: ExtensionAPI) => void {
	return (pi) => {
		if (process.env.PI_SUBAGENT_CHILD === "1") {
			// pi-subagents 的 runner 子进程会加载 ambient extensions，并用该环境变量
			// 标记自身（其父扩展据此保持惰性）。pi-undo 只在主会话生效：避免子进程
			// 重复 capture 争用 workspace lock。子代理的工作区变更由父会话的 run 级
			// before/after 快照覆盖，撤回语义不受影响。
			return;
		}
		let runtime: PiUndoRuntime | undefined;
		let runtimeContext: ExtensionContext | undefined;
		let generation = 0;
		let deferredPrompts: DeferredPrompt[] = [];
		let replaying: DeferredPrompt | undefined;
		let acceptedReplay: DeferredPrompt | undefined;
		let activeCommands = new Set<symbol>();
		let activeAction: "undo" | "redo" | undefined;
		let captureFailureNotified = false;
		let recoveryHintNotified = false;
		let clearTreeWatch: (() => void) | undefined;
		let initializing = false;

		const initialize = async (context: ExtensionContext): Promise<void> => {
			if (initializing) return;
			initializing = true;
			try {
				clearTreeWatch?.();
				if (runtime?.dispose !== undefined) await runtime.dispose();
				else await runtime?.controller.dispose?.();
			} catch (error) {
				initializing = false;
				runtime?.reporter.setRecoveryRequired(errorMessage(error));
				context.ui.notify(`旧任务尚未安全结束：${errorMessage(error)}`, "error");
				return;
			}
			const currentGeneration = ++generation;
			restoreDeferredPrompts(context);
			runtimeContext = context;
			deferredPrompts = [];
			replaying = undefined;
			acceptedReplay = undefined;
			activeCommands = new Set<symbol>();
			activeAction = undefined;
			captureFailureNotified = false;
			recoveryHintNotified = false;
			try {
				const next = await runtimeFactory(context, pi);
				if (currentGeneration !== generation) return;
				runtime = next;
				await next.controller.recover();
				const history = next.controller.history();
				if (history.locked) {
					next.reporter.setRecoveryRequired(
						next.controller.recoveryReason?.() ?? next.recovery?.reason ?? "pending journal",
						next.recovery,
					);
				} else next.reporter.setReady(history.undoCount, history.redoCount);
			} catch (error) {
				if (currentGeneration !== generation) return;
				runtime = undefined;
				new StatusReporter(context).setRecoveryRequired(errorMessage(error));
			} finally {
				initializing = false;
			}
		};

		const dispatchDeferredPrompt = (active: PiUndoRuntime, expectedGeneration: number): void => {
			if (
				initializing || expectedGeneration !== generation || runtime !== active || activeCommands.size > 0 ||
				replaying !== undefined || deferredPrompts.length === 0 || active.controller.history().locked
			) return;
			const prompt = deferredPrompts[0]!;
			replaying = prompt;
			acceptedReplay = undefined;
			queueMicrotask(() => {
				if (initializing || expectedGeneration !== generation || runtime !== active || replaying !== prompt) return;
				try {
					pi.sendUserMessage(prompt.images === undefined || prompt.images.length === 0
						? prompt.text
						: [{ type: "text" as const, text: prompt.text }, ...prompt.images]);
				} catch (error) {
					replaying = undefined;
					acceptedReplay = undefined;
					deferredPrompts.shift();
					restoreEditorText(runtimeContext, prompt.text);
					runtimeContext?.ui.notify(`Unable to replay queued prompt: ${errorMessage(error)}`, "warning");
				}
			});
		};

		const restoreDeferredPrompts = (context: ExtensionContext | undefined): void => {
			if (deferredPrompts.length === 0) return;
			const prompts = deferredPrompts.splice(0);
			replaying = undefined;
			acceptedReplay = undefined;
			restoreEditorText(context, prompts.map((prompt) => prompt.text).join("\n\n"));
			if (prompts.some((prompt) => (prompt.images?.length ?? 0) > 0)) {
				context?.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
			}
		};

		const resumeDeferredPrompts = (
			active: PiUndoRuntime,
			expectedGeneration: number,
			lockedReason: string,
		): void => {
			if (expectedGeneration !== generation || runtime !== active) return;
			const history = active.controller.history();
			if (history.locked) {
				active.reporter.setRecoveryRequired(
					active.controller.recoveryReason?.() ?? lockedReason,
				);
				restoreDeferredPrompts(runtimeContext);
			} else {
				active.reporter.setReady(history.undoCount, history.redoCount);
				dispatchDeferredPrompt(active, expectedGeneration);
			}
		};

		const runCommand = async (
			action: "undo" | "redo",
			context: ExtensionCommandContext,
		): Promise<void> => {
			const active = runtime;
			const commandGeneration = generation;
			if (active === undefined || initializing) {
				context.ui.notify("pi-undo session unavailable", "warning");
				return;
			}
			const commandToken = Symbol(action);
			const commandSet = activeCommands;
			// 只有当前执行命令能安装/清理导航上下文；busy 的第二个命令不得覆盖或清空它。
			const ownsCommandContext = commandSet.size === 0;
			commandSet.add(commandToken);
			if (ownsCommandContext) activeAction = action;
			if (ownsCommandContext) active.reporter.setPhase(action === "undo" ? "undoing" : "redoing");
			if (ownsCommandContext) active.setCommandContext?.(context);
			const commandStarted = performance.now();
			let result: OperationResult;
			try {
				result = action === "undo"
					? await active.controller.undo()
					: await active.controller.redo();
			} finally {
				if (ownsCommandContext) active.setCommandContext?.(undefined);
				commandSet.delete(commandToken);
				if (commandSet === activeCommands && commandSet.size === 0) activeAction = undefined;
			}
			if (commandGeneration !== generation || runtime !== active) return;
			const history = active.controller.history();
			if (history.locked && result.code === "busy") {
				result = {
					...result,
					code: "recovery_required",
					message: active.controller.recoveryReason?.() ?? "pending journal",
				};
			}
			if (!ownsCommandContext) {
				context.ui.notify(`${result.code} files:${result.changedFiles}`, "warning");
				return;
			}
			active.reporter.result(result, performance.now() - commandStarted);
			if (result.code === "recovery_required" && !recoveryHintNotified) {
				recoveryHintNotified = true;
				context.ui.notify("pi-undo: run /undo-recover to retry recovery", "info");
			}
			const hasDeferredPrompt = deferredPrompts.length > 0 || replaying !== undefined;
			if (
				action === "undo" && result.code === "ok" && result.refillPrompt !== undefined &&
				!hasDeferredPrompt
			) active.reporter.refillPrompt(result.refillPrompt);
			resumeDeferredPrompts(active, commandGeneration, result.message ?? result.code);
		};

		const runDiff = async (args: string, context: ExtensionCommandContext): Promise<void> => {
			const active = runtime;
			if (active === undefined || active.diffSource === undefined) {
				context.ui.notify("pi-undo session unavailable", "warning");
				return;
			}
			const checkpoints = active.controller.listCheckpoints();
			if (checkpoints.length === 0) {
				context.ui.notify("No recorded Agent runs to diff", "info");
				return;
			}
			const trimmed = args.trim();
			const position = trimmed === "" ? 1 : Number(trimmed);
			if (!Number.isInteger(position) || position < 1 || position > checkpoints.length) {
				context.ui.notify(`Invalid run number; recorded runs: ${checkpoints.length}`, "warning");
				return;
			}
			// 1 = 最近一次 run（栈顶）。
			const checkpoint = checkpoints[checkpoints.length - position]!;
			let diffs;
			try {
				diffs = await computeCheckpointDiff(active.diffSource, checkpoint);
			} catch (error) {
				context.ui.notify(`Unable to load diff: ${sanitizeDisplayText(errorMessage(error), 120)}`, "error");
				return;
			}
			if (runtime !== active) return;
			if (diffs.length === 0) {
				context.ui.notify("This run changed no files", "info");
				return;
			}
			const label = runLabel(checkpoint.rawPrompt, position);
			if (context.mode !== "tui") {
				context.ui.notify(`${label} — ${formatDiffSummary(diffs)}`, "info");
				return;
			}
			await browseDiff(context, label, diffs);
		};

		pi.registerCommand("undo", {
			description: "Undo the last completed Agent run",
			handler: async (_args: string, context: ExtensionCommandContext) => runCommand("undo", context),
		});
		pi.registerCommand("redo", {
			description: "Redo the last undone Agent run",
			handler: async (_args: string, context: ExtensionCommandContext) => runCommand("redo", context),
		});
		pi.registerCommand("diff", {
			description: "Review files changed by an Agent run (latest, or /diff N)",
			handler: async (args: string, context: ExtensionCommandContext) => runDiff(args, context),
		});
		pi.registerCommand("undo-cancel", {
			description: "安全停止正在执行的撤回或重做",
			handler: async (_args: string, context: ExtensionCommandContext) => {
				const requested = runtime?.controller.cancelOperation?.() ?? false;
				context.ui.notify(requested ? "已请求停止，正在等待写入结束并恢复一致状态" : "当前没有可取消的撤回操作", "info");
			},
		});
		pi.registerCommand("undo-recover", {
			description: "Re-run pi-undo recovery and refresh undo history",
			handler: async (_args: string, context: ExtensionCommandContext) => {
				if (activeCommands.size > 0) {
					context.ui.notify("pi-undo: wait for the current operation to finish before recovering", "warning");
					return;
				}
				// 原地重建 runtime：等价于重启窗口，复用启动 recovery 语义。
				await initialize(context);
				const active = runtime;
				if (active === undefined) return;
				const history = active.controller.history();
				if (history.locked) {
					context.ui.notify(
						`pi-undo: still locked (${active.controller.recoveryReason?.() ?? "pending journal"}); resolve the blocking session, then run /undo-recover again`,
						"warning",
					);
					return;
				}
				context.ui.notify(`pi-undo: recovery complete (undo:${history.undoCount} redo:${history.redoCount})`, "info");
			},
		});

		pi.on("session_start", async (_event: unknown, context: ExtensionContext) => initialize(context));
		pi.on("input", async (event: InputEvent, context: ExtensionContext) => {
			const active = runtime;
			if (initializing) {
				restoreEditorText(context, event.text);
				return { action: "handled" as const };
			}
			if (active === undefined) return { action: "continue" as const };
			const inputContext = { streaming: event.streamingBehavior !== undefined };
			const result = active.controller.beginInput !== undefined
				? active.controller.beginInput(event.text, inputContext)
				: await active.controller.prepareInput(event.text, inputContext);
			const replay = replaying;
			if (result.action === "defer") {
				if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
					replaying = undefined;
					acceptedReplay = undefined;
					active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
					return { action: "handled" as const };
				}
				if (event.text.trimStart().startsWith("/")) {
					restoreEditorText(context, event.text);
					context.ui.notify("Command input preserved until undo/redo completes", "info");
					return { action: "handled" as const };
				}
				deferredPrompts.push({
					text: event.text,
					...(event.images === undefined ? {} : { images: event.images.map((image) => ({ ...image })) }),
				});
				active.reporter.setPhase(`${activeAction ?? "operation"} queued:${deferredPrompts.length}`);
				return { action: "handled" as const };
			}
			if (replay !== undefined && event.source === "extension" && samePrompt(event.text, event.images, replay)) {
				if (result.action === "continue") {
					acceptedReplay = replay;
				} else {
					removeDeferredPrompt(deferredPrompts, replay);
					replaying = undefined;
					acceptedReplay = undefined;
					restoreEditorText(context, replay.text);
					if ((replay.images?.length ?? 0) > 0) {
						context.ui.notify("Queued prompt text restored; image attachments must be reattached", "warning");
					}
				}
			}
			if (
				active.controller.beginInput === undefined &&
				result.action === "continue" && active.controller.captureFailed() && !captureFailureNotified
			) {
				captureFailureNotified = true;
				const reason = active.controller.captureFailureReason();
				context.ui.notify(
					`pi-undo: pre-input snapshot failed${reason === undefined || reason.length === 0 ? "" : ` (${reason})`}; this run will not be undoable`,
					"warning",
				);
			}
			return result;
		});
		pi.on("message_end", async (event, context: ExtensionContext) => {
			if (event.message.role !== "user") return;
			const active = runtime;
			const messageGeneration = generation;
			if (active === undefined || active.controller.commitInput === undefined) return;
			await active.controller.commitInput();
			if (runtime !== active || generation !== messageGeneration) return;
			if (active.controller.captureFailed() && !captureFailureNotified) {
				captureFailureNotified = true;
				const reason = active.controller.captureFailureReason();
				context.ui.notify(
					`pi-undo: pre-input snapshot failed${reason === undefined || reason.length === 0 ? "" : ` (${reason})`}; this run will not be undoable`,
					"warning",
				);
			}
		});
		pi.on("before_agent_start", async () => {
			const active = runtime;
			const startGeneration = generation;
			if (active === undefined) return;
			await active.controller.beforeAgentStart();
			const replay = replaying;
			if (
				runtime === active && generation === startGeneration && replay !== undefined &&
				acceptedReplay === replay
			) {
				removeDeferredPrompt(deferredPrompts, replay);
				replaying = undefined;
				acceptedReplay = undefined;
			}
		});
		pi.on("agent_settled", async () => {
			const active = runtime;
			const settledGeneration = generation;
			if (active === undefined) return;
			await active.controller.agentSettled();
			if (runtime !== active || generation !== settledGeneration) return;
			if (activeCommands.size === 0) resumeDeferredPrompts(
				active,
				settledGeneration,
				active.controller.recoveryReason?.() ?? "session state ambiguous",
			);
		});
		pi.on("session_before_tree", async (event: PiSessionBeforeTreeEvent, context: ExtensionContext) => {
			if (runtime === undefined || initializing) return { cancel: true };
			if (runtime.isInternalNavigation?.()) return undefined;
			const active = runtime;
			if (event.signal?.aborted) return { cancel: true };
			clearTreeWatch?.();
			let timer: ReturnType<typeof setInterval> | undefined;
			const clear = (): void => {
				if (timer !== undefined) clearInterval(timer);
				event.signal?.removeEventListener("abort", cancel);
				if (clearTreeWatch === clear) clearTreeWatch = undefined;
			};
			const cancel = (): void => { void active.controller.cancelTree?.().catch(() => {}); };
			clearTreeWatch = clear;
			event.signal?.addEventListener("abort", cancel, { once: true });
			const result = await active.controller.beforeTree({ targetLeafId: event.preparation.targetId, signal: event.signal });
			if (result !== undefined || event.signal?.aborted) {
				await active.controller.cancelTree?.();
				clear();
				return { cancel: true };
			}
			// 摘要错误或其他扩展取消可能没有 session_tree；Pi 离开导航后才终结准备事务。
			timer = setInterval(() => {
				if (!context.isIdle()) return;
				clear();
				void active.controller.cancelTree?.().then(() => {
					resumeDeferredPrompts(active, generation, "tree_navigation_incomplete");
				});
			}, 100);
			timer.unref();
			return undefined;
		});
		pi.on("session_tree", async (event: PiSessionTreeEvent) => {
			const active = runtime;
			if (active?.isInternalNavigation?.()) return;
			const treeGeneration = generation;
			clearTreeWatch?.();
			await active?.controller.afterTree(active.normalizeTreeEvent?.(event) ?? {
				newLeafId: event.newLeafId,
				navigationTargetLeafId: event.summaryEntry === undefined ? event.newLeafId : event.summaryEntry.parentId,
			});
			if (active !== undefined && runtime === active && generation === treeGeneration) {
				resumeDeferredPrompts(
					active,
					treeGeneration,
					active.controller.recoveryReason?.() ?? "session state ambiguous",
				);
			}
		});
		pi.on("session_shutdown", async () => {
			clearTreeWatch?.();
			generation += 1;
			deferredPrompts = [];
			replaying = undefined;
			acceptedReplay = undefined;
			activeCommands = new Set<symbol>();
			activeAction = undefined;
			runtimeContext = undefined;
			if (runtime?.dispose !== undefined) await runtime.dispose();
			else await runtime?.controller.dispose?.();
			await runtime?.controller.cancelTree?.();
			runtime?.reporter.clear();
			runtime = undefined;
		});
	};
}

function samePrompt(
	text: string,
	images: readonly DeferredImage[] | undefined,
	prompt: DeferredPrompt,
): boolean {
	if (text !== prompt.text || (images?.length ?? 0) !== (prompt.images?.length ?? 0)) return false;
	return (images ?? []).every((image, index) => JSON.stringify(image) === JSON.stringify(prompt.images?.[index]));
}

function removeDeferredPrompt(prompts: DeferredPrompt[], prompt: DeferredPrompt): void {
	const index = prompts.indexOf(prompt);
	if (index >= 0) prompts.splice(index, 1);
}

function restoreEditorText(context: ExtensionContext | undefined, text: string): void {
	if (context === undefined || text.length === 0) return;
	const current = context.ui.getEditorText();
	if (current === text || current.startsWith(`${text}\n\n`)) return;
	context.ui.setEditorText(current.length === 0 ? text : `${text}\n\n${current}`);
}

function errorMessage(error: unknown): string {
	return error instanceof Error ? error.message : String(error);
}

function runLabel(rawPrompt: string, position: number): string {
	const singleLine = sanitizeDisplayText(rawPrompt, 1_000);
	const preview = singleLine.length > 48 ? `${singleLine.slice(0, 47)}…` : singleLine;
	return preview === "" ? `Run #${position}` : `Run #${position}: ${preview}`;
}

export default createPiUndoExtension(createPiUndoRuntime);