// 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 Tool, type UserMessage } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { BorderedLoader } from "@mariozechner/pi-coding-agent";
import {
	type Component,
	Editor,
	type EditorTheme,
	Key,
	matchesKey,
	truncateToWidth,
	type TUI,
	visibleWidth,
	wrapTextWithAnsi,
} from "@mariozechner/pi-tui";
import { Type } from "@sinclair/typebox";

// Structured output format for question extraction
interface ExtractedQuestion {
	question: string;
	context?: string;
}

interface ExtractionResult {
	questions: ExtractedQuestion[];
}

const SYSTEM_PROMPT = `You extract questions from assistant messages so a user can answer them without reading the full message. Call the extract_questions tool with your results.

Each extracted question must be a standalone object with a "question" string and an optional "context" string. The user will read ONLY the extracted questions, not the original message, so both the question and context must be fully self-contained.

Rules:
- Each question must make sense on its own without referring back to the message (no "Issue 1", "the above", "as mentioned")
- Context should include all relevant details, options, and reasoning needed to answer — it replaces reading the original message
- Merge duplicate or overlapping questions into one (e.g. an analysis section and a plan section about the same topic become one question)
- Capture both explicit questions ("Which approach?") and implicit ones (proposals that need approval, recommendations that need confirmation)
- Keep questions in logical order
- If no questions are found, call the tool with an empty questions array

<examples>
<example>
<description>Simple explicit question</description>
<input>I can set up the database with either PostgreSQL or MySQL. The existing infrastructure uses PostgreSQL for two other services. Which database should I use?</input>
<output>
{"question": "Which database should be used for the new service?", "context": "Options are PostgreSQL or MySQL. The existing infrastructure already uses PostgreSQL for two other services."}
</output>
</example>

<example>
<description>Implicit proposal needing approval</description>
<input>I looked at the failing test. The flakiness comes from a race condition in the cleanup goroutine. I think the simplest fix is to add a sync.WaitGroup so the test waits for cleanup to finish before asserting. Want me to go ahead with that?</input>
<output>
{"question": "Should the flaky test be fixed by adding a sync.WaitGroup to wait for the cleanup goroutine before asserting?", "context": "The test failure is caused by a race condition in the cleanup goroutine. The WaitGroup approach ensures cleanup completes before assertions run."}
</output>
</example>

<example>
<description>Long message with analysis and plan sections that overlap — questions must be merged, deduplicated, and fully standalone</description>
<input>
## Issue 1: Non-deterministic map iteration for overrides → rawArgs

The problem is in RunE at line ~75: for k, vals := range overrides gives random order each run.

Three good approaches:
1. Sort the keys before iterating. Simplest, most common Go pattern.
2. Use []restic.Flag or []struct{key, vals} instead of a map. More invasive but preserves construction order.
3. Skip the map→rawArgs→parsePassthrough round-trip entirely. Instead of serializing overrides to rawArgs and then having runCommand parse them back, just pass the overrides map directly to runCommand and merge there. This eliminates the round-trip and the ordering problem disappears.

Recommend approach 3 as best: it kills the root cause.

## Issue 2: Double-pointers (**screens.Snapshot etc.)

The problem is buildCommandScreens needs to both return screens AND communicate which specific screen instances were created back to runInteractive. Double-pointers are gnarly.

Simpler approach: return a struct holding typed screen pointers. The ResolveFunc closure captures the struct.

## Issue 3: Picker height overflow

The session already subtracts chrome height via adjustedSizeMsg() before forwarding WindowSizeMsg to screens. The legacy pickerModel does its own subtraction because it runs standalone. Assessment: this finding is a false positive.

## Issue 2.1: huh form boilerplate

The buildForm/Init/Update boilerplate is repetitive across Overwrite, Target, Preset, Snapshot. Code is simple and clear in each screen. Recommendation: skip — borderline DRY.

## Issue 2.2: drain test helpers

The drain* test helpers are copy-pasted with only the type changed. Perfect case for Go generics: func drain[S ui.Screen](s S, cmd tea.Cmd) (S, tea.Cmd)

## Issue 3.1: Snapshot theme handler clears user input

When BackgroundColorMsg arrives during phaseSnapshotManual, buildManualForm() resets s.entered = "". Fix: preserve s.entered before rebuilding. Also needs rebuildSelectForm() call during phaseSnapshotSelecting.

## Issue 4: WindowSizeMsg missed during phaseFileLoading

If the screen is loading when the initial WindowSizeMsg arrives, it only goes to the spinner and the size is lost. Fix: store the last WindowSizeMsg on the FilePicker struct, apply when picker is built.

## Issue 5.1: Same as Issue 1

Same root cause — map iteration in RunE. Fix for Issue 1 handles this.

## Issue 6.1: Duplicated constant

Both internal/config and internal/restic define the same commandSuffix constant.

## Issue 6.2: Inconsistent empty-string semantics

rc.Environ treats "" as present (checks map key existence), but os.Getenv checks != "" so an empty env var counts as absent.

## Plan

1. Remove legacy restore prompts (promptRestore, promptSnapshotID, promptFileSelection, printRestoreSummary, standalone pickerModel)
2. Eliminate map→rawArgs round-trip (Issues 1 and 5.1)
3. Replace double-pointers with return struct (Issue 2)
4. Extract generic drain test helper (Issue 2.2)
5. Fix snapshot theme handler (Issue 3.1)
6. Store WindowSizeMsg during loading (Issue 4)
7. Share commandSuffix constant (Issue 6.1)
8. Fix inconsistent empty-string check (Issue 6.2)

NOT planning to do Issue 2.1 (huh form boilerplate extraction).
</input>
<output>
{"question": "How should the non-deterministic map iteration in RunE be fixed, where 'for k, vals := range overrides' produces random argument order each run?", "context": "Three approaches: (1) Sort map keys before iterating — simplest, small code change. (2) Replace map[string][]string with a slice of key-value pairs to preserve construction order. (3) Eliminate the map→rawArgs→parsePassthrough round-trip entirely by passing overrides directly to runCommand — this also fixes the identical issue in rawArgs parsing. Recommended: approach 3."}
{"question": "Should the double-pointer output parameters (**screens.Snapshot, etc.) in buildCommandScreens be replaced with a returned struct?", "context": "Currently uses double-pointers so runInteractive can read results from specific screen instances. A commandScreens struct holding typed screen pointers would be returned instead, and the ResolveFunc closure captures it. Idiomatic Go — return a struct instead of output parameters."}
{"question": "The picker height overflow appears to be a false positive because the session already subtracts chrome height via adjustedSizeMsg() before forwarding WindowSizeMsg to screens. Safe to skip?", "context": "session.go subtracts 5 chrome lines (breadcrumb + title + help bar) before forwarding. The legacy pickerModel does its own subtraction because it runs standalone with tea.NewProgram. The screen adapter doesn't need additional subtraction."}
{"question": "Should the repeated huh form boilerplate (buildForm/Init/Update) across Overwrite, Target, Preset, and Snapshot screens be extracted into a shared helper?", "context": "Each screen repeats same form setup pattern with slight behavior variations. Code is currently simple and self-contained in each screen. Recommended: skip — borderline DRY."}
{"question": "Should the copy-pasted type-specific drain test helpers be replaced with a single generic drain[S ui.Screen] function?", "context": "Multiple test files have identical drain helpers differing only in type parameter (e.g. drainSnapshot, drainFilePicker). A single generic function would replace all of them."}
{"question": "Should the snapshot theme handler be fixed to preserve user input when the terminal theme changes?", "context": "When BackgroundColorMsg arrives during phaseSnapshotManual, buildManualForm() resets s.entered to empty, wiping user input. Fix: preserve s.entered before rebuild. Also needs rebuildSelectForm() call during phaseSnapshotSelecting."}
{"question": "Should FilePicker store the last WindowSizeMsg to prevent size loss during the loading phase?", "context": "If the screen is loading when the initial WindowSizeMsg arrives, the message only goes to the spinner and size is lost. When buildPicker runs later, picker doesn't know terminal size. Fix: add lastSize field, apply when picker built."}
{"question": "Should the duplicated commandSuffix constant be extracted to a shared package?", "context": "Both internal/config and internal/restic define the same constant."}
{"question": "Should os.Getenv() != '' checks be replaced with os.LookupEnv so empty environment variables count as present?", "context": "rc.Environ treats empty string as present (checks map key existence), but os.Getenv checks != '' so an empty env var counts as absent. Using os.LookupEnv would make behavior consistent — explicitly setting RESTIC_REPOSITORY='' would count as configured."}
{"question": "Should the legacy interactive restore prompts (promptRestore, promptSnapshotID, promptFileSelection, printRestoreSummary, standalone pickerModel) be removed from root.go?", "context": "The TUI session now handles the full restore flow through screens, making these standalone prompt functions unused."}
</output>
</example>
</examples>`;

