import { join } from "node:path";
import type { AssistantMessage, Message, Model, TextContent, ToolCall, ToolResultMessage } from "@earendil-works/pi-ai";
import {
	AuthStorage,
	type CreateAgentSessionOptions,
	createAgentSession,
	createCodingTools,
	getAgentDir,
	ModelRegistry,
	SessionManager,
	SettingsManager,
} from "@earendil-works/pi-coding-agent";
import { createStructuredOutputTool, type StructuredOutputCapture } from "./structured-output.js";

/** Real token/cost usage for one subagent run, read from the SDK session. */
export interface AgentUsage {
	input: number;
	output: number;
	total: number;
	cost: number;
}

/**
 * One step in the subagent's tool-call chain — an assistant tool_use paired
 * (when available) with its tool_result. Persisted on the journal so
 * pi-lookback can render the full chain under each workflow agent instead of
 * just its final return value.
 */
export interface ToolCallRecord {
	id: string;
	name: string;
	arguments: Record<string, unknown>;
	/** Stringified tool output (truncated to MAX_TOOL_RESULT_BYTES). Absent if no matching result was found. */
	result?: string;
	resultTruncated?: boolean;
	isError?: boolean;
}

/** Per-record cap on stringified tool result. Keeps journals from blowing up on giant Read/Bash outputs. */
const MAX_TOOL_RESULT_BYTES = 32 * 1024;
/** Hard cap on chain length per agent — a 1000-tool runaway shouldn't bloat the journal. */
const MAX_TOOL_CHAIN_ENTRIES = 500;

/** Reasoning-effort levels accepted by the SDK's createAgentSession({ thinkingLevel }). */
export type ThinkingEffort = "off" | "minimal" | "low" | "medium" | "high" | "xhigh";
export const THINKING_EFFORT_VALUES: readonly ThinkingEffort[] = [
	"off",
	"minimal",
	"low",
	"medium",
	"high",
	"xhigh",
];

export interface AgentRunOptions {
	label?: string;
	/** A plain JSON Schema; when present the agent must emit structured_output. */
	schema?: object;
	instructions?: string;
	signal?: AbortSignal;
	/** `provider/modelId` or bare `modelId`; session default when omitted/unresolved. */
	model?: string;
	/** Reasoning effort / thinking level for this run; SDK default when omitted. */
	effort?: ThinkingEffort;
	/** Run in a different cwd (e.g. an isolated worktree). */
	cwd?: string;
	onUsage?: (u: AgentUsage) => void;
	/** Fires once at the end of the run (success or error) with the full tool-call chain. */
	onToolChain?: (chain: ToolCallRecord[]) => void;
	onModelFallback?: (spec: string) => void;
}

export interface WorkflowAgentOptions {
	cwd?: string;
	/** The session's main model spec (`provider/modelId`) — the inherited default. */
	mainModel?: string;
	session?: Partial<CreateAgentSessionOptions>;
}

/**
 * Spawns real Pi subagent sessions. One instance per workflow run; .run() is
 * called once per agent() call in the script. Faithful to CC's subagent: the
 * final text IS the return value, or — with a schema — the validated
 * structured_output args, with bounded repair before failing.
 */
export class WorkflowAgent {
	private readonly cwd: string;
	private readonly mainModel?: string;
	private readonly baseSession: Partial<CreateAgentSessionOptions>;
	private registry?: ModelRegistry;

	constructor(options: WorkflowAgentOptions = {}) {
		this.cwd = options.cwd ?? process.cwd();
		this.mainModel = options.mainModel;
		this.baseSession = options.session ?? {};
	}

	private getRegistry(): ModelRegistry {
		if (!this.registry) {
			const dir = getAgentDir();
			const auth = AuthStorage.create(join(dir, "auth.json"));
			this.registry = ModelRegistry.create(auth, join(dir, "models.json"));
		}
		return this.registry;
	}

	private resolveModel(spec: string): Model<unknown> | undefined {
		const registry = this.getRegistry();
		const slash = spec.indexOf("/");
		if (slash > 0) return registry.find(spec.slice(0, slash), spec.slice(slash + 1)) as Model<unknown> | undefined;
		return (registry.getAvailable().find((m) => m.id === spec) ??
			registry.getAll().find((m) => m.id === spec)) as Model<unknown> | undefined;
	}

