import { randomUUID } from "node:crypto";
import type { ClusterConfig } from "./types.ts";
import { reviewWorkerResult } from "./reviewer.ts";
import { persistSnapshot, persistTaskOutput, type RunPersistence } from "./state.ts";
import {
	COMPLETE_EXECUTION_DATA_VERSION,
	LEVEL_ORDER,
	emptyUsage,
	isTerminalStatus,
	type ClusterEvent,
	type ClusterInput,
	type ClusterLevel,
	type ClusterRunResult,
	type ClusterState,
	type ReviewResult,
	type TaskAttempt,
	type TaskRuntime,
	type TaskLevelSelection,
	type UserDecision,
	type UserDecisionRequest,
	type WorkerResult,
} from "./types.ts";
import { clusterUsage } from "./usage.ts";
import { applyPiRunEvent, isSuccessfulWorker, runPiJson, type ProcessPauseControl } from "./worker.ts";

export interface SchedulerOptions {
	runId?: string;
	input: ClusterInput;
	cwd: string;
	config: ClusterConfig;
	reviewerModel: string;
	reviewerThinkingLevel?: ClusterConfig["reviewer"]["thinkingLevel"];
	signal?: AbortSignal;
	persistence?: RunPersistence;
	levelSelections?: Map<string, TaskLevelSelection>;
	onChange?: (state: ClusterState) => void;
	onPersist?: (state: ClusterState, result?: ClusterRunResult) => void;
	onRunFinished?: (state: ClusterState, result: ClusterRunResult) => Promise<void>;
	onUserDecision?: (request: UserDecisionRequest) => void;
	apiRetryDelaysMs?: readonly number[];
}

export class ClusterScheduler {
	readonly state: ClusterState;
	private readonly config: ClusterConfig;
	private readonly cwd: string;
	private readonly reviewerModel: string;
	private readonly reviewerThinkingLevel: SchedulerOptions["reviewerThinkingLevel"];
	private readonly parentSignal?: AbortSignal;
	private readonly controller = new AbortController();
	private readonly pauseListeners = new Set<(paused: boolean) => void>();
	private readonly pauseControl: ProcessPauseControl;
	private readonly persistence?: RunPersistence;
	private readonly levelSelections: Map<string, TaskLevelSelection>;
	private readonly onChange?: SchedulerOptions["onChange"];
	private readonly onPersist?: SchedulerOptions["onPersist"];
	private readonly onRunFinished?: SchedulerOptions["onRunFinished"];
	private readonly onUserDecision?: SchedulerOptions["onUserDecision"];
	private readonly apiRetryDelaysMs?: readonly number[];
	private readonly active = new Map<string, Promise<void>>();
	private readonly decisions = new Map<string, { resolve: (decision: UserDecision) => void }>();
	private resumeWaiter: (() => void) | undefined;
	private hasRun = false;
	private completion?: Promise<ClusterRunResult>;

	constructor(options: SchedulerOptions) {
		this.config = options.config;
		this.cwd = options.cwd;
		this.reviewerModel = options.reviewerModel;
		this.reviewerThinkingLevel = options.reviewerThinkingLevel;
		this.parentSignal = options.signal;
		this.persistence = options.persistence;
		this.levelSelections = options.levelSelections ?? new Map();
		this.onChange = options.onChange;
		this.onPersist = options.onPersist;
		this.onRunFinished = options.onRunFinished;
		this.onUserDecision = options.onUserDecision;
		this.apiRetryDelaysMs = options.apiRetryDelaysMs;
		const tasks = options.input.tasks.map((task) => {
			const selection = this.levelSelections.get(task.id);
			const selectedLevel = selection?.initialLevel ?? task.level;
			const initialLevel = LEVEL_ORDER[Math.max(LEVEL_ORDER.indexOf(task.level), LEVEL_ORDER.indexOf(selectedLevel))] ?? task.level;
			return {
				id: task.id,
				title: task.title,
				taskType: task.taskType,
				task: task.task,
				acceptanceCriteria: [...task.acceptanceCriteria],
				dependsOn: [...(task.dependsOn ?? [])],
				cwd: task.cwd,
				requestedLevel: task.level,
				initialLevel,
				levelSelectionReason: selection?.levelSelectionReason ?? `主 agent 请求 ${task.level}；无可用学习证据。`,
				level: initialLevel,
				status: "queued" as const,
				attempts: [],
			};
		});
		this.state = {
			runId: options.runId ?? randomUUID(),
			goal: options.input.goal,
			cwd: options.cwd,
			executionDataVersion: COMPLETE_EXECUTION_DATA_VERSION,
			status: "running",
			tasks,
			events: tasks.map((task) => ({
				timestamp: Date.now(),
				taskId: task.id,
				kind: "state" as const,
				message: `任务 ${task.id} 等级选择：请求 ${task.requestedLevel}，实际初始 ${task.initialLevel}。${task.levelSelectionReason}`,
			})),
			startedAt: Date.now(),
			paused: false,
			pausePeriods: [],
		};
		this.pauseControl = {
			isPaused: () => this.state.paused,
			onChange: (listener) => {
				this.pauseListeners.add(listener);
				return () => this.pauseListeners.delete(listener);
			},
		};
		if (this.parentSignal) {
			if (this.parentSignal.aborted) this.cancel();
			else this.parentSignal.addEventListener("abort", () => this.cancel(), { once: true });
		}
	}