/**
 * Tool definition for structured question extraction. Models produce
 * structured tool-call arguments far more reliably than free-form JSON
 * in a text response.
 */
const QUESTION_EXTRACTION_TOOL: Tool = {
	name: "extract_questions",
	description:
		"Extract questions from the assistant's message. Each question is a self-contained object the user will read instead of the original message.",
	parameters: Type.Object({
		questions: Type.Array(
			Type.Object({
				question: Type.String({
					description: "A complete, standalone question. Must make sense without reading the original message.",
				}),
				context: Type.Optional(
					Type.String({
						description:
							"Self-contained context with all details, options, and reasoning needed to answer the question.",
					}),
				),
			}),
			{
				description:
					"Array of question objects. Merge overlapping questions. Each item MUST be an object with 'question' and optional 'context' strings, NOT a plain string.",
			},
		),
	}),
};

// Preferred model for extraction — lightweight is fine for structured JSON output.
// Falls back to the session model if this one isn't available.
const PREFERRED_EXTRACTION_MODEL_ID = "nemotron-3-super-120b-a12b";

/**
 * Resolve the extraction model: prefer a lightweight model from the registry,
 * fall back to the current session model.
 */
function resolveExtractionModel(ctx: ExtensionContext, fallback: NonNullable<ExtensionContext["model"]>) {
	const available = ctx.modelRegistry.getAvailable();
	const preferred = available.find((m) => m.id === PREFERRED_EXTRACTION_MODEL_ID);
	return preferred ?? fallback;
}

