/**
 * Extension loaded into sub-agents.
 * - Shows agent identity + available tools as a styled widget above the editor (toggle with Ctrl+J)
 * - Provides a `subagent_done` tool for interactive agents to self-terminate
 */
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Box, Text } from "@earendil-works/pi-tui";
import { Type } from "@sinclair/typebox";
import { writeFileSync } from "node:fs";
import {
	appendPersistentTaskEvent,
	consumePersistentTaskInbox,
	readPersistentDeliveryLedger,
	readPersistentTaskEvents,
} from "./session.ts";
import { createSubagentActivityRecorder } from "./activity.ts";
import { isString } from "./type-guards.ts";

export function shouldMarkUserTookOver(agentStarted: boolean): boolean {
	return agentStarted;
}

export function shouldAutoExitOnAgentEnd(
	_userTookOver: boolean,
	messages: any[] | undefined,
): boolean {
	// Manual input should not strand an auto-exit subagent. If the latest agent
	// turn completed normally, close the session. Escape/abort still leaves it
	// open for inspection or another prompt.
	//
	// stopReason: "error" (e.g. exhausted retries on a provider overload) also
	// returns true — we want to shut down so the parent is woken up — but we
	// pair this with findLatestAssistantError() so the parent learns it was an
	// error, not a clean completion.
	if (messages) {
		for (let i = messages.length - 1; i >= 0; i--) {
			const msg = messages[i];
			if (msg?.role === "assistant") {
				return msg.stopReason !== "aborted";
			}
		}
	}

	return true;
}

export interface SubagentErrorInfo {
	errorMessage: string;
	stopReason: "error";
}

/**
 * If the last assistant message in the turn ended with `stopReason: "error"`
 * (typically auto-retry exhausted on an overload / rate limit / server error),
 * return its error info so the parent orchestrator can surface a clear
 * failure instead of silently treating the run as completed.
 *
 * Returns `null` when the latest assistant turn completed normally or was
 * aborted by the user (handled separately by shouldAutoExitOnAgentEnd).
 */
export function findLatestAssistantError(
	messages: any[] | undefined,
): SubagentErrorInfo | null {
	if (!messages) return null;
	for (let i = messages.length - 1; i >= 0; i--) {
		const msg = messages[i];
		if (msg?.role !== "assistant") continue;
		if (msg.stopReason !== "error") return null;
		const raw = isString(msg.errorMessage) ? msg.errorMessage.trim() : "";
		return {
			errorMessage:
				raw ||
				"Subagent agent loop ended with stopReason=error (no errorMessage field).",
			stopReason: "error",
		};
	}
	return null;
}

export function buildCompletionSidecar(
	messages: any[] | undefined,
):
	| { type: "done" }
	| { type: "error"; errorMessage: string; stopReason: "error" } {
	const errorInfo = findLatestAssistantError(messages);
	return errorInfo ? { type: "error", ...errorInfo } : { type: "done" };
}

export function buildPersistentTaskEvent(task: string, generation: string) {
	return {
		version: 1 as const,
		type: "task-done" as const,
		task,
		generation,
		at: new Date().toISOString(),
	};
}

export function isPersistentStopDirective(
	inbox: ReturnType<typeof consumePersistentTaskInbox>,
): boolean {
	return inbox?.type === "stop";
}

function findUnsettledPersistentTask(
	sessionFile: string,
	generation: string,
): string {
	const settledTasks = new Set(
		readPersistentTaskEvents(sessionFile)
			.filter((event) => event.generation === generation)
			.map((event) => event.task),
	);
	return (
		readPersistentDeliveryLedger(sessionFile)
			.filter(
				(entry) =>
					entry.generation === generation &&
					entry.outcome === "dispatched" &&
					!settledTasks.has(entry.task),
			)
			.at(-1)?.task ?? ""
	);
}

export function parseDeniedTools(rawValue: string | undefined): string[] {
	return (rawValue ?? "")
		.split(",")
		.map((value) => value.trim())
		.filter(Boolean);
}