	get signal(): AbortSignal {
		return this.controller.signal;
	}

	private notifyPauseChanged(): void {
		for (const listener of this.pauseListeners) listener(this.state.paused);
	}

	private setPaused(paused: boolean, timestamp = Date.now()): void {
		if (this.state.paused === paused) return;
		if (paused) {
			(this.state.pausePeriods ??= []).push({ startedAt: timestamp });
		} else {
			const period = this.state.pausePeriods?.findLast((item) => item.finishedAt === undefined);
			if (period) period.finishedAt = timestamp;
		}
		this.state.paused = paused;
		this.notifyPauseChanged();
	}

	cancel(): void {
		const wasActive = this.state.status === "running" || this.state.status === "paused";
		const finishedAt = Date.now();
		this.setPaused(false, finishedAt);
		if (!this.controller.signal.aborted) this.controller.abort();
		if (wasActive) {
			this.state.status = "cancelled";
			this.state.finishedAt = finishedAt;
			for (const task of this.state.tasks) {
				if (!isTerminalStatus(task.status)) {
					task.status = "cancelled";
					task.finishedAt = finishedAt;
				}
			}
		}
		for (const waiter of this.decisions.values()) waiter.resolve({ action: "abandon" });
		this.decisions.clear();
		this.resumeWaiter?.();
		this.resumeWaiter = undefined;
		this.emit("control", "已请求取消集群", undefined, true);
	}

	async waitForCompletion(): Promise<void> {
		await this.completion?.catch(() => undefined);
	}

	togglePause(): void {
		if (this.controller.signal.aborted) return;
		if (this.state.tasks.some((task) => task.status === "paused_for_user")) return;
		if (this.state.status !== "running" && this.state.status !== "paused") return;
		this.setPaused(!this.state.paused);
		this.state.status = this.state.paused ? "paused" : "running";
		this.emit("control", this.state.paused ? "已暂停集群" : "已恢复集群", undefined, true);
		if (!this.state.paused) {
			this.resumeWaiter?.();
			this.resumeWaiter = undefined;
		}
	}

	decide(taskId: string, action: UserDecision["action"]): void {
		const waiter = this.decisions.get(taskId);
		if (!waiter) return;
		this.decisions.delete(taskId);
		if (this.decisions.size === 0) {
			this.setPaused(false);
			if (this.state.status === "paused") this.state.status = "running";
			this.resumeWaiter?.();
			this.resumeWaiter = undefined;
		}
		this.emit("control", `用户对任务 ${taskId} 选择了 ${action}`, taskId, true);
		waiter.resolve({ action });
	}

	async run(): Promise<ClusterRunResult> {
		if (this.hasRun) throw new Error("一个调度器实例只能运行一次");
		this.hasRun = true;
		this.completion = this.runInternal();
		return this.completion;
	}