/**
 * Parse a tool call response into an ExtractionResult.
 */
function parseToolCallResult(response: Awaited<ReturnType<typeof complete>>): ExtractionResult | null {
	const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_questions");

	if (!toolCall || toolCall.type !== "toolCall") {
		console.error("Model did not call extract_questions:", response.content);
		return null;
	}

	const args = toolCall.arguments as Record<string, unknown>;
	if (!Array.isArray(args.questions)) {
		console.error("[answer] expected questions array, got:", typeof args.questions, args);
		return null;
	}

	// Validate each question item — model may return plain strings instead
	// of objects if it ignores the schema, so we handle both gracefully.
	const questions: ExtractedQuestion[] = [];
	for (const item of args.questions) {
		if (typeof item === "string") {
			// Model returned bare string instead of object — use as question text
			questions.push({ question: item });
			continue;
		}
		if (typeof item !== "object" || item === null) continue;
		const obj = item as Record<string, unknown>;
		if (typeof obj.question !== "string") {
			console.error("[answer] skipping item with no 'question' string:", Object.keys(obj));
			continue;
		}
		questions.push({
			question: obj.question,
			context: typeof obj.context === "string" ? obj.context : undefined,
		});
	}

	return { questions };
}

/**
 * Interactive Q&A component for answering extracted questions
 */
class QnAComponent implements Component {
	private questions: ExtractedQuestion[];
	private answers: string[];
	private currentIndex: number = 0;
	private editor: Editor;
	private tui: TUI;
	private onDone: (result: string | null) => void;
	private showingConfirmation: boolean = false;
	private notesMode: boolean = false;
	private notesText: string = "";
	private notesEditor: Editor;