export default function (pi: ExtensionAPI) {
	let toolNames: string[] = [];
	let denied: string[] = [];
	let expanded = false;

	// Read subagent identity from env vars (set by parent orchestrator)
	const subagentName = process.env.PI_SUBAGENT_NAME ?? "";
	const subagentAgent = process.env.PI_SUBAGENT_AGENT ?? "";
	const deniedToolsValue = process.env.PI_DENY_TOOLS;
	const autoExit = process.env.PI_SUBAGENT_AUTO_EXIT === "1";
	const persistent = process.env.PI_SUBAGENT_PERSISTENT === "1";
	const generation = process.env.PI_SUBAGENT_GENERATION_ID ?? "";
	const sessionFile = process.env.PI_SUBAGENT_SESSION;
	const initialTask = process.env.PI_SUBAGENT_TASK_ID ?? "";
	let currentTask = persistent
		? findUnsettledPersistentTask(sessionFile ?? "", generation) ||
			(initialTask &&
			!readPersistentTaskEvents(sessionFile ?? "").some(
				(event) =>
					event.task === initialTask && event.generation === generation,
			)
				? initialTask
				: "")
		: "";
	const recorder = createSubagentActivityRecorder({
		runningChildId: process.env.PI_SUBAGENT_ID,
		activityFile: process.env.PI_SUBAGENT_ACTIVITY_FILE,
	});

	function renderWidget(ctx: { ui: { setWidget: Function } }, _theme: any) {
		ctx.ui.setWidget(
			"subagent-tools",
			(_tui: any, theme: any) => {
				const box = new Box(1, 0, (text: string) =>
					theme.bg("toolSuccessBg", text),
				);

				const label = subagentAgent || subagentName;
				const agentTag = label
					? theme.bold(theme.fg("accent", `[${label}]`))
					: "";

				if (expanded) {
					// Expanded: full tool list + denied
					const countInfo = theme.fg("dim", ` — ${toolNames.length} available`);
					const hint = theme.fg("muted", "  (Ctrl+J to collapse)");

					const toolList = toolNames
						.map((name: string) => theme.fg("dim", name))
						.join(theme.fg("muted", ", "));

					let deniedLine = "";
					if (denied.length > 0) {
						const deniedList = denied
							.map((name: string) => theme.fg("error", name))
							.join(theme.fg("muted", ", "));
						deniedLine = "\n" + theme.fg("muted", "denied: ") + deniedList;
					}

					const content = new Text(
						`${agentTag}${countInfo}${hint}\n${toolList}${deniedLine}`,
						0,
						0,
					);
					box.addChild(content);
				} else {
					// Collapsed: one-line summary
					const countInfo = theme.fg("dim", ` — ${toolNames.length} tools`);
					const deniedInfo =
						denied.length > 0
							? theme.fg("dim", " · ") +
								theme.fg("error", `${denied.length} denied`)
							: "";
					const hint = theme.fg("muted", "  (Ctrl+J to expand)");

					const content = new Text(
						`${agentTag}${countInfo}${deniedInfo}${hint}`,
						0,
						0,
					);
					box.addChild(content);
				}

				return box;
			},
			{ placement: "aboveEditor" },
		);
	}

	let userTookOver = false;
	let agentStarted = false;
	let latestAgentMessages: any[] | undefined;
	let completionFinalized = false;
	let sessionContext: { shutdown(): void } | undefined;

	// Show widget + status bar on session start
	pi.on("session_start", (_event, ctx) => {
		sessionContext = ctx;
		recorder.sessionStart();
		const tools = pi.getAllTools();
		toolNames = tools.map((t) => t.name).sort();
		denied = parseDeniedTools(deniedToolsValue);

		renderWidget(ctx, null);
	});

	pi.on("input", () => {
		recorder.input();
		// Ignore the initial task message that starts an autonomous subagent.
		// Only inputs after the first agent run has started count as user takeover.
		if (!shouldMarkUserTookOver(agentStarted)) return;
		userTookOver = true;
	});

	pi.on("before_agent_start", () => {
		recorder.beforeAgentStart();
	});

	pi.on("agent_start", () => {
		agentStarted = true;
		recorder.agentStart();
	});

	pi.on("agent_end", (event) => {
		latestAgentMessages = event.messages;
		recorder.agentEndWaiting();
		if (autoExit) {
			// Reset any recorded manual input marker. Auto-exit is decided by whether
			// the latest agent turn completed normally, not by who initiated it.
			userTookOver = false;
		}
	});

	pi.on("agent_settled", (_event, ctx) => {
		if (persistent && !completionFinalized) {
			let messages = latestAgentMessages;
			try {
				const branchMessages = ctx.sessionManager
					.getBranch()
					.flatMap((entry) =>
						entry.type === "message" ? [entry.message] : [],
					);
				if (branchMessages.length > 0) messages = branchMessages;
			} catch {
				// Fall back to the latest low-level run when session evidence is unavailable.
			}
			if (currentTask && shouldAutoExitOnAgentEnd(userTookOver, messages)) {
				appendPersistentTaskEvent(
					process.env.PI_SUBAGENT_SESSION ?? "",
					buildPersistentTaskEvent(currentTask, generation),
				);
				currentTask = "";
				recorder.agentEndWaiting();
			}
			return;
		}
		if (!autoExit || completionFinalized) return;

		let messages = latestAgentMessages;
		try {
			const branchMessages = ctx.sessionManager
				.getBranch()
				.flatMap((entry) => (entry.type === "message" ? [entry.message] : []));
			if (branchMessages.length > 0) messages = branchMessages;
		} catch {
			// Fall back to the latest low-level run when session evidence is unavailable.
		}

		if (!shouldAutoExitOnAgentEnd(userTookOver, messages)) return;
		completionFinalized = true;

		// Surface a settled stopReason: "error" to the parent via the .exit
		// sidecar. Transient errors followed by retry or compaction never reach
		// this point as the latest assistant message.
		const sessionFile = process.env.PI_SUBAGENT_SESSION;
		if (sessionFile) {
			try {
				writeFileSync(
					`${sessionFile}.exit`,
					JSON.stringify(buildCompletionSidecar(messages)),
				);
			} catch {
				// Best effort — the watcher can still detect the terminal sentinel
				// after shutdown if the completion sidecar cannot be written.
			}
		}

		recorder.agentEndDone();
		ctx.shutdown();
	});

	pi.on("turn_start", (event) => {
		recorder.turnStart(event.turnIndex);
	});

	pi.on("turn_end", (event) => {
		recorder.turnEnd(event.turnIndex);
	});

	pi.on("before_provider_request", () => {
		recorder.beforeProviderRequest();
	});

	pi.on("after_provider_response", () => {
		recorder.afterProviderResponse();
	});

	pi.on("message_update", (event) => {
		recorder.messageUpdate(event.assistantMessageEvent?.type);
	});

	pi.on("tool_execution_start", (event) => {
		recorder.toolExecutionStart(event.toolCallId, event.toolName);
	});

	pi.on("tool_call", (event) => {
		recorder.toolCall(event.toolCallId, event.toolName);
	});

	pi.on("tool_execution_update", (event) => {
		recorder.toolExecutionUpdate(event.toolCallId, event.toolName);
	});

	pi.on("tool_result", (event) => {
		recorder.toolResult(event.toolCallId, event.toolName);
	});

	pi.on("tool_execution_end", (event) => {
		recorder.toolExecutionEnd(event.toolCallId, event.toolName);
	});

	pi.on("session_shutdown", (event) => {
		if (inboxPoller) clearInterval(inboxPoller);
		recorder.sessionShutdown(event.reason);
	});

	// Toggle expand/collapse with Ctrl+J
	pi.registerShortcut("ctrl+j", {
		description: "Toggle subagent tools widget",
		handler: (ctx) => {
			expanded = !expanded;
			renderWidget(ctx, null);
		},
	});

	pi.registerTool({
		name: "caller_ping",
		label: "Caller Ping",
		description:
			"Send a help request to the parent agent and exit this session. " +
			"The parent will be notified with your message and can resume this session with a response. " +
			"Use when you're stuck, need clarification, or need the parent to take action.",
		parameters: Type.Object({
			message: Type.String({ description: "What you need help with" }),
		}),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const sessionFile = process.env.PI_SUBAGENT_SESSION;
			if (!sessionFile) {
				throw new Error(
					"caller_ping is only available in subagent contexts. " +
						"PI_SUBAGENT_SESSION environment variable is not set.",
				);
			}

			recorder.callerPing();
			if (persistent) {
				appendPersistentTaskEvent(sessionFile, {
					type: "help-request",
					task: currentTask,
					generation,
					message: params.message,
				});
				return {
					content: [
						{
							type: "text",
							text: "Help request sent. Stay available for subagent_send.",
						},
					],
					details: {},
				};
			}
			const exitData = {
				type: "ping" as const,
				name: process.env.PI_SUBAGENT_NAME ?? "subagent",
				message: params.message,
			};
			writeFileSync(`${sessionFile}.exit`, JSON.stringify(exitData));
			completionFinalized = true;

			ctx.shutdown();
			return {
				content: [
					{
						type: "text",
						text: "Ping sent. Session will exit and parent will be notified.",
					},
				],
				details: {},
			};
		},
	});

	let inboxPoller: ReturnType<typeof setInterval> | undefined;
	if (persistent) {
		inboxPoller = setInterval(() => {
			try {
				const sessionFile = process.env.PI_SUBAGENT_SESSION;
				if (!sessionFile || currentTask) return;
				const inbox = consumePersistentTaskInbox(sessionFile);
				if (!inbox) return;
				if (isPersistentStopDirective(inbox)) {
					try {
						writeFileSync(
							`${sessionFile}.exit`,
							JSON.stringify({ type: "done" }),
						);
					} catch {
						// The parent can still confirm the shell exit marker.
					}
					completionFinalized = true;
					recorder.subagentDone();
					sessionContext?.shutdown();
					return;
				}
				currentTask = inbox.task;
				pi.sendUserMessage(inbox.message);
			} catch {
				// A malformed or transiently unavailable inbox must not kill the poller.
			}
		}, 1000);
	}

	if (autoExit) return;

	pi.registerTool({
		name: "subagent_done",
		label: "Subagent Done",
		description:
			"Call this tool when you have completed your task. " +
			"It will close this session and return your results to the main session. " +
			"Your LAST assistant message before calling this becomes the summary returned to the caller.",
		parameters: Type.Object({}),
		async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
			const sessionFile = process.env.PI_SUBAGENT_SESSION;
			recorder.subagentDone();
			if (persistent && sessionFile && currentTask) {
				appendPersistentTaskEvent(
					sessionFile,
					buildPersistentTaskEvent(currentTask, generation),
				);
				currentTask = "";
				return {
					content: [
						{
							type: "text",
							text: "Task complete. Staying available for subagent_send.",
						},
					],
					details: {},
				};
			}
			if (sessionFile) {
				writeFileSync(`${sessionFile}.exit`, JSON.stringify({ type: "done" }));
			}
			completionFinalized = true;
			ctx.shutdown();
			return {
				content: [{ type: "text", text: "Shutting down subagent session." }],
				details: {},
			};
		},
	});
}