pi-gui/ui/src/components/ModelSelector.tsx

import { useEffect, useId, useMemo, useRef, useState } from "react";
import type { AgentState, Model } from "../workspace/types";

type SelectableModel = Model & { provider: string; id: string };
type ModelGroup = { provider: string; models: SelectableModel[] };

function isSelectableModel(model: Model): model is SelectableModel {
	return Boolean(model.provider?.trim() && model.id?.trim());
}

function modelName(model: SelectableModel) {
	return model.name?.trim() || model.id;
}

function modelMatches(model: SelectableModel, query: string) {
	if (!query) return true;
	return [model.provider, model.id, model.name]
		.filter((value): value is string => Boolean(value))
		.some((value) => value.toLocaleLowerCase().includes(query));
}

function groupModels(models: SelectableModel[], query: string): ModelGroup[] {
	const byProvider = new Map<string, SelectableModel[]>();
	for (const model of models) {
		if (!modelMatches(model, query)) continue;
		const providerModels = byProvider.get(model.provider) ?? [];
		providerModels.push(model);
		byProvider.set(model.provider, providerModels);
	}
	return [...byProvider.entries()]
		.sort(([left], [right]) => left.localeCompare(right))
		.map(([provider, providerModels]) => ({
			provider,
			models: providerModels.sort((left, right) =>
				modelName(left).localeCompare(modelName(right)),
			),
		}));
}

function ProviderGroup({
	group,
	current,
	headingId,
	onSelect,
	disabled,
}: {
	group: ModelGroup;
	current?: AgentState["model"];
	headingId: string;
	onSelect: (model: SelectableModel) => void;
	disabled: boolean;
}) {
	return (
		<div
			className="model-provider-group"
			role="group"
			aria-labelledby={headingId}
		>
			<div className="model-provider-heading" id={headingId}>
				{group.provider}
			</div>
			{group.models.map((model) => {
				const name = modelName(model);
				return (
					<button
						type="button"
						className="model-option"
						role="option"
						aria-selected={
							model.provider === current?.provider && model.id === current?.id
						}
						key={`${model.provider}/${model.id}`}
						disabled={disabled}
						onClick={() => onSelect(model)}
					>
						<span>{name}</span>
						{name !== model.id && <small>{model.id}</small>}
					</button>
				);
			})}
		</div>
	);
}

export function ModelSelector({
	models,
	current,
	onSelect,
	pending = false,
}: {
	models: Model[];
	current?: AgentState["model"];
	onSelect: (model: Model) => Promise<boolean>;
	pending?: boolean;
}) {
	const [open, setOpen] = useState(false);
	const [query, setQuery] = useState("");
	const [selecting, setSelecting] = useState(false);
	const rootRef = useRef<HTMLDivElement>(null);
	const triggerRef = useRef<HTMLButtonElement>(null);
	const searchRef = useRef<HTMLInputElement>(null);
	const labelId = useId();
	const valueId = useId();
	const menuId = useId();
	const availableModels = useMemo(
		() => models.filter(isSelectableModel),
		[models],
	);
	const groups = useMemo(
		() => groupModels(availableModels, query.trim().toLocaleLowerCase()),
		[availableModels, query],
	);
	const resultCount = groups.reduce(
		(count, group) => count + group.models.length,
		0,
	);
	const selectedModel = availableModels.find(
		(model) => model.provider === current?.provider && model.id === current?.id,
	);
	const selectedName = selectedModel
		? modelName(selectedModel)
		: current?.name?.trim() || current?.id;
	const selectedLabel = current?.provider
		? `${current.provider} · ${selectedName ?? "Current model"}`
		: "Current model";
	const busy = pending || selecting;

	useEffect(() => {
		if (!open) return;
		const focusFrame = window.requestAnimationFrame(() =>
			searchRef.current?.focus(),
		);
		const closeOnEscape = (event: KeyboardEvent) => {
			if (event.key !== "Escape") return;
			event.preventDefault();
			setOpen(false);
			window.requestAnimationFrame(() => triggerRef.current?.focus());
		};
		const closeOutside = (event: PointerEvent) => {
			if (!rootRef.current?.contains(event.target as Node)) setOpen(false);
		};
		window.addEventListener("keydown", closeOnEscape);
		window.addEventListener("pointerdown", closeOutside);
		return () => {
			window.cancelAnimationFrame(focusFrame);
			window.removeEventListener("keydown", closeOnEscape);
			window.removeEventListener("pointerdown", closeOutside);
		};
	}, [open]);

	const chooseModel = async (model: SelectableModel) => {
		if (busy) return;
		setSelecting(true);
		try {
			if (!(await onSelect(model))) return;
			setOpen(false);
			window.requestAnimationFrame(() => triggerRef.current?.focus());
		} finally {
			setSelecting(false);
		}
	};

	return (
		<div className="model-control" aria-busy={busy}>
			<span id={labelId}>Model</span>
			<div className="model-selector" ref={rootRef}>
				<button
					ref={triggerRef}
					type="button"
					className="model-trigger"
					aria-expanded={open}
					aria-haspopup="dialog"
					aria-controls={open ? menuId : undefined}
					aria-labelledby={`${labelId} ${valueId}`}
					disabled={availableModels.length === 0 || busy}
					onClick={() => {
						setQuery("");
						setOpen((value) => !value);
					}}
				>
					<span id={valueId}>{selectedLabel}</span>
				</button>
				{open && (
					<div
						id={menuId}
						className="model-menu"
						role="dialog"
						aria-label="Choose model"
					>
						<label className="sr-only" htmlFor={`${menuId}-search`}>
							Search models
						</label>
						<input
							ref={searchRef}
							id={`${menuId}-search`}
							type="search"
							value={query}
							disabled={busy}
							onChange={(event) => setQuery(event.currentTarget.value)}
							placeholder="Search provider, name, or model ID"
							aria-controls={`${menuId}-results`}
						/>
						<p className="model-result-count" role="status" aria-live="polite">
							{busy
								? "Updating Pi model…"
								: `${resultCount} model${resultCount === 1 ? "" : "s"}`}
						</p>
						{groups.length ? (
							<div
								id={`${menuId}-results`}
								className="model-results"
								role="listbox"
								aria-label="Available models"
							>
								{groups.map((group, index) => (
									<ProviderGroup
										key={group.provider}
										group={group}
										current={current}
										headingId={`${menuId}-provider-${index}`}
										onSelect={(model) => void chooseModel(model)}
										disabled={busy}
									/>
								))}
							</div>
						) : (
							<p className="model-empty muted">No models match “{query.trim()}”.</p>
						)}
					</div>
				)}
			</div>
		</div>
	);
}