	// Cache
	private cachedWidth?: number;
	private cachedLines?: string[];

	// Colors - using proper reset sequences
	private dim = (s: string) => `\x1b[2m${s}\x1b[0m`;
	private bold = (s: string) => `\x1b[1m${s}\x1b[0m`;
	private cyan = (s: string) => `\x1b[36m${s}\x1b[0m`;
	private green = (s: string) => `\x1b[32m${s}\x1b[0m`;
	private yellow = (s: string) => `\x1b[33m${s}\x1b[0m`;
	private gray = (s: string) => `\x1b[90m${s}\x1b[0m`;

	constructor(questions: ExtractedQuestion[], tui: TUI, onDone: (result: string | null) => void) {
		this.questions = questions;
		this.answers = questions.map(() => "");
		this.tui = tui;
		this.onDone = onDone;

		// Create a minimal theme for the editor
		const editorTheme: EditorTheme = {
			borderColor: this.dim,
			selectList: {
				selectedPrefix: this.cyan,
				selectedText: (s: string) => `\x1b[44m${s}\x1b[0m`,
				description: this.gray,
				scrollInfo: this.dim,
				noMatch: this.dim,
			},
		};

		this.editor = new Editor(tui, editorTheme);
		// Disable the editor's built-in submit (which clears the editor)
		// We'll handle Enter ourselves to preserve the text
		this.editor.disableSubmit = true;
		this.editor.onChange = () => {
			this.invalidate();
			this.tui.requestRender();
		};

		this.notesEditor = new Editor(tui, editorTheme);
		this.notesEditor.onSubmit = (value) => {
			this.notesText = value.trim();
			this.notesMode = false;
			this.invalidate();
			this.tui.requestRender();
		};
	}

	private saveCurrentAnswer(): void {
		this.answers[this.currentIndex] = this.editor.getText();
	}

	private navigateTo(index: number): void {
		if (index < 0 || index >= this.questions.length) return;
		this.saveCurrentAnswer();
		this.currentIndex = index;
		this.editor.setText(this.answers[index] || "");
		this.invalidate();
	}

	private submit(): void {
		this.saveCurrentAnswer();

		// Build the response text
		const parts: string[] = [];
		parts.push(`<qna>`);
		for (let i = 0; i < this.questions.length; i++) {
			const q = this.questions[i];
			const a = this.answers[i]?.trim() || "(no answer)";
			parts.push(`<q n="${i}">${q.question}</q>`);
			parts.push(`<a n="${i}">${a}</a>`);
		}
		parts.push(`</qna>`);
		if (this.notesText) {
			parts.push(`\n<note>${this.notesText}</note>`);
		}

		this.onDone(parts.join("\n").trim());
	}

	private cancel(): void {
		this.onDone(null);
	}

	invalidate(): void {
		this.cachedWidth = undefined;
		this.cachedLines = undefined;
	}

	handleInput(data: string): void {
		// Notes mode: route to notes editor
		if (this.notesMode) {
			if (matchesKey(data, Key.escape)) {
				this.notesMode = false;
				this.notesEditor.setText(this.notesText);
				this.invalidate();
				this.tui.requestRender();
				return;
			}
			this.notesEditor.handleInput(data);
			this.invalidate();
			this.tui.requestRender();
			return;
		}

		// Handle confirmation dialog
		if (this.showingConfirmation) {
			if (matchesKey(data, Key.enter) || data.toLowerCase() === "y") {
				this.submit();
				return;
			}
			if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c")) || data.toLowerCase() === "n") {
				this.showingConfirmation = false;
				this.invalidate();
				this.tui.requestRender();
				return;
			}
			return;
		}

		// Global navigation and commands
		if (matchesKey(data, Key.escape) || matchesKey(data, Key.ctrl("c"))) {
			this.cancel();
			return;
		}

