// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
// SPDX-FileCopyrightText: Armin Ronacher <armin.ronacher@active-4.com>
//
// SPDX-License-Identifier: Apache-2.0

/**
 * Q&A extraction hook - extracts questions from assistant responses
 *
 * Custom interactive TUI for answering questions.
 *
 * Demonstrates the "prompt generator" pattern with custom TUI:
 * 1. /answer command gets the last assistant message
 * 2. Shows a spinner while extracting questions as structured JSON
 * 3. Presents an interactive TUI to navigate and answer questions
 * 4. Submits the compiled answers when done
 */

import { complete, type UserMessage } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { BorderedLoader } from "@mariozechner/pi-coding-agent";
import { parseToolCallResult, resolveExtractionModel } from "./extract.js";
import { type ExtractionResult, QUESTION_EXTRACTION_TOOL, SYSTEM_PROMPT } from "./prompt.js";
import { QnAComponent } from "./QnAComponent.js";

export default function (pi: ExtensionAPI) {
	const answerHandler = async (ctx: ExtensionContext, instruction?: string) => {
		if (!ctx.hasUI) {
			ctx.ui.notify("answer requires interactive mode", "error");
			return;
		}

		if (!ctx.model) {
			ctx.ui.notify("No model selected", "error");
			return;
		}

		// Find the last assistant message on the current branch
		const branch = ctx.sessionManager.getBranch();
		let lastAssistantText: string | undefined;

		for (let i = branch.length - 1; i >= 0; i--) {
			const entry = branch[i];
			if (entry.type === "message") {
				const msg = entry.message;
				if ("role" in msg && msg.role === "assistant") {
					if (msg.stopReason !== "stop") {
						ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error");
						return;
					}
					const textParts = msg.content
						.filter((c): c is { type: "text"; text: string } => c.type === "text")
						.map((c) => c.text);
					if (textParts.length > 0) {
						lastAssistantText = textParts.join("\n");
						break;
					}
				}
			}
		}

		if (!lastAssistantText) {
			ctx.ui.notify("No assistant messages found", "error");
			return;
		}

		// Capture narrowed values so closures can use them without non-null assertions.
		const sessionModel = ctx.model;
		const assistantText = lastAssistantText;

		// Run extraction with loader UI.
		// The result distinguishes user cancellation from errors so we can
		// show the right message after the custom UI closes.
		type ExtractionOutcome =
			| { kind: "ok"; result: ExtractionResult }
			| { kind: "cancelled" }
			| { kind: "error"; message: string };

		const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
			const extractionModel = resolveExtractionModel(ctx, sessionModel);
			const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);

			// Guard against double-completion: loader.onAbort fires on user
			// cancel, but the in-flight promise may also resolve/reject after
			// the abort.  Only the first call to finish() takes effect.
			let finished = false;
			const finish = (result: ExtractionOutcome) => {
				if (finished) return;
				finished = true;
				done(result);
			};

			loader.onAbort = () => finish({ kind: "cancelled" });

			const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
				if (!auth.ok) {
					return { kind: "model_error" as const, model };
				}
				const instructionBlock = instruction
					? `<user_instruction>\n${instruction}\n</user_instruction>\n\nExtract questions from the assistant's message based on the user's instruction above.`
					: "Extract questions from the assistant's message for the user to fill out.";

				const userMessage: UserMessage = {
					role: "user",
					content: [
						{
							type: "text",
							text: `<last_assistant_message>\n${assistantText}\n</last_assistant_message>\n\n${instructionBlock}`,
						},
					],
					timestamp: Date.now(),
				};

				const response = await complete(
					model,
					{
						systemPrompt: SYSTEM_PROMPT,
						messages: [userMessage],
						tools: [QUESTION_EXTRACTION_TOOL],
					},
					{ apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal },
				);

				if (response.stopReason === "aborted") {
					return { kind: "cancelled" as const };
				}

				if (response.stopReason === "error" || response.content.length === 0) {
					return { kind: "model_error" as const, model };
				}

				const parsed = parseToolCallResult(response);
				if (!parsed) {
					return {
						kind: "parse_error" as const,
						message: `${model.name} did not call extract_questions tool`,
					};
				}
				return { kind: "ok" as const, result: parsed };
			};

			const doExtract = async () => {
				let result = await tryExtract(extractionModel);

				// If the preferred model errored and it's not already the
				// session model, fall back to the session model.
				if (result.kind === "model_error" && extractionModel.id !== sessionModel.id) {
					ctx.ui.notify(`${extractionModel.name} unavailable, falling back to ${sessionModel.name}...`, "warning");
					result = await tryExtract(sessionModel);
				}

				switch (result.kind) {
					case "ok":
						return finish({ kind: "ok", result: result.result });
					case "cancelled":
						return finish({ kind: "cancelled" });
					case "model_error":
						return finish({
							kind: "error",
							message: `${result.model.name} returned an error with no content`,
						});
					case "parse_error":
						return finish({ kind: "error", message: result.message });
				}
			};

			doExtract().catch((err) =>
				finish({
					kind: "error",
					message: `${err?.message ?? err}`,
				}),
			);

			return loader;
		});

		if (outcome.kind === "cancelled") {
			ctx.ui.notify("Cancelled", "info");
			return;
		}
		if (outcome.kind === "error") {
			ctx.ui.notify(`Extraction failed: ${outcome.message}`, "error");
			return;
		}

		const extractionResult = outcome.result;

		if (extractionResult.questions.length === 0) {
			ctx.ui.notify("No questions found in the last message", "info");
			return;
		}

		// Show the Q&A component
		const answersResult = await ctx.ui.custom<string | null>((tui, _theme, _kb, done) => {
			return new QnAComponent(extractionResult.questions, tui, done);
		});

		if (answersResult === null) {
			ctx.ui.notify("Cancelled", "info");
			return;
		}

		// Send the answers directly as a message and trigger a turn
		pi.sendMessage(
			{
				customType: "answers",
				content: `I answered your questions in the following way:\n\n${answersResult}`,
				display: true,
			},
			{ triggerTurn: true },
		);
	};

	pi.registerCommand("answer", {
		description:
			"Extract questions from last assistant message into interactive Q&A. Optional: provide instructions for what to extract.",
		handler: (args, ctx) => answerHandler(ctx, args?.trim() || undefined),
	});

	pi.registerShortcut("ctrl+.", {
		description: "Extract and answer questions",
		handler: answerHandler,
	});
}