	private async runInternal(): Promise<ClusterRunResult> {
		this.emit("run", `集群启动：${this.state.goal}`, undefined, true);
		if (this.persistence) await persistSnapshot(this.persistence, this.state);

		while (!this.controller.signal.aborted) {
			await this.waitIfPaused();
			if (this.controller.signal.aborted) break;
			this.markBlockedTasks();
			const ready = this.state.tasks.filter(
				(task) => task.status === "queued" && task.dependsOn.every((dependency) => this.taskById(dependency)?.status === "completed"),
			);
			const slots = Math.max(0, this.config.maxConcurrency - this.active.size);
			for (const task of ready.slice(0, slots)) this.launch(task);

			if (this.active.size > 0) {
				const finishedTaskId = await Promise.race(
					Array.from(this.active.entries()).map(async ([taskId, promise]) => {
						await promise;
						return taskId;
					}),
				);
				this.active.delete(finishedTaskId);
				continue;
			}

			const unfinished = this.state.tasks.some((task) => !isTerminalStatus(task.status));
			if (!unfinished) break;
			this.markBlockedTasks();
			const stillUnfinished = this.state.tasks.some((task) => !isTerminalStatus(task.status));
			if (stillUnfinished) {
				this.state.status = "failed";
				this.emit("error", "任务图没有可执行任务，可能存在循环依赖", undefined, true);
				break;
			}
		}

		if (this.controller.signal.aborted) await Promise.allSettled(this.active.values());
		for (const task of this.state.tasks) {
			if (!isTerminalStatus(task.status)) {
				task.status = "cancelled";
				task.finishedAt ??= Date.now();
			}
		}
		if (this.controller.signal.aborted) this.state.status = "cancelled";
		else if (this.state.tasks.every((task) => task.status === "completed")) this.state.status = "completed";
		else if (this.state.tasks.some((task) => ["failed", "blocked", "timed_out"].includes(task.status))) this.state.status = "failed";
		else this.state.status = "completed";
		this.state.finishedAt = Date.now();
		this.setPaused(false, this.state.finishedAt);
		const result = this.buildResult();
		this.emit("run", `集群结束：${result.status}`, undefined, true);
		if (this.persistence) await persistSnapshot(this.persistence, this.state, result);
		this.onPersist?.(this.state, result);
		try {
			await this.onRunFinished?.(this.state, result);
		} catch (error) {
			this.emit("error", `全局学习证据写入失败：${error instanceof Error ? error.message : String(error)}`, undefined, true);
		}
		return result;
	}

	private launch(task: TaskRuntime): void {
		const promise = this.executeTask(task).catch((error) => {
			task.status = "failed";
			task.error = error instanceof Error ? error.message : String(error);
			task.finishedAt = Date.now();
			this.emit("error", `任务 ${task.id} 执行异常：${task.error}`, task.id, true);
		});
		this.active.set(task.id, promise);
	}