	async run(prompt: string, options: AgentRunOptions = {}): Promise<unknown> {
		const capture: StructuredOutputCapture = { called: false, value: undefined };
		const runCwd = options.cwd ?? this.cwd;

		const tools = createCodingTools(runCwd);
		if (options.schema) {
			tools.push(createStructuredOutputTool({ schema: options.schema, capture }));
		}

		// Resolve model: explicit spec > session default. An unresolved spec falls
		// back to the default (with a callback) rather than failing the run.
		const spec = options.model ?? this.mainModel;
		let resolvedModel: Model<unknown> | undefined;
		if (options.model) {
			resolvedModel = this.resolveModel(options.model);
			if (!resolvedModel) options.onModelFallback?.(options.model);
		} else if (spec) {
			resolvedModel = this.resolveModel(spec);
		}

		const agentDir = getAgentDir();
		const { session } = await createAgentSession({
			cwd: runCwd,
			agentDir,
			sessionManager: SessionManager.inMemory(),
			// Real SettingsManager so the subagent inherits the user's default
			// provider/model (inMemory would fall back to an unauthed model).
			settingsManager: SettingsManager.create(this.cwd, agentDir),
			customTools: tools,
			...this.baseSession,
			...(resolvedModel ? { model: resolvedModel } : {}),
			...(options.effort ? { thinkingLevel: options.effort } : {}),
		});

		let removeAbort: (() => void) | undefined;
		try {
			if (options.signal?.aborted) throw new Error("Subagent was aborted");
			if (options.signal) {
				const onAbort = () => void session.abort();
				options.signal.addEventListener("abort", onAbort, { once: true });
				removeAbort = () => options.signal?.removeEventListener("abort", onAbort);
			}

			await session.prompt(this.buildPrompt(prompt, options, Boolean(options.schema)));
			if (options.signal?.aborted) throw new Error("Subagent was aborted");

			if (options.schema) {
				return await this.resolveStructured(session, capture, options);
			}
			return this.lastAssistantText(session.messages);
		} finally {
			removeAbort?.();
			if (options.onUsage) {
				try {
					const { tokens, cost } = session.getSessionStats();
					options.onUsage({ input: tokens.input, output: tokens.output, total: tokens.total, cost });
				} catch {
					// usage is best-effort
				}
			}
			if (options.onToolChain) {
				try {
					options.onToolChain(extractToolChain(session.messages as Message[]));
				} catch {
					// chain capture is best-effort; never fail the run because we couldn't summarize it
				}
			}
			session.dispose();
		}
	}

	private async resolveStructured(
		session: { prompt(t: string): Promise<void>; setActiveToolsByName?(n: string[]): void; messages: unknown[] },
		capture: StructuredOutputCapture,
		options: AgentRunOptions,
	): Promise<unknown> {
		if (capture.called) return capture.value;
		try {
			session.setActiveToolsByName?.(["structured_output"]);
		} catch {
			// re-prompt alone still drives most models to comply
		}
		for (let attempt = 0; attempt < 2 && !capture.called; attempt++) {
			if (options.signal?.aborted) throw new Error("Subagent was aborted");
			await session.prompt(
				"You did not call the structured_output tool. Call structured_output now as your only action, with every required field filled in. Do not write a prose answer.",
			);
		}
		if (capture.called) return capture.value;
		throw new Error(`agent "${options.label ?? "?"}" did not produce valid structured_output after repair attempts`);
	}

	private buildPrompt(prompt: string, options: AgentRunOptions, structured: boolean): string {
		const parts = [
			options.instructions,
			options.label ? `Task label: ${options.label}` : undefined,
			prompt,
			"Your final text IS the return value of this subagent — return raw data/conclusions, not a human-facing message.",
		].filter(Boolean) as string[];
		if (structured) {
			parts.push(
				[
					"Final output contract:",
					"- Your final action MUST be a single structured_output tool call.",
					"- Its arguments are the return value of this subagent.",
					"- Inspect files / run commands first if needed, then call structured_output exactly once.",
				].join("\n"),
			);
		}
		return parts.join("\n\n");
	}

	private lastAssistantText(messages: unknown[]): string {
		for (let i = messages.length - 1; i >= 0; i--) {
			const message = messages[i] as Partial<AssistantMessage> | undefined;
			if (message?.role !== "assistant" || !Array.isArray(message.content)) continue;
			const text = message.content
				.filter((p): p is TextContent => (p as { type?: string }).type === "text")
				.map((p) => p.text)
				.join("");
			if (text.trim()) return text;
		}
		return "";
	}
}

/**
 * Walk session.messages in order and emit one ToolCallRecord per assistant
 * tool_use, joined with its matching toolResult (by toolCallId). Best-effort:
 * unmatched calls still appear with no result; orphan results are dropped
 * (they have no caller to attribute to). Bounded by MAX_TOOL_CHAIN_ENTRIES.
 */
export function extractToolChain(messages: Message[]): ToolCallRecord[] {
	if (!Array.isArray(messages)) return [];
	const results = new Map<string, ToolResultMessage>();
	for (const m of messages) {
		if (m && (m as Message).role === "toolResult") {
			const r = m as ToolResultMessage;
			if (r.toolCallId) results.set(r.toolCallId, r);
		}
	}
	const chain: ToolCallRecord[] = [];
	for (const m of messages) {
		if (!m || (m as Message).role !== "assistant") continue;
		const content = (m as AssistantMessage).content;
		if (!Array.isArray(content)) continue;
		for (const block of content) {
			if ((block as { type?: string }).type !== "toolCall") continue;
			const call = block as ToolCall;
			const record: ToolCallRecord = {
				id: call.id,
				name: call.name,
				arguments: call.arguments ?? {},
			};
			const hit = results.get(call.id);
			if (hit) {
				const text = (hit.content || [])
					.filter((p): p is TextContent => (p as { type?: string }).type === "text")
					.map((p) => p.text)
					.join("\n");
				if (text.length > MAX_TOOL_RESULT_BYTES) {
					record.result = text.slice(0, MAX_TOOL_RESULT_BYTES);
					record.resultTruncated = true;
				} else {
					record.result = text;
				}
				if (hit.isError) record.isError = true;
			}
			chain.push(record);
			if (chain.length >= MAX_TOOL_CHAIN_ENTRIES) return chain;
		}
	}
	return chain;
}