import * as fs from "node:fs";
import * as os from "node:os";
import * as vm from "node:vm";
import type { AgentRunOptions, AgentUsage, ThinkingEffort, ToolCallRecord, WorkflowAgent } from "./agent.js";
import { createLimiter, type Limiter } from "./limiter.js";
import { hashString, type JournalEntry } from "./persistence.js";
import { type WorkflowMeta, parseWorkflowScript } from "./parse.js";

export const MAX_CONCURRENCY = 16;
export const MAX_AGENTS_PER_RUN = 1000;
export const MAX_FANOUT_ITEMS = 4096;
export const DEFAULT_AGENT_TIMEOUT_MS = 300_000;

/** Options the script passes to agent(). Mirrors CC's documented opts. */
export interface AgentOptions {
	label?: string;
	phase?: string;
	schema?: object;
	model?: string;
	effort?: ThinkingEffort;
	isolation?: "worktree";
	agentType?: string;
	timeoutMs?: number;
}

/** Shared across a run and any nested workflow() so the global caps/budget hold across nesting. */
export interface SharedRuntime {
	limiter: Limiter;
	agentCount: number;
	spent: number;
	depth: number;
}

export interface WorkflowRunResult<T = unknown> {
	runId: string;
	meta: WorkflowMeta;
	result: T;
	agentCount: number;
	tokensTotal: number;
	cost: number;
	logs: string[];
	phases: string[];
	durationMs: number;
	journal: JournalEntry[];
}

export interface WorkflowRunOptions {
	runId: string;
	cwd?: string;
	args?: unknown;
	signal?: AbortSignal;
	concurrency?: number;
	maxAgents?: number;
	agentTimeoutMs?: number;
	tokenBudget?: number | null;
	mainModel?: string;
	/** Injectable for tests; defaults to a real subagent runner. */
	agent?: WorkflowAgent;
	/** Pre-existing journal (resume): callIndex -> {hash, result}. */
	resumeJournal?: Map<number, { hash: string; result: unknown }>;
	/** Resolve a saved workflow name to its script (for workflow('name')). */
	loadSavedWorkflow?: (name: string) => string | undefined;
	/** Shared caps/budget (set when nested inside another workflow()). */
	sharedRuntime?: SharedRuntime;
	// Live callbacks (used by the manager for snapshots + journal persistence).
	onAgentStart?: (e: { index: number; label: string; phase?: string; model?: string }) => void;
	onAgentEnd?: (e: { index: number; label: string; phase?: string; result: unknown; tokens: number; error?: string; toolCalls?: ToolCallRecord[] }) => void;
	onJournal?: (e: JournalEntry) => void;
	onPhase?: (title: string) => void;
	onLog?: (msg: string) => void;
}

/**
 * The determinism prelude, executed inside the vm realm BEFORE the user body.
 * Neuters the nondeterministic builtins that would break resume — matching CC's
 * contract (Date.now()/Math.random()/argless new Date() throw). Uses the vm
 * realm's own Math/Date so it adds no host-Function escape. Not a security
 * sandbox — a best-effort guard against accidental nondeterminism.
 */
const DETERMINISM_PRELUDE = [
	'"use strict";',
	'Math.random = () => { throw new Error("Math.random() is unavailable in a workflow (it breaks resume); vary by index or pass via args"); };',
	"{",
	"  const RealDate = Date;",
	'  const fail = (w) => { throw new Error(w + " is unavailable in a workflow (it breaks resume); pass a timestamp via args"); };',
	"  const SafeDate = function (...a) {",
	'    if (!new.target) fail("Date()");',
	'    if (a.length === 0) fail("new Date()");',
	"    return Reflect.construct(RealDate, a, SafeDate);",
	"  };",
	"  SafeDate.UTC = RealDate.UTC; SafeDate.parse = RealDate.parse;",
	'  SafeDate.now = () => fail("Date.now()");',
	"  SafeDate.prototype = RealDate.prototype;",
	"  globalThis.Date = SafeDate;",
	"}",
].join("\n");

interface RunState {
	logs: string[];
	phases: string[];
	currentPhase?: string;
	callSeq: number;
	firstMiss: number;
}

