pi-ext

import type { ExtensionContext, Theme } from "@earendil-works/pi-coding-agent";
import {
	Editor,
	type KeybindingsManager,
	type SelectItem,
	type TUI,
	type TuiMouseEvent,
	type TuiMouseEventResult,
	truncateToWidth,
	visibleWidth,
	wrapTextWithAnsi,
} from "@earendil-works/pi-tui";
import type { AskAnswer, AskQuestion } from "./ask-tool.js";

const OTHER_VALUE = "__other__";
const OTHER_ITEM: SelectItem = {
	value: OTHER_VALUE,
	label: "Other…",
	description: "Type your own answer",
};

export async function askInTui(
	title: string,
	question: AskQuestion,
	ctx: ExtensionContext,
	signal?: AbortSignal,
): Promise<AskAnswer | undefined> {
	const promise = ctx.ui.custom<AskAnswer | undefined>(
		(tui, theme, keybindings, done) =>
			new AskQuestionOverlay(tui, theme, keybindings, done, title, question),
	);
	return await abortable(promise, signal);
}

function abortable<T>(promise: Promise<T>, signal?: AbortSignal): Promise<T> {
	if (!signal) return promise;
	return new Promise<T>((resolve, reject) => {
		const onAbort = () => resolve(undefined as T);
		if (signal.aborted) {
			resolve(undefined as T);
			return;
		}
		signal.addEventListener("abort", onAbort, { once: true });
		promise.then(
			(value) => {
				signal.removeEventListener("abort", onAbort);
				resolve(value);
			},
			(reason) => {
				signal.removeEventListener("abort", onAbort);
				reject(reason);
			},
		);
	});
}

type OverlayMode = "list" | "other" | "notes";

class AskQuestionOverlay {
	private settled = false;
	private mode: OverlayMode = "list";
	private readonly list: OptionList;
	private readonly otherInput: Editor;
	private readonly notesInput: Editor;
	private otherText = "";
	private otherSelected = false;
	private focusedState = true;
	private listRow = 0;
	private listHeight = 0;
	private inputRow: number | undefined;
	private inputHeight = 0;

	constructor(
		private readonly tui: TUI,
		private readonly theme: Theme,
		private readonly keybindings: KeybindingsManager,
		private readonly done: (answer: AskAnswer | undefined) => void,
		private readonly title: string,
		private readonly question: AskQuestion,
	) {
		const items: SelectItem[] = [
			...question.options.map((option) => ({
				value: option.label,
				label: option.label,
				description: option.description,
			})),
			OTHER_ITEM,
		];
		this.list = new OptionList(
			items,
			question.multiSelect,
			theme,
			() => this.otherSelected,
			() => this.activateCurrent(),
		);
		const editorTheme = {
			borderColor: (text: string) => theme.fg("border", text),
			selectList: {
				selectedPrefix: (text: string) => theme.fg("accent", text),
				selectedText: (text: string) => theme.fg("accent", text),
				description: (text: string) => theme.fg("dim", text),
				scrollInfo: (text: string) => theme.fg("dim", text),
				noMatch: (text: string) => theme.fg("warning", text),
			},
		};
		this.otherInput = new Editor(tui, editorTheme, { paddingX: 1 });
		this.otherInput.onSubmit = (value) => this.submitOther(value);
		this.notesInput = new Editor(tui, editorTheme, { paddingX: 1 });
		this.notesInput.onSubmit = (value) => this.submitFromNotes(value);
	}

	get focused(): boolean {
		return this.focusedState;
	}

	set focused(value: boolean) {
		this.focusedState = value;
		this.syncInputFocus();
	}

	handleInput(data: string): void {
		if (this.mode === "other" || this.mode === "notes") {
			if (this.keybindings.matches(data, "tui.select.cancel")) {
				this.setMode("list");
				return;
			}
			this.activeInput().handleInput(data);
			this.tui.requestRender();
			return;
		}
		if (
			this.keybindings.matches(data, "tui.select.up") ||
			this.keybindings.matches(data, "tui.select.down")
		) {
			this.list.move(this.keybindings.matches(data, "tui.select.up") ? -1 : 1);
		} else if (this.keybindings.matches(data, "tui.select.confirm")) {
			this.confirm();
		} else if (this.keybindings.matches(data, "tui.input.tab")) {
			this.setMode("notes");
		} else if (this.keybindings.matches(data, "tui.select.cancel")) {
			this.finish(undefined);
		} else if (data === " " && this.question.multiSelect) {
			this.toggleCurrent();
		}
		this.tui.requestRender();
	}

