/**
 * UOK Execution Graph Scheduler v2
 *
 * Purpose: extend the v1 scheduler with cancellation tokens, progress
 * streaming, and worker pool enforcement for background /tasks work.
 * Maintains backward compatibility with v1's serial/parallel modes.
 *
 * Consumer: /tasks background work, parallel unit dispatch, and UOK kernel
 * when running multiple gates concurrently.
 */

import { debugLog } from "../debug-logger.js";
import { getUnitTimeout } from "../preferences.js";
import {
	persistFullGraph,
	persistGraphNode,
	persistGraphSnapshot,
	persistProgressEvent,
} from "./execution-graph-persist.js";
import { buildTaskRecord, gateOutcomesToTaskState } from "./task-state.js";

const DEFAULT_UNIT_TIMEOUT_MS = 300_000; // 5 minutes fallback
const DEFAULT_SCHEDULER_TIMEOUT_MS = 3_600_000; // 1 hour default for background work

/**
 * Cancellation token for cooperative task abort.
 */
export class CancellationToken {
	#cancelled = false;
	#reason = null;
	#listeners = new Set();

	cancel(reason = "cancelled") {
		if (this.#cancelled) return;
		this.#cancelled = true;
		this.#reason = reason;
		for (const fn of this.#listeners) {
			try {
				fn(reason);
			} catch {
				/* ignore */
			}
		}
		this.#listeners.clear();
	}

	get isCancelled() {
		return this.#cancelled;
	}

	get reason() {
		return this.#reason;
	}

	onCancel(fn) {
		if (this.#cancelled) {
			try {
				fn(this.#reason);
			} catch {
				/* ignore */
			}
			return () => {};
		}
		this.#listeners.add(fn);
		return () => this.#listeners.delete(fn);
	}

	throwIfCancelled() {
		if (this.#cancelled) {
			const err = new Error(this.#reason ?? "cancelled");
			err.name = "CancellationError";
			throw err;
		}
	}
}

/**
 * Progress stream for reporting task state changes back to parent session.
 */
export class ProgressStream {
	#listeners = new Set();
	#history = [];
	#maxHistory = 100;

	emit(event) {
		const enriched = {
			...event,
			ts: new Date().toISOString(),
		};
		this.#history.push(enriched);
		if (this.#history.length > this.#maxHistory) {
			this.#history = this.#history.slice(-this.#maxHistory);
		}
		for (const fn of this.#listeners) {
			try {
				fn(enriched);
			} catch {
				/* ignore */
			}
		}
	}

	onProgress(fn) {
		this.#listeners.add(fn);
		// Replay recent history so new listener catches up
		for (const ev of this.#history) {
			try {
				fn(ev);
			} catch {
				/* ignore */
			}
		}
		return () => this.#listeners.delete(fn);
	}

	getHistory() {
		return [...this.#history];
	}
}

/**
 * Worker pool with limit enforcement and graceful drain.
 */
export class WorkerPool {
	#maxWorkers;
	#active = new Map(); // workerId -> { nodeId, startedAt, token }
	#queue = [];
	#counter = 0;

	constructor(maxWorkers = 2) {
		this.#maxWorkers = Math.max(1, Math.min(32, maxWorkers));
	}

	get maxWorkers() {
		return this.#maxWorkers;
	}

	get activeCount() {
		return this.#active.size;
	}

	get queuedCount() {
		return this.#queue.length;
	}

	get isAtCapacity() {
		return this.#active.size >= this.#maxWorkers;
	}

	async acquire(nodeId, token) {
		return new Promise((resolve, reject) => {
			const tryAcquire = () => {
				if (token?.isCancelled) {
					reject(new Error(token.reason ?? "cancelled"));
					return;
				}
				if (this.#active.size < this.#maxWorkers) {
					this.#counter++;
					const workerId = `w-${this.#counter}`;
					const startedAt = Date.now();
					this.#active.set(workerId, { nodeId, startedAt, token });
					resolve({
						workerId,
						release: () => this.#release(workerId),
					});
					return;
				}
				this.#queue.push({ nodeId, token, resolve, reject });
			};

			tryAcquire();
		});
	}

	#release(workerId) {
		this.#active.delete(workerId);
		// Drain queue
		while (this.#queue.length > 0 && this.#active.size < this.#maxWorkers) {
			const next = this.#queue.shift();
			if (next.token?.isCancelled) {
				next.reject(new Error(next.token.reason ?? "cancelled"));
				continue;
			}
			this.#counter++;
			const newWorkerId = `w-${this.#counter}`;
			const startedAt = Date.now();
			this.#active.set(newWorkerId, {
				nodeId: next.nodeId,
				startedAt,
				token: next.token,
			});
			next.resolve({
				workerId: newWorkerId,
				release: () => this.#release(newWorkerId),
			});
		}
	}

	cancelAll(reason = "pool-drained") {
		// Reject all queued items first
		for (const item of this.#queue) {
			if (item.token && !item.token.isCancelled) item.token.cancel(reason);
			item.reject(new Error(reason));
		}
		this.#queue = [];
		// Then cancel active workers
		for (const [, { token }] of this.#active) {
			if (token && !token.isCancelled) token.cancel(reason);
		}
	}

	getActive() {
		return Array.from(this.#active.entries()).map(([workerId, info]) => ({
			workerId,
			nodeId: info.nodeId,
			durationMs: Date.now() - info.startedAt,
		}));
	}
}