	private async executeTask(task: TaskRuntime): Promise<void> {
		task.startedAt = Date.now();
		let levelIndex = LEVEL_ORDER.indexOf(task.level);
		let feedback = "";

		while (!this.controller.signal.aborted && levelIndex < LEVEL_ORDER.length) {
			const level = LEVEL_ORDER[levelIndex];
			let retries = 0;
			while (!this.controller.signal.aborted) {
				const levelConfig = this.config.levels[level];
				const attempt: TaskAttempt = {
					level,
					attempt: task.attempts.filter((item) => item.level === level).length + 1,
					model: levelConfig.model,
					startedAt: Date.now(),
				};
				task.level = level;
				task.status = "running";
				task.attempts.push(attempt);
				this.emit("state", `任务 ${task.id} 使用 ${level} 级 worker 开始第 ${attempt.attempt} 次尝试`, task.id, true);
				const liveWorker: WorkerResult = {
					exitCode: -1,
					output: "",
					stderr: "",
					toolCalls: [],
					model: levelConfig.model,
					usage: emptyUsage(),
					history: [],
					apiRetries: [],
				};
				attempt.worker = liveWorker;
				this.refresh();
				const worker = await runPiJson({
					cwd: task.cwd ?? this.cwd,
					model: levelConfig.model,
					thinkingLevel: levelConfig.thinkingLevel,
					tools: levelConfig.tools,
					systemPrompt: this.workerSystemPrompt(level),
					prompt: this.workerPrompt(task, feedback),
					timeoutMs: levelConfig.timeoutMs ?? 900_000,
					signal: this.controller.signal,
					pauseControl: this.pauseControl,
					apiRetryDelaysMs: this.apiRetryDelaysMs,
					onEvent: (event) => {
						applyPiRunEvent(liveWorker, event);
						task.output = liveWorker.output.slice(-20_000);
						if (event.apiRetry) {
							this.emit("worker", `任务 ${task.id} 的 worker 遇到瞬时 API 错误，将在 ${event.apiRetry.delayMs}ms 后进行第 ${event.apiRetry.attempt} 次 API 重试`, task.id, true);
						} else if (event.tool) this.emit("worker", `${event.tool.name} 正在执行`, task.id);
						else this.refresh();
					},
				});
				attempt.worker = worker;
				task.output = worker.output;
				this.emit("worker", `任务 ${task.id} 的 ${level} 级 worker 已返回`, task.id, true);
				if (this.persistence && worker.output) task.outputPath = await persistTaskOutput(this.persistence, task, worker.output);

				task.status = "reviewing";
				this.emit("review", `任务 ${task.id} 正在审核`, task.id, true);
				let review: ReviewResult;
				if (isSuccessfulWorker(worker)) {
					const liveReviewer: WorkerResult = {
						exitCode: -1,
						output: "",
						stderr: "",
						toolCalls: [],
						model: this.reviewerModel,
						usage: emptyUsage(),
						history: [],
						apiRetries: [],
					};
					attempt.reviewer = liveReviewer;
					this.refresh();
					const execution = await reviewWorkerResult(task, worker, {
						cwd: task.cwd ?? this.cwd,
						model: this.reviewerModel,
						thinkingLevel: this.reviewerThinkingLevel,
						tools: this.config.reviewer.tools,
						timeoutMs: this.config.reviewer.timeoutMs ?? 180_000,
						signal: this.controller.signal,
						pauseControl: this.pauseControl,
						apiRetryDelaysMs: this.apiRetryDelaysMs,
						onEvent: (event) => {
							applyPiRunEvent(liveReviewer, event);
							task.output = liveReviewer.output.slice(-20_000);
							if (event.apiRetry) {
								this.emit("review", `任务 ${task.id} 的审核器遇到瞬时 API 错误，将在 ${event.apiRetry.delayMs}ms 后进行第 ${event.apiRetry.attempt} 次 API 重试`, task.id, true);
							} else this.refresh();
						},
					});
					attempt.reviewer = execution.reviewer;
					review = execution.review;
				} else {
					review = this.failedWorkerReview(worker, retries < this.config.maxRetriesPerLevel);
				}
				attempt.review = review;
				task.review = review;
				attempt.finishedAt = Date.now();
				this.emit("review", `任务 ${task.id} 审核结论：${review.decision}，${review.reason}`, task.id, true);

				if (this.controller.signal.aborted) {
					task.status = "cancelled";
					task.finishedAt = Date.now();
					return;
				}
				if (review.decision === "timeout") {
					task.status = "timed_out";
					task.error = review.reason;
					task.finishedAt = Date.now();
					this.emit("error", `任务 ${task.id} 超时：${review.reason}`, task.id, true);
					return;
				}
				if (review.decision === "error") {
					task.status = "failed";
					task.error = review.reason;
					task.finishedAt = Date.now();
					this.emit("error", `任务 ${task.id} 审核失败：${review.reason}`, task.id, true);
					return;
				}

				if (review.decision === "pass") {
					task.status = "completed";
					task.finishedAt = Date.now();
					this.emit("state", `任务 ${task.id} 已完成`, task.id, true);
					return;
				}

				feedback = review.nextInstruction || review.reason;
				if (review.decision === "retry" && retries < this.config.maxRetriesPerLevel) {
					retries += 1;
					this.emit("state", `任务 ${task.id} 将在 ${level} 级重试`, task.id, true);
					continue;
				}

				if ((review.decision === "retry" || review.decision === "escalate") && levelIndex < LEVEL_ORDER.length - 1) {
					const nextLevel = LEVEL_ORDER[levelIndex + 1];
					this.emit("control", `任务 ${task.id} 从 ${level} 级升级到 ${nextLevel} 级`, task.id, true);
					levelIndex += 1;
					break;
				}

				const decision = await this.waitForUser(task, review);
				if (decision.action === "accept") {
					task.status = "completed";
					task.finishedAt = Date.now();
					return;
				}
				if (decision.action === "retry") {
					retries = 0;
					continue;
				}
				if (decision.action === "escalate" && levelIndex < LEVEL_ORDER.length - 1) {
					levelIndex += 1;
					break;
				}
				task.status = this.controller.signal.aborted ? "cancelled" : "failed";
				task.error = decision.action === "abandon" ? "用户放弃任务" : review.reason;
				task.finishedAt = Date.now();
				return;
			}
		}

		if (this.controller.signal.aborted) {
			task.status = "cancelled";
			task.finishedAt = Date.now();
		}
	}