	handleMouse(event: TuiMouseEvent): TuiMouseEventResult | undefined {
		if (event.y >= this.listRow && event.y < this.listRow + this.listHeight) {
			return this.list.handleMouse(
				localMouseEvent(event, this.listRow, this.listHeight),
			);
		}
		if (
			this.inputRow !== undefined &&
			event.y >= this.inputRow &&
			event.y < this.inputRow + this.inputHeight
		) {
			const result = this.activeInput().handleMouse(
				localMouseEvent(event, this.inputRow, this.inputHeight),
			);
			return result ? { ...result, focus: true } : undefined;
		}
		return undefined;
	}

	render(width: number): string[] {
		const lines = [
			...wrapStyled(
				this.theme.fg("accent", this.theme.bold(this.title)),
				width,
			),
			...wrapStyled(this.theme.fg("toolOutput", this.question.question), width),
			"",
		];
		this.listRow = lines.length;
		const listLines = this.list.render(width);
		this.listHeight = listLines.length;
		lines.push(...listLines);
		this.inputRow = undefined;
		this.inputHeight = 0;
		if (this.mode === "other" || this.mode === "notes") {
			lines.push(
				"",
				...wrapStyled(
					this.theme.fg("accent", this.mode === "other" ? "Answer:" : "Notes:"),
					width,
				),
			);
			this.inputRow = lines.length;
			const inputLines = this.activeInput().render(width);
			this.inputHeight = inputLines.length;
			lines.push(...inputLines);
		}
		lines.push(
			"",
			...wrapStyled(this.theme.fg("dim", this.footerHint()), width),
		);
		return lines;
	}

	invalidate(): void {
		this.list.invalidate();
		this.otherInput.invalidate();
		this.notesInput.invalidate();
	}

	private footerHint(): string {
		if (this.mode === "other") return "enter submit · esc back";
		if (this.mode === "notes") return "enter submit · esc back to options";
		return this.question.multiSelect
			? "space/click toggle · enter submit · tab notes · esc cancel"
			: "enter/click select · tab notes · esc cancel";
	}

	private confirm(): void {
		if (this.list.current()?.value === OTHER_VALUE) {
			this.setMode("other");
			return;
		}
		this.finish(
			this.answer(
				this.question.multiSelect
					? this.multiAnswers()
					: this.list.current()
						? [this.list.current()?.label ?? ""]
						: [],
				this.notesDraft(),
			),
		);
	}

	private activateCurrent(): void {
		if (this.question.multiSelect) this.toggleCurrent();
		else this.confirm();
	}

	private toggleCurrent(): void {
		const current = this.list.current();
		if (!current) return;
		if (current.value === OTHER_VALUE) {
			this.setMode("other");
			return;
		}
		this.list.toggle();
	}

	private submitOther(value: string): void {
		const text = value.trim();
		if (!text) return;
		if (this.question.multiSelect) {
			this.otherText = text;
			this.otherSelected = true;
			this.setMode("list");
			return;
		}
		this.finish(this.answer([text], this.notesDraft()));
	}

	private submitFromNotes(value: string): void {
		const notes = value.trim();
		const current = this.list.current();
		const answers = this.question.multiSelect
			? this.multiAnswers()
			: current?.value === OTHER_VALUE
				? this.otherText
					? [this.otherText]
					: []
				: current
					? [current.label]
					: [];
		if (answers.length === 0 && !notes) return;
		this.finish(this.answer(answers, notes || undefined));
	}

	private multiAnswers(): string[] {
		const labels = this.list.selectedLabels();
		return this.otherSelected ? [...labels, this.otherText] : labels;
	}

	private answer(answers: string[], notes?: string): AskAnswer {
		return notes ? { answers, notes } : { answers };
	}

