pi-ext

import {
	defineTool,
	type ExtensionContext,
	keyHint,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { askInTui } from "./ask-overlay.js";
import { span } from "./src/debug.ts";

const DESCRIPTION =
	"Ask the user 1-4 clarifying questions with multiple-choice options and wait for answers in interactive sessions. " +
	"Non-interactive sessions return unavailable without asking. " +
	"This is the equivalent of Claude Code's AskUserQuestion and Codex's request_user_input. " +
	"Use it to gather preferences or requirements, clarify ambiguous instructions, or get decisions on implementation choices as you work. " +
	"The user can always answer with free-form text instead of picking an option.";

export const OTHER_LABEL = "Other (free text)";

const parameters = Type.Object({
	questions: Type.Array(
		Type.Object({
			id: Type.Optional(
				Type.String({
					description:
						"Stable snake_case identifier for mapping answers. Derived from the question if omitted.",
				}),
			),
			question: Type.String({
				description:
					"The complete question to ask. Clear, specific, ends with a question mark.",
			}),
			header: Type.Optional(
				Type.String({
					description:
						'Very short label shown in the UI (max 12 chars), e.g. "Auth method". Derived if omitted.',
				}),
			),
			multiSelect: Type.Optional(
				Type.Boolean({
					default: false,
					description: "Allow selecting multiple options instead of just one.",
				}),
			),
			options: Type.Array(
				Type.Object({
					label: Type.String({
						description: "Concise display text (1-5 words).",
					}),
					description: Type.String({
						description:
							"One short sentence explaining impact/tradeoff if selected.",
					}),
				}),
				{
					minItems: 2,
					maxItems: 4,
					description:
						'Distinct, mutually exclusive choices (unless multiSelect). Put the recommended option first and suffix its label with "(Recommended)". Do not include an "Other" option; free-form input is always available.',
				},
			),
		}),
		{
			minItems: 1,
			maxItems: 4,
			description: "Questions to ask the user (1-4). Prefer fewer.",
		},
	),
});

export type AskOption = { label: string; description: string };

export type AskQuestionInput = {
	id?: string;
	question: string;
	header?: string;
	multiSelect?: boolean;
	options: AskOption[];
};

export type AskQuestion = {
	key: string;
	header: string;
	question: string;
	multiSelect: boolean;
	options: AskOption[];
};

export type AskAnswer = { answers: string[]; notes?: string };

export type AskFlowResult =
	| { status: "answered"; answers: Record<string, AskAnswer> }
	| { status: "cancelled" };

type AskDetails =
	| {
			kind: "ask";
			status: "answered";
			questions: Array<AskQuestion & { answer?: AskAnswer }>;
	  }
	| { kind: "ask"; status: "cancelled" }
	| { kind: "ask"; status: "unavailable" }
	| { kind: "ask"; status: "error"; message: string };

const OTHER_PATTERN = /^\s*other\b/i;
const CANCELLED_MESSAGE = "The user interrupted. Stop and wait for direction.";
const UNAVAILABLE_MESSAGE =
	"No interactive UI is available in this mode. Proceed with explicit stated assumptions, or ask in plain text in your reply.";

export function normalizeQuestions(raw: AskQuestionInput[]): AskQuestion[] {
	if (!Array.isArray(raw) || raw.length === 0)
		throw new Error("ask_user_question requires 1-4 questions.");
	const usedKeys = new Set<string>();
	return raw.map((question, index) => {
		const options = (question.options ?? []).filter(
			(option) => !OTHER_PATTERN.test(option.label),
		);
		if (options.length === 0)
			throw new Error(
				`Question ${index + 1} has no usable options; provide 2-4 distinct options and no "Other".`,
			);
		const baseKey = (question.id ?? question.question).trim();
		let key = baseKey;
		let suffix = 2;
		while (usedKeys.has(key)) key = `${baseKey}-${suffix++}`;
		usedKeys.add(key);
		return {
			key,
			header: (question.header ?? question.id ?? question.question)
				.trim()
				.slice(0, 12),
			question: question.question,
			multiSelect: question.multiSelect === true,
			options,
		};
	});
}

export async function runAskFlow(
	questions: AskQuestion[],
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<AskFlowResult> {
	const answers: Record<string, AskAnswer> = {};
	for (const [index, question] of questions.entries()) {
		const title = `Question ${index + 1}/${questions.length} · ${question.header}`;
		const answer =
			ctx.mode === "tui"
				? await askInTui(title, question, ctx, signal)
				: question.multiSelect
					? await askMultiSelect(title, question, ctx, signal)
					: await askSingleSelect(title, question, ctx, signal);
		if (!answer) return { status: "cancelled" };
		answers[question.key] = answer;
	}
	return { status: "answered", answers };
}

async function askSingleSelect(
	title: string,
	question: AskQuestion,
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<AskAnswer | undefined> {
	const entries = question.options.map(
		(option) => `${option.label}: ${option.description}`,
	);
	entries.push(OTHER_LABEL);
	const choice = await ctx.ui.select(title, entries, { signal });
	if (choice === undefined) return undefined;
	const index = entries.indexOf(choice);
	if (index < 0 || index >= question.options.length)
		return await askOther(title, ctx, signal);
	return { answers: [question.options[index].label] };
}

async function askMultiSelect(
	title: string,
	question: AskQuestion,
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<AskAnswer | undefined> {
	const selected: string[] = [];
	for (const option of question.options) {
		if (await ctx.ui.confirm(title, `Include ${option.label}?`, { signal }))
			selected.push(option.label);
	}
	const wrapUp = await ctx.ui.select(title, ["Done", OTHER_LABEL, "Cancel"], {
		signal,
	});
	if (wrapUp === undefined || wrapUp === "Cancel") return undefined;
	if (wrapUp === OTHER_LABEL) {
		const other = await askOther(title, ctx, signal);
		if (!other) return undefined;
		selected.push(...other.answers);
	}
	return { answers: selected };
}

async function askOther(
	title: string,
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<AskAnswer | undefined> {
	const text = await ctx.ui.input(
		`${title} · your answer`,
		"Type your answer",
		{ signal },
	);
	const trimmed = text?.trim();
	if (!trimmed) return undefined;
	return { answers: [trimmed] };
}

function textResult(payload: unknown): string {
	return JSON.stringify(payload);
}

function cancelledResult() {
	return {
		content: [
			{
				type: "text" as const,
				text: textResult({ status: "cancelled", message: CANCELLED_MESSAGE }),
			},
		],
		details: { kind: "ask", status: "cancelled" as const },
		terminate: true,
	};
}

function unavailableResult() {
	return {
		content: [
			{
				type: "text" as const,
				text: textResult({
					status: "unavailable",
					message: UNAVAILABLE_MESSAGE,
				}),
			},
		],
		details: { kind: "ask", status: "unavailable" as const },
	};
}

function errorResult(message: string) {
	return {
		content: [
			{
				type: "text" as const,
				text: textResult({ status: "error", message }),
			},
		],
		details: { kind: "ask", status: "error" as const, message },
	};
}

function answeredResult(
	questions: AskQuestion[],
	answers: Record<string, AskAnswer>,
) {
	return {
		content: [
			{
				type: "text" as const,
				text: textResult({ status: "answered", answers }),
			},
		],
		details: {
			kind: "ask" as const,
			status: "answered" as const,
			questions: questions.map((question) => ({
				...question,
				answer: answers[question.key],
			})),
		},
	};
}

export function createAskTool() {
	return defineTool({
		name: "ask_user_question",
		label: "ask",
		description: DESCRIPTION,
		promptSnippet: "Ask the human structured multiple-choice questions",
		promptGuidelines: [
			"Use ask_user_question only when the human must decide; verify or safely infer what you can.",
			'For ask_user_question, provide 2-4 distinct options per question, recommended option first, its label suffixed with "(Recommended)".',
			'Never add an "Other" option to ask_user_question; the user always has free-form input available.',
		],
		parameters,
		renderCall(args, theme, context) {
			const text =
				(context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
			const questions = Array.isArray(args.questions) ? args.questions : [];
			const count = questions.length;
			let output = `${theme.fg("toolTitle", theme.bold("ask_user_question"))} ${theme.fg("toolOutput", `· ${count} question${count === 1 ? "" : "s"}`)}`;
			if (!context.expanded) {
				output += theme.fg(
					"dim",
					` (${keyHint("app.tools.expand", "to expand")})`,
				);
			} else {
				for (const question of questions) {
					output += `\n${theme.fg("accent", String(question.header ?? ""))} ${theme.fg("toolOutput", String(question.question ?? ""))}`;
					for (const option of question.options ?? []) {
						output += `\n  ${theme.fg("toolOutput", String(option.label ?? ""))} ${theme.fg("dim", String(option.description ?? ""))}`;
					}
				}
			}
			text.setText(output);
			return text;
		},
		renderResult(result, { expanded }, theme, context) {
			const text =
				(context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
			const details = result.details as AskDetails | undefined;
			let output: string;
			if (details?.kind !== "ask") {
				output = theme.fg("toolOutput", "ask_user_question finished");
			} else if (details.status === "answered") {
				const answered = details.questions.filter((q) => q.answer).length;
				if (!expanded) {
					output = theme.fg(
						"toolOutput",
						`${answered}/${details.questions.length} answered`,
					);
				} else {
					const lines = details.questions.map((question) => {
						const answer = question.answer
							? [question.answer.answers.join(", "), question.answer.notes]
									.filter(Boolean)
									.join(" · ")
							: "unanswered";
						return `\n${theme.fg("toolOutput", question.question)} ${theme.fg("accent", answer)}`;
					});
					output = `${theme.fg("toolOutput", `${answered}/${details.questions.length} answered`)}${lines.join("")}`;
				}
			} else if (details.status === "cancelled") {
				output = theme.fg("warning", "cancelled");
			} else if (details.status === "unavailable") {
				output = theme.fg("dim", "no interactive UI available");
			} else {
				output = theme.fg("warning", `error: ${details.message}`);
			}
			text.setText(output);
			return text;
		},
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			const finish = span?.("ask.flow");
			let questions: AskQuestion[];
			try {
				questions = normalizeQuestions(params.questions);
			} catch (error) {
				finish?.("finish", { status: "invalid" });
				return errorResult(
					error instanceof Error ? error.message : String(error),
				);
			}
			if (!ctx.hasUI) {
				finish?.("finish", { count: questions.length, status: "unavailable" });
				return unavailableResult();
			}
			let flow: AskFlowResult;
			try {
				flow = await runAskFlow(questions, ctx, signal);
			} catch {
				finish?.("error");
				return cancelledResult();
			}
			finish?.("finish", { count: questions.length, status: flow.status });
			if (flow.status === "cancelled") return cancelledResult();
			return answeredResult(questions, flow.answers);
		},
	});
}