		// Activate notes editor (available when not typing an answer)
		if (matchesKey(data, Key.alt("n"))) {
			this.notesMode = true;
			this.notesEditor.setText(this.notesText);
			this.invalidate();
			this.tui.requestRender();
			return;
		}

		// Tab / Shift+Tab for navigation
		if (matchesKey(data, Key.tab)) {
			if (this.currentIndex < this.questions.length - 1) {
				this.navigateTo(this.currentIndex + 1);
				this.tui.requestRender();
			}
			return;
		}
		if (matchesKey(data, Key.shift("tab"))) {
			if (this.currentIndex > 0) {
				this.navigateTo(this.currentIndex - 1);
				this.tui.requestRender();
			}
			return;
		}

		// Arrow up/down for question navigation when editor is empty
		// (Editor handles its own cursor navigation when there's content)
		if (matchesKey(data, Key.up) && this.editor.getText() === "") {
			if (this.currentIndex > 0) {
				this.navigateTo(this.currentIndex - 1);
				this.tui.requestRender();
				return;
			}
		}
		if (matchesKey(data, Key.down) && this.editor.getText() === "") {
			if (this.currentIndex < this.questions.length - 1) {
				this.navigateTo(this.currentIndex + 1);
				this.tui.requestRender();
				return;
			}
		}

		// Handle Enter ourselves (editor's submit is disabled)
		// Plain Enter moves to next question or shows confirmation on last question
		// Shift+Enter adds a newline (handled by editor)
		if (matchesKey(data, Key.enter) && !matchesKey(data, Key.shift("enter"))) {
			this.saveCurrentAnswer();
			if (this.currentIndex < this.questions.length - 1) {
				this.navigateTo(this.currentIndex + 1);
			} else {
				// On last question - show confirmation
				this.showingConfirmation = true;
			}
			this.invalidate();
			this.tui.requestRender();
			return;
		}