	private notesDraft(): string | undefined {
		return this.notesInput.getText().trim() || undefined;
	}

	private activeInput(): Editor {
		return this.mode === "other" ? this.otherInput : this.notesInput;
	}

	private setMode(mode: OverlayMode): void {
		this.mode = mode;
		this.syncInputFocus();
		this.tui.requestRender();
	}

	private syncInputFocus(): void {
		this.otherInput.focused = this.focusedState && this.mode === "other";
		this.notesInput.focused = this.focusedState && this.mode === "notes";
	}

	private finish(answer: AskAnswer | undefined): void {
		if (this.settled) return;
		this.settled = true;
		this.done(answer);
	}
}

type ItemRows = { start: number; end: number };

class OptionList {
	private index = 0;
	private readonly toggled = new Set<string>();
	private itemRows: ItemRows[] = [];

	constructor(
		private readonly items: SelectItem[],
		private readonly multiSelect: boolean,
		private readonly theme: Theme,
		private readonly isOtherSelected: () => boolean,
		private readonly onActivate: () => void,
	) {}

	current(): SelectItem | undefined {
		return this.items[this.index];
	}

	move(delta: number): void {
		this.index =
			(this.index + delta + this.items.length) % Math.max(this.items.length, 1);
	}

	toggle(): void {
		const value = this.current()?.value;
		if (value === undefined || value === OTHER_VALUE) return;
		if (this.toggled.has(value)) this.toggled.delete(value);
		else this.toggled.add(value);
	}

	selectedLabels(): string[] {
		return this.items
			.filter(
				(item) => this.toggled.has(item.value) && item.value !== OTHER_VALUE,
			)
			.map((item) => item.label);
	}

	handleMouse(event: TuiMouseEvent): TuiMouseEventResult | undefined {
		if (event.type === "wheel" && event.wheelDelta) {
			this.move(event.wheelDelta < 0 ? -1 : 1);
			return { handled: true, render: true };
		}
		if (
			event.button !== "left" ||
			(event.type !== "press" && event.type !== "click")
		)
			return undefined;
		const itemIndex = this.itemRows.findIndex(
			(rows) => event.y >= rows.start && event.y < rows.end,
		);
		if (itemIndex < 0) return undefined;
		this.index = itemIndex;
		if (event.type === "click") this.onActivate();
		return { handled: true, focus: true, render: true };
	}

	invalidate(): void {}

	render(width: number): string[] {
		const lines: string[] = [];
		this.itemRows = [];
		for (const [index, item] of this.items.entries()) {
			const start = lines.length;
			const cursor = index === this.index;
			const checked =
				item.value === OTHER_VALUE
					? this.isOtherSelected()
					: this.toggled.has(item.value);
			const prefix = this.multiSelect
				? `${cursor ? "›" : " "} [${checked ? "x" : " "}] `
				: `${cursor ? "›" : " "} `;
			const labelStyle = cursor ? "accent" : "toolOutput";
			lines.push(
				...wrapIndented(this.theme.fg(labelStyle, item.label), width, prefix),
			);
			if (item.description) {
				lines.push(
					...wrapIndented(this.theme.fg("dim", item.description), width, "  "),
				);
			}
			this.itemRows.push({ start, end: lines.length });
		}
		return lines;
	}
}

function wrapStyled(text: string, width: number): string[] {
	if (width <= 0) return [""];
	return wrapTextWithAnsi(text, width).map((line) =>
		truncateToWidth(line, width, ""),
	);
}

function wrapIndented(text: string, width: number, prefix: string): string[] {
	const prefixWidth = visibleWidth(prefix);
	if (width <= prefixWidth) return [truncateToWidth(prefix, width, "")];
	const continuation = " ".repeat(prefixWidth);
	return wrapTextWithAnsi(text, width - prefixWidth).map((line, index) =>
		truncateToWidth(`${index === 0 ? prefix : continuation}${line}`, width, ""),
	);
}

function localMouseEvent(
	event: TuiMouseEvent,
	row: number,
	height: number,
): TuiMouseEvent {
	return { ...event, y: event.y - row, height };
}