/**
 * v2 Scheduler with cancellation, progress streaming, and worker pool.
 *
 * Extends v1 semantics:
 * - serial mode: unchanged deterministic topological order
 * - parallel mode: worker pool enforces max concurrency, cancellation
 *   stops new work, progress streams report state changes
 */
export class ExecutionGraphSchedulerV2 {
	handlers = new Map();
	#pool = null;
	#progress = new ProgressStream();
	#tasks = new Map(); // nodeId -> task record
	#db = null;
	#graphId = null;

	registerHandler(kind, handler) {
		this.handlers.set(kind, handler);
	}

	get progress() {
		return this.#progress;
	}

	getTask(nodeId) {
		return this.#tasks.get(nodeId);
	}

	getAllTasks() {
		return Array.from(this.#tasks.values());
	}

	async run(nodes, options = {}) {
		const {
			parallel = false,
			maxWorkers = 2,
			globalToken = new CancellationToken(),
			parentStream = null,
			timeoutMs = DEFAULT_SCHEDULER_TIMEOUT_MS,
			db = null,
			graphId = `uok-graph-${Date.now()}`,
			graphMeta = {},
		} = options;

		this.#tasks.clear();
		this.#db = db;
		this.#graphId = db ? graphId : null;

		// Wire progress stream to parent if provided
		if (parentStream) {
			this.#progress.onProgress((ev) => parentStream.emit(ev));
		}
		if (this.#db && this.#graphId) {
			this.#persistInitialGraph(nodes, graphMeta);
			this.#progress.onProgress((ev) => this.#persistProgress(ev));
		}

		const sorted = this.#topologicalSort(nodes);
		const conflicts = this.#detectFileConflicts(nodes);

		// Global timeout: auto-cancel if scheduler runs too long
		const timeoutToken = new CancellationToken();
		const timeoutHandle = setTimeout(() => {
			timeoutToken.cancel(`scheduler-timeout-${timeoutMs}ms`);
		}, timeoutMs);

		// Link timeout to global token
		timeoutToken.onCancel((reason) => {
			if (globalToken && !globalToken.isCancelled) {
				globalToken.cancel(reason);
			}
		});

		try {
			if (!parallel) {
				const result = await this.#runSerial(sorted, globalToken, conflicts);
				this.#persistGraphStatus("completed", sorted.length);
				return result;
			}
			this.#pool = new WorkerPool(maxWorkers);
			const result = await this.#runParallel(
				sorted,
				globalToken,
				conflicts,
				maxWorkers,
			);
			this.#persistGraphStatus("completed", sorted.length);
			return result;
		} finally {
			clearTimeout(timeoutHandle);
		}
	}