		// Pass to editor
		this.editor.handleInput(data);
		this.invalidate();
		this.tui.requestRender();
	}

	render(width: number): string[] {
		if (this.cachedLines && this.cachedWidth === width) {
			return this.cachedLines;
		}

		const lines: string[] = [];
		const boxWidth = Math.min(width - 4, 120); // Allow wider box
		const contentWidth = boxWidth - 4; // 2 chars padding on each side

		// Helper to create horizontal lines (dim the whole thing at once)
		const horizontalLine = (count: number) => "─".repeat(count);

		// Helper to create a box line
		const boxLine = (content: string, leftPad: number = 2): string => {
			const paddedContent = " ".repeat(leftPad) + content;
			const contentLen = visibleWidth(paddedContent);
			const rightPad = Math.max(0, boxWidth - contentLen - 2);
			return this.dim("│") + paddedContent + " ".repeat(rightPad) + this.dim("│");
		};

		const emptyBoxLine = (): string => {
			return this.dim("│") + " ".repeat(boxWidth - 2) + this.dim("│");
		};

		const padToWidth = (line: string): string => {
			const len = visibleWidth(line);
			return line + " ".repeat(Math.max(0, width - len));
		};

		// Title
		lines.push(padToWidth(this.dim(`╭${horizontalLine(boxWidth - 2)}╮`)));
		const title = `${this.bold(this.cyan("Questions"))} ${this.dim(`(${this.currentIndex + 1}/${this.questions.length})`)}`;
		lines.push(padToWidth(boxLine(title)));
		lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));

		// Progress indicator
		const progressParts: string[] = [];
		for (let i = 0; i < this.questions.length; i++) {
			const answered = (this.answers[i]?.trim() || "").length > 0;
			const current = i === this.currentIndex;
			if (current) {
				progressParts.push(this.cyan("●"));
			} else if (answered) {
				progressParts.push(this.green("●"));
			} else {
				progressParts.push(this.dim("○"));
			}
		}
		lines.push(padToWidth(boxLine(progressParts.join(" "))));
		lines.push(padToWidth(emptyBoxLine()));

		// Current question
		const q = this.questions[this.currentIndex];
		const questionText = `${this.bold("Q:")} ${q.question}`;
		const wrappedQuestion = wrapTextWithAnsi(questionText, contentWidth);
		for (const line of wrappedQuestion) {
			lines.push(padToWidth(boxLine(line)));
		}

		// Context if present
		if (q.context) {
			lines.push(padToWidth(emptyBoxLine()));
			const contextText = this.gray(`> ${q.context}`);
			const wrappedContext = wrapTextWithAnsi(contextText, contentWidth - 2);
			for (const line of wrappedContext) {
				lines.push(padToWidth(boxLine(line)));
			}
		}

		lines.push(padToWidth(emptyBoxLine()));

		// Render the editor component (multi-line input) with padding
		// Skip the first and last lines (editor's own border lines)
		const answerPrefix = this.bold("A: ");
		const editorWidth = contentWidth - 4 - 3; // Extra padding + space for "A: "
		const editorLines = this.editor.render(editorWidth);
		for (let i = 1; i < editorLines.length - 1; i++) {
			if (i === 1) {
				// First content line gets the "A: " prefix
				lines.push(padToWidth(boxLine(answerPrefix + editorLines[i])));
			} else {
				// Subsequent lines get padding to align with the first line
				lines.push(padToWidth(boxLine(`   ${editorLines[i]}`)));
			}
		}

		// Notes section
		if (this.notesMode) {
			lines.push(padToWidth(emptyBoxLine()));
			const notesLabel = `${this.cyan("✎")} ${this.bold("Note:")} ${this.dim("(Enter to save, Esc to discard)")}`;
			lines.push(padToWidth(boxLine(notesLabel)));
			const notesEditorWidth = contentWidth - 4;
			const notesEditorLines = this.notesEditor.render(notesEditorWidth);
			for (let i = 1; i < notesEditorLines.length - 1; i++) {
				lines.push(padToWidth(boxLine(`  ${notesEditorLines[i]}`)));
			}
		} else if (this.notesText) {
			lines.push(padToWidth(emptyBoxLine()));
			const savedNote = `${this.cyan("✎")} ${this.gray(this.notesText)}`;
			const wrappedNote = wrapTextWithAnsi(savedNote, contentWidth);
			for (const line of wrappedNote) {
				lines.push(padToWidth(boxLine(line)));
			}
			lines.push(padToWidth(boxLine(this.dim("Alt+N to edit note"))));
		} else {
			lines.push(padToWidth(emptyBoxLine()));
			lines.push(padToWidth(boxLine(this.dim("Alt+N to add a note"))));
		}

		lines.push(padToWidth(emptyBoxLine()));

		// Confirmation dialog or footer with controls
		if (this.showingConfirmation) {
			lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
			const confirmMsg = `${this.yellow("Submit all answers?")} ${this.dim("(Enter/y to confirm, Esc/n to cancel)")}`;
			lines.push(padToWidth(boxLine(truncateToWidth(confirmMsg, contentWidth))));
		} else {
			lines.push(padToWidth(this.dim(`├${horizontalLine(boxWidth - 2)}┤`)));
			const controls = `${this.dim("Tab/Enter")} next · ${this.dim("Shift+Tab")} prev · ${this.dim("Shift+Enter")} newline · ${this.dim("Alt+N")} note · ${this.dim("Esc")} cancel`;
			lines.push(padToWidth(boxLine(truncateToWidth(controls, contentWidth))));
		}
		lines.push(padToWidth(this.dim(`╰${horizontalLine(boxWidth - 2)}╯`)));

		this.cachedWidth = width;
		this.cachedLines = lines;
		return lines;
	}
}

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}...`);
			loader.onAbort = () => done({ 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 done({ kind: "ok", result: result.result });
					case "cancelled":
						return done({ kind: "cancelled" });
					case "model_error":
						return done({
							kind: "error",
							message: `${result.model.name} returned an error with no content`,
						});
					case "parse_error":
						return done({ kind: "error", message: result.message });
				}
			};

			doExtract().catch((err) =>
				done({
					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,
	});
}