export async function runWorkflow<T = unknown>(
	script: string,
	options: WorkflowRunOptions,
): Promise<WorkflowRunResult<T>> {
	const started = Date.now();
	const { meta, body } = parseWorkflowScript(script);
	const maxAgents = options.maxAgents ?? MAX_AGENTS_PER_RUN;
	const agentTimeoutMs = options.agentTimeoutMs ?? DEFAULT_AGENT_TIMEOUT_MS;
	const baseCwd = options.cwd ?? process.cwd();
	const tokenBudget = options.tokenBudget ?? null;

	const state: RunState = {
		logs: [],
		// Default the current phase to the first declared one, so agents created
		// before an explicit phase() call still group under a declared phase.
		phases: meta.phases?.[0]?.title ? [meta.phases[0].title] : [],
		currentPhase: meta.phases?.[0]?.title,
		callSeq: 0,
		firstMiss: Number.POSITIVE_INFINITY,
	};

	// Load the SDK-backed runner lazily, only when no runner was injected — so the
	// pure runtime (and its tests) never pull in the pi SDK unless an agent spawns.
	let agentRunner = options.agent;
	const getRunner = async (): Promise<WorkflowAgent> => {
		if (!agentRunner) {
			const { WorkflowAgent } = await import("./agent.js");
			agentRunner = new WorkflowAgent({ cwd: baseCwd, mainModel: options.mainModel });
		}
		return agentRunner;
	};

	const concurrency = Math.max(
		1,
		Math.min(options.concurrency ?? Math.max(1, cpuCount() - 2), MAX_CONCURRENCY),
	);
	const shared: SharedRuntime = options.sharedRuntime ?? {
		limiter: createLimiter(concurrency),
		agentCount: 0,
		spent: 0,
		depth: 0,
	};
	const journal: JournalEntry[] = [];
	let costTotal = 0;
	let tokensTotal = 0;

	const throwIfAborted = () => {
		if (options.signal?.aborted) throw new Error("workflow aborted");
	};

	const log = (message: string) => {
		const text = String(message);
		state.logs.push(text);
		options.onLog?.(text);
	};

	const phase = (title: string) => {
		state.currentPhase = String(title);
		if (!state.phases.includes(state.currentPhase)) state.phases.push(state.currentPhase);
		options.onPhase?.(state.currentPhase);
	};

	const budget = Object.freeze({
		total: tokenBudget,
		spent: () => shared.spent,
		remaining: () => (tokenBudget == null ? Number.POSITIVE_INFINITY : Math.max(0, tokenBudget - shared.spent)),
	});

	/** Resolve a per-phase model override from meta.phases[].model, else meta.model. */
	const phaseModel = (phaseTitle?: string): string | undefined => {
		if (phaseTitle && meta.phases) {
			const hit = meta.phases.find((p) => p.title === phaseTitle);
			if (hit?.model) return hit.model;
		}
		return meta.model;
	};

	/** Resolve a per-phase effort override from meta.phases[].effort, else meta.effort. */
	const phaseEffort = (phaseTitle?: string): ThinkingEffort | undefined => {
		if (phaseTitle && meta.phases) {
			const hit = meta.phases.find((p) => p.title === phaseTitle);
			if (hit?.effort) return hit.effort as ThinkingEffort;
		}
		return meta.effort as ThinkingEffort | undefined;
	};

	const agent = async (prompt: string, agentOptions: AgentOptions = {}) => {
		throwIfAborted();
		if (typeof prompt !== "string") throw new TypeError("agent(prompt, opts): prompt must be a string");

		if (shared.agentCount >= maxAgents) {
			throw new Error(`agent limit exceeded (${maxAgents}) — a runaway-loop backstop`);
		}
		if (tokenBudget != null && budget.remaining() <= 0) {
			throw new Error("workflow token budget exhausted");
		}

		const assignedPhase = agentOptions.phase ?? state.currentPhase;
		const model = agentOptions.model ?? phaseModel(assignedPhase);
		const effort = agentOptions.effort ?? phaseEffort(assignedPhase);

		// Deterministic resume key: assigned at lexical call time, before the
		// limiter, so parallel()/pipeline() fan-out is reproducible.
		const callIndex = state.callSeq++;
		const callHash = hashString(
			JSON.stringify({
				prompt,
				model: model ?? null,
				effort: effort ?? null,
				phase: assignedPhase ?? null,
				schema: agentOptions.schema ?? null,
				agentType: agentOptions.agentType ?? null,
				isolation: agentOptions.isolation ?? null,
			}),
		);

		// Reserve the slot synchronously (atomic with the gate above) so a
		// parallel() fan-out can't all observe the same count and overshoot.
		shared.agentCount++;
		const label = (agentOptions.label?.trim() || (assignedPhase ? `${assignedPhase} agent ${shared.agentCount}` : `agent ${shared.agentCount}`)).slice(0, 80);

		// Longest-unchanged-prefix resume: replay cached result only while this
		// call's index is before the first changed/new call.
		const cached = options.resumeJournal?.get(callIndex);
		if (cached && cached.hash === callHash && callIndex < state.firstMiss) {
			options.onAgentStart?.({ index: callIndex, label, phase: assignedPhase, model });
			options.onAgentEnd?.({ index: callIndex, label, phase: assignedPhase, result: cached.result, tokens: 0 });
			journal.push({ index: callIndex, hash: callHash, result: cached.result });
			return cached.result;
		}
		if (!cached || cached.hash !== callHash) state.firstMiss = Math.min(state.firstMiss, callIndex);

		return shared.limiter(async () => {
			throwIfAborted();
			options.onAgentStart?.({ index: callIndex, label, phase: assignedPhase, model });

			let worktree: import("./worktree.js").Worktree | undefined;
			if (agentOptions.isolation === "worktree") {
				const { createWorktree } = await import("./worktree.js");
				worktree = await createWorktree(baseCwd, `${options.runId}-${callIndex}-${label}`);
				if (!worktree.isolated) log(`isolation ignored for "${label}" (${worktree.reason})`);
			}
			const runCwd = worktree?.isolated ? worktree.cwd : undefined;

			let usage: AgentUsage | undefined;
			let toolChain: ToolCallRecord[] | undefined;
			const record = (result: unknown): number => {
				const tokens = usage && usage.total > 0 ? usage.total : estimateTokens(result) + estimateTokens(prompt);
				if (usage) costTotal += usage.cost;
				tokensTotal += tokens;
				shared.spent += tokens;
				return tokens;
			};

			try {
				throwIfAborted();
				const runOpts: AgentRunOptions = {
					label,
					schema: agentOptions.schema,
					signal: options.signal,
					model,
					effort,
					cwd: runCwd,
					instructions: assignedPhase ? `You are in workflow phase: ${assignedPhase}.` : undefined,
					onUsage: (u) => {
						usage = u;
					},
					onToolChain: (chain) => {
						toolChain = chain;
					},
					onModelFallback: (spec) => log(`${label}: model "${spec}" unavailable — using session default`),
				};
				const runner = await getRunner();
				const result = await withTimeout(
					runner.run(prompt, runOpts),
					agentOptions.timeoutMs ?? agentTimeoutMs,
					`agent "${label}" timed out`,
				);
				throwIfAborted();
				// A model that resolves but isn't in the user's enabledModels runs an
				// empty, zero-token turn (pi gates it silently). Surface that instead of
				// returning a confusing "" so the author knows to enable the model.
				if (usage && usage.total === 0 && (result === "" || result == null)) {
					log(
						`agent "${label}" returned no output and used 0 tokens — model "${model ?? "(session default)"}" may not be enabled (check settings.enabledModels / /models).`,
					);
				}
				const tokens = record(result);
				const entry: JournalEntry = { index: callIndex, hash: callHash, result };
				if (toolChain?.length) entry.toolCalls = toolChain;
				journal.push(entry);
				options.onJournal?.(entry);
				options.onAgentEnd?.({ index: callIndex, label, phase: assignedPhase, result, tokens, toolCalls: toolChain });
				return result;
			} catch (err) {
				if (options.signal?.aborted) throw err;
				const message = (err as Error).message ?? String(err);
				const tokens = record(null);
				log(`agent "${label}" failed: ${message}`);
				// CC contract: don't cache failures on the journal (so resume retries them).
				// We still surface the captured tool chain live via onAgentEnd so observers
				// (and pi-lookback streaming the snapshot) can show what the agent did before it died.
				options.onAgentEnd?.({ index: callIndex, label, phase: assignedPhase, result: null, tokens, error: message, toolCalls: toolChain });
				// CC contract: a failed/dead agent resolves to null (filter with .filter(Boolean)).
				return null;
			} finally {
				if (worktree?.isolated) {
					const { removeWorktree } = await import("./worktree.js");
					await removeWorktree(worktree, baseCwd);
				}
			}
		});
	};

	const parallel = async (thunks: Array<() => Promise<unknown>>) => {
		throwIfAborted();
		if (!Array.isArray(thunks)) throw new TypeError("parallel() expects an array of functions");
		if (thunks.length > MAX_FANOUT_ITEMS) {
			throw new Error(`parallel() accepts at most ${MAX_FANOUT_ITEMS} items (got ${thunks.length})`);
		}
		if (thunks.some((t) => typeof t !== "function")) {
			throw new TypeError("parallel() expects functions, not promises. Wrap each call: () => agent(...)");
		}
		return Promise.all(
			thunks.map(async (thunk, index) => {
				try {
					return await thunk();
				} catch (err) {
					if (options.signal?.aborted) throw err;
					log(`parallel[${index}] failed: ${(err as Error).message}`);
					return null; // CC: a throwing thunk resolves to null; the call never rejects.
				}
			}),
		);
	};

	const pipeline = async (
		items: unknown[],
		...stages: Array<(prev: unknown, original: unknown, index: number) => unknown>
	) => {
		throwIfAborted();
		if (!Array.isArray(items)) throw new TypeError("pipeline() expects an array as the first argument");
		if (items.length > MAX_FANOUT_ITEMS) {
			throw new Error(`pipeline() accepts at most ${MAX_FANOUT_ITEMS} items (got ${items.length})`);
		}
		if (stages.some((s) => typeof s !== "function")) {
			throw new TypeError("pipeline() stages must be functions");
		}
		// No barrier between stages: each item flows through all stages independently.
		return Promise.all(
			items.map(async (item, index) => {
				let value: unknown = item;
				for (const stage of stages) {
					try {
						throwIfAborted();
						value = await stage(value, item, index);
					} catch (err) {
						if (options.signal?.aborted) throw err;
						log(`pipeline[${index}] dropped at a stage: ${(err as Error).message}`);
						return null; // a throwing stage drops the item and skips its remaining stages
					}
				}
				return value;
			}),
		);
	};

	// One level of nesting only. The child shares caps/budget/limiter/signal.
	const workflow = async (nameOrRef: string | { scriptPath: string }, childArgs?: unknown) => {
		throwIfAborted();
		if (shared.depth >= 1) throw new Error("workflow() nesting is one level only");
		let childScript: string | undefined;
		if (typeof nameOrRef === "string") {
			childScript = options.loadSavedWorkflow?.(nameOrRef);
			if (!childScript) throw new Error(`unknown saved workflow "${nameOrRef}"`);
		} else if (nameOrRef && typeof nameOrRef.scriptPath === "string") {
			try {
				childScript = fs.readFileSync(nameOrRef.scriptPath, "utf8");
			} catch (err) {
				throw new Error(`could not read scriptPath: ${(err as Error).message}`);
			}
		} else {
			throw new TypeError("workflow(nameOrRef): pass a saved-workflow name or { scriptPath }");
		}
		const child = await runWorkflow(childScript, {
			...options,
			args: childArgs,
			resumeJournal: undefined,
			sharedRuntime: { ...shared, depth: shared.depth + 1 },
			runId: `${options.runId}.nested`,
		});
		return child.result;
	};

	// ── Build the sandbox and execute the body in an async context ──
	const sandbox: Record<string, unknown> = {
		agent,
		parallel,
		pipeline,
		phase,
		log,
		workflow,
		args: options.args,
		budget,
		cwd: baseCwd,
		process: { cwd: () => baseCwd },
		console: { log: (...a: unknown[]) => log(a.map(String).join(" ")) },
	};
	const context = vm.createContext(sandbox);
	vm.runInContext(DETERMINISM_PRELUDE, context, { filename: "better-workflows:prelude" });

	let result: T;
	try {
		const wrapped = `(async () => {\n${body}\n})()`;
		const promise = vm.runInContext(wrapped, context, { filename: "better-workflows:script" });
		result = (await promise) as T;
	} catch (err) {
		if (options.signal?.aborted) throw new Error("workflow aborted");
		throw err;
	}

	if (shared.agentCount === 0) {
		throw new Error("workflow scripts must call agent() at least once; this run declared phases but spawned no agents");
	}

	return {
		runId: options.runId,
		meta,
		result,
		agentCount: shared.agentCount,
		tokensTotal,
		cost: costTotal,
		logs: state.logs,
		phases: state.phases,
		durationMs: Date.now() - started,
		journal,
	};
}

function cpuCount(): number {
	try {
		return os.cpus().length || 8;
	} catch {
		return 8;
	}
}

function estimateTokens(value: unknown): number {
	if (value == null) return 0;
	const text = typeof value === "string" ? value : JSON.stringify(value);
	return Math.ceil((text?.length ?? 0) / 4); // ~4 chars/token heuristic when the provider reports none
}

async function withTimeout<T>(promise: Promise<T>, ms: number, message: string): Promise<T> {
	if (!ms || ms <= 0) return promise;
	let timer: ReturnType<typeof setTimeout>;
	const timeout = new Promise<never>((_, reject) => {
		timer = setTimeout(() => reject(new Error(message)), ms);
	});
	try {
		return await Promise.race([promise, timeout]);
	} finally {
		clearTimeout(timer!);
	}
}