	async #runSerial(sorted, token, conflicts) {
		const order = [];
		for (const node of sorted) {
			token.throwIfCancelled();
			const result = await this.#executeNode(node, token, null);
			order.push(result.nodeId);
			this.#progress.emit({
				type: "node-complete",
				nodeId: result.nodeId,
				state: result.state,
				durationMs: result.durationMs,
			});
		}
		return { order, conflicts, tasks: this.getAllTasks() };
	}

	async #runParallel(sorted, token, conflicts, maxWorkers) {
		const done = new Set();
		const order = [];
		const promises = new Map(); // nodeId -> promise

		this.#progress.emit({
			type: "scheduler-start",
			nodeCount: sorted.length,
			maxWorkers,
		});

		while (done.size < sorted.length) {
			token.throwIfCancelled();

			const ready = sorted.filter(
				(n) =>
					!done.has(n.id) &&
					!promises.has(n.id) &&
					n.dependsOn.every((dep) => done.has(dep)),
			);

			if (ready.length === 0 && promises.size === 0) {
				throw new Error(
					"Execution graph deadlock detected: no ready nodes and graph not complete",
				);
			}

			// Start as many ready nodes as pool allows
			for (const node of ready) {
				if (this.#pool.isAtCapacity) break;
				if (promises.has(node.id)) continue;

				const nodeToken = new CancellationToken();
				// Link to global token
				token.onCancel((reason) => nodeToken.cancel(reason));

				const acquirePromise = this.#pool
					.acquire(node.id, nodeToken)
					.then(async ({ workerId, release }) => {
						try {
							const result = await this.#executeNode(node, nodeToken, workerId);
							done.add(node.id);
							order.push(result.nodeId);
							promises.delete(node.id);
							this.#progress.emit({
								type: "node-complete",
								nodeId: result.nodeId,
								state: result.state,
								durationMs: result.durationMs,
								workerId,
							});
							return result;
						} finally {
							release();
						}
					})
					.catch((err) => {
						if (err.name === "CancellationError") {
							this.#progress.emit({
								type: "node-cancelled",
								nodeId: node.id,
								reason: err.message,
							});
						} else {
							this.#progress.emit({
								type: "node-error",
								nodeId: node.id,
								error: err.message,
							});
						}
						promises.delete(node.id);
						throw err;
					});

				promises.set(node.id, acquirePromise);
			}

			// Wait for at least one node to complete before checking ready again
			if (promises.size > 0) {
				await Promise.race(promises.values());
			}
		}

		this.#progress.emit({
			type: "scheduler-complete",
			nodeCount: sorted.length,
			order,
		});

		return { order, conflicts, tasks: this.getAllTasks() };
	}

	async #executeNode(node, token, workerId) {
		const startedAt = Date.now();
		this.#progress.emit({
			type: "node-start",
			nodeId: node.id,
			workerId,
		});

		const handler = this.handlers.get(node.kind);
		let gateResults = [];
		let error = null;

		// Per-node timeout: use unit timeout from preferences or fallback
		const unitTimeoutMs = getUnitTimeout() * 1000 || DEFAULT_UNIT_TIMEOUT_MS;
		const nodeToken = new CancellationToken();
		const nodeTimeoutHandle = setTimeout(() => {
			nodeToken.cancel(`unit-timeout-${unitTimeoutMs}ms`);
		}, unitTimeoutMs);

		// Link node timeout to parent token
		token.onCancel((reason) => {
			if (!nodeToken.isCancelled) nodeToken.cancel(reason);
		});

		try {
			// Check both tokens
			token.throwIfCancelled();
			nodeToken.throwIfCancelled();
			if (handler) {
				const result = await handler(node, nodeToken, this.#progress);
				if (result && Array.isArray(result.gateResults)) {
					gateResults = result.gateResults;
				}
			}
		} catch (err) {
			error = err;
			if (err.name === "CancellationError") {
				gateResults = [
					{
						gateId: "scheduler",
						outcome: "fail",
						failureClass: "cancelled",
						rationale: err.message,
					},
				];
			} else {
				gateResults = [
					{
						gateId: "scheduler",
						outcome: "fail",
						failureClass: "execution",
						rationale: err.message,
					},
				];
			}
		} finally {
			clearTimeout(nodeTimeoutHandle);
		}

		const durationMs = Date.now() - startedAt;
		const state = gateOutcomesToTaskState(gateResults);

		const task = buildTaskRecord({
			nodeId: node.id,
			unitType: node.metadata?.unitType ?? node.kind,
			unitId: node.metadata?.unitId ?? node.id,
			milestoneId: node.metadata?.milestoneId,
			sliceId: node.metadata?.sliceId,
			taskId: node.metadata?.taskId,
			title: node.metadata?.title,
			gateResults,
			startedAt: new Date(startedAt).toISOString(),
			endedAt: new Date().toISOString(),
			durationMs,
			workerId,
		});

		this.#tasks.set(node.id, task);
		this.#persistTask(node, task, gateResults, error);

		if (error && error.name !== "CancellationError") {
			this.#persistGraphStatus("failed", this.#tasks.size);
			throw error;
		}

		return { nodeId: node.id, state, durationMs, gateResults };
	}

	#topologicalSort(nodes) {
		const nodeMap = new Map(nodes.map((n) => [n.id, n]));
		const inDegree = new Map(nodes.map((n) => [n.id, 0]));
		for (const node of nodes) {
			for (const dep of node.dependsOn) {
				if (nodeMap.has(dep)) {
					inDegree.set(node.id, (inDegree.get(node.id) ?? 0) + 1);
				}
			}
		}
		const queue = nodes
			.filter((n) => (inDegree.get(n.id) ?? 0) === 0)
			.sort((a, b) => a.id.localeCompare(b.id));
		const ordered = [];
		while (queue.length > 0) {
			const current = queue.shift();
			ordered.push(current);
			for (const next of nodes) {
				if (!next.dependsOn.includes(current.id)) continue;
				const deg = (inDegree.get(next.id) ?? 0) - 1;
				inDegree.set(next.id, deg);
				if (deg === 0) {
					queue.push(next);
					queue.sort((a, b) => a.id.localeCompare(b.id));
				}
			}
		}
		if (ordered.length !== nodes.length) {
			throw new Error("Execution graph has cyclic dependencies");
		}
		return ordered;
	}

	#detectFileConflicts(nodes) {
		const conflicts = [];
		for (let i = 0; i < nodes.length; i++) {
			const a = nodes[i];
			const writesA = new Set(a.writes ?? []);
			if (writesA.size === 0) continue;
			for (let j = i + 1; j < nodes.length; j++) {
				const b = nodes[j];
				for (const file of b.writes ?? []) {
					if (writesA.has(file)) {
						conflicts.push({ nodeA: a.id, nodeB: b.id, file });
					}
				}
			}
		}
		return conflicts;
	}

	#persistInitialGraph(nodes, graphMeta) {
		try {
			persistFullGraph(this.#db, this.#graphId, nodes, {
				...graphMeta,
				status: "active",
			});
		} catch (err) {
			debugLog(
				"scheduler-v2",
				`failed to persist initial graph: ${err?.message ?? String(err)}`,
			);
		}
	}

	#persistProgress(event) {
		if (!this.#db || !this.#graphId) return;
		try {
			persistProgressEvent(this.#db, this.#graphId, event);
		} catch (err) {
			debugLog(
				"scheduler-v2",
				`failed to persist progress event: ${err?.message ?? String(err)}`,
			);
		}
	}

	#persistTask(node, task, gateResults, error) {
		if (!this.#db || !this.#graphId) return;
		try {
			persistGraphNode(this.#db, this.#graphId, {
				id: task.id,
				kind: node.kind,
				unitType: task.unitType,
				unitId: task.unitId,
				milestoneId: task.milestoneId,
				sliceId: task.sliceId,
				taskId: task.taskId,
				title: task.title,
				dependsOn: node.dependsOn ?? [],
				writes: node.writes ?? [],
				state: task.state,
				workerId: task.workerId,
				startedAt: task.startedAt,
				endedAt: task.endedAt,
				durationMs: task.durationMs,
				costUsd: task.costUsd,
				modelId: task.modelId,
				gateResults,
				error: error ? (error.message ?? String(error)) : null,
			});
		} catch (err) {
			debugLog(
				"scheduler-v2",
				`failed to persist task ${task.id}: ${err?.message ?? String(err)}`,
			);
		}
	}

	#persistGraphStatus(status, nodeCount) {
		if (!this.#db || !this.#graphId) return;
		try {
			persistGraphSnapshot(this.#db, {
				id: this.#graphId,
				nodeCount,
				createdAt: new Date().toISOString(),
				status,
			});
		} catch (err) {
			debugLog(
				"scheduler-v2",
				`failed to persist graph status: ${err?.message ?? String(err)}`,
			);
		}
	}
}