	private async waitForUser(task: TaskRuntime, review: ReviewResult): Promise<UserDecision> {
		if (this.controller.signal.aborted) return { action: "abandon" };
		task.status = "paused_for_user";
		this.setPaused(true);
		this.state.status = "paused";
		const request: UserDecisionRequest = {
			runId: this.state.runId,
			taskId: task.id,
			title: task.title,
			level: task.level,
			review,
		};
		const decision = new Promise<UserDecision>((resolve) => {
			this.decisions.set(task.id, { resolve });
		});
		this.emit("control", `任务 ${task.id} 需要用户决策：${review.reason}`, task.id, true);
		if (this.onUserDecision) {
			try {
				this.onUserDecision(request);
			} catch {
				this.decide(task.id, "abandon");
			}
		} else {
			queueMicrotask(() => this.decide(task.id, "abandon"));
		}
		return decision;
	}

	private async waitIfPaused(): Promise<void> {
		if (!this.state.paused || this.controller.signal.aborted) return;
		await new Promise<void>((resolve) => {
			const finish = () => {
				if (this.resumeWaiter === finish) this.resumeWaiter = undefined;
				resolve();
			};
			this.resumeWaiter = finish;
			if (!this.state.paused || this.controller.signal.aborted) finish();
		});
	}

	private taskById(taskId: string): TaskRuntime | undefined {
		return this.state.tasks.find((task) => task.id === taskId);
	}

	private markBlockedTasks(): void {
		for (const task of this.state.tasks) {
			if (task.status !== "queued") continue;
			const dependency = task.dependsOn.find((id) => {
				const dependencyTask = this.taskById(id);
				return !dependencyTask || ["failed", "blocked", "cancelled", "timed_out"].includes(dependencyTask.status);
			});
			if (dependency) {
				task.status = "blocked";
				task.error = `依赖任务 ${dependency} 未完成`;
				this.emit("state", `任务 ${task.id} 被阻塞：${task.error}`, task.id, true);
			}
		}
	}

	private failedWorkerReview(worker: WorkerResult, canRetry: boolean): ReviewResult {
		const reason = worker.errorMessage || worker.stderr.slice(-4_000) || `worker 退出码为 ${worker.exitCode}`;
		return {
			decision: worker.exitCode === 124 ? "timeout" : canRetry ? "retry" : "escalate",
			reason: worker.exitCode === 124 ? `worker 超时：${reason}` : reason,
			missingCriteria: [],
			nextInstruction: "请检查 worker 进程错误并重新完成任务。",
		};
	}

	private workerSystemPrompt(level: ClusterLevel): string {
		return [
			"你是 Pi 子 agent 集群中的执行 worker。",
			`当前等级：${level}。`,
			"只处理分配给你的任务，不要重新拆分任务，也不要自行升级模型。",
			"先检查现有代码和上下文，再实施必要修改。",
			"完成后用中文简洁汇报：做了什么、修改了哪些文件、运行了哪些测试、仍有哪些风险。",
		].join("\n");
	}

	private workerPrompt(task: TaskRuntime, feedback: string): string {
		return [
			`总目标：${this.state.goal}`,
			`任务 ID：${task.id}`,
			`任务标题：${task.title}`,
			`任务类型：${task.taskType}`,
			`任务要求：${task.task}`,
			"验收标准：",
			...task.acceptanceCriteria.map((criterion, index) => `${index + 1}. ${criterion}`),
			feedback ? `上一次审核反馈：${feedback}` : "",
			"请直接完成任务，并在结束时输出可供审核器核对的结果摘要。",
		].filter(Boolean).join("\n");
	}

	private buildResult(): ClusterRunResult {
		const usage = clusterUsage(this.state.tasks);
		const lines = this.state.tasks.map((task) => {
			const reason = task.review?.reason || task.error || "";
			return `- ${task.id}: ${task.status}${reason ? `，${reason}` : ""}`;
		});
		return {
			runId: this.state.runId,
			status: this.state.status,
			summary: [`目标：${this.state.goal}`, ...lines].join("\n"),
			tasks: this.state.tasks,
			usage,
		};
	}

	private refresh(): void {
		this.onChange?.(this.state);
	}

	private emit(kind: ClusterEvent["kind"], message: string, taskId?: string, forcePersist = false): void {
		this.state.events.push({ timestamp: Date.now(), taskId, kind, message });
		if (this.state.events.length > 120) this.state.events.splice(0, this.state.events.length - 120);
		this.onChange?.(this.state);
		if (forcePersist) {
			this.onPersist?.(this.state);
			if (this.persistence) void persistSnapshot(this.persistence, this.state);
		}
	}
}