Model Selector

"use client"

import * as React from "react"
import { CheckIcon, ChevronsUpDownIcon, SearchIcon } from "lucide-react"

import { cn } from "@/lib/utils"
import { Badge } from "@/components/ui/badge"
import { Button } from "@/components/ui/button"
import {
  Command,
  CommandEmpty,
  CommandGroup,
  CommandInput,
  CommandItem,
  CommandList,
} from "@/components/ui/command"
import {
  Popover,
  PopoverContent,
  PopoverTrigger,
} from "@/components/ui/popover"

export interface ModelSelectorValue {
  provider: string
  model: string
}

export interface ModelSelectorModel {
  provider: string
  providerLabel?: string
  model: string
  label?: string
  description?: string
  capabilities?: string[]
}

export interface ModelSelectorProps {
  models: ModelSelectorModel[]
  value?: ModelSelectorValue | null
  onValueChange?: (value: ModelSelectorValue, model: ModelSelectorModel) => void
  placeholder?: string
  searchPlaceholder?: string
  emptyText?: string
  disabled?: boolean
  className?: string
  contentClassName?: string
  showCapabilities?: boolean
}

function serializeValue(value: ModelSelectorValue) {
  return `${value.provider}:${value.model}`
}

function getModelLabel(model: ModelSelectorModel) {
  return model.label ?? model.model
}

export function ModelSelector({
  models,
  value,
  onValueChange,
  placeholder = "Select model",
  searchPlaceholder = "Search models...",
  emptyText = "No models found.",
  disabled = false,
  className,
  contentClassName,
  showCapabilities = true,
}: ModelSelectorProps) {
  const [open, setOpen] = React.useState(false)
  const selected = value
    ? models.find(
        (model) =>
          model.provider === value.provider && model.model === value.model
      )
    : null
  const groups = React.useMemo(() => {
    const grouped = new Map<string, ModelSelectorModel[]>()

    for (const model of models) {
      const key = model.providerLabel ?? model.provider
      grouped.set(key, [...(grouped.get(key) ?? []), model])
    }

    return Array.from(grouped.entries())
  }, [models])

  return (
    <Popover open={open} onOpenChange={setOpen}>
      <PopoverTrigger asChild>
        <Button
          type="button"
          variant="outline"
          role="combobox"
          aria-expanded={open}
          disabled={disabled}
          className={cn(
            "h-10 w-full justify-between rounded-md px-3 font-normal",
            !selected && "text-muted-foreground",
            className
          )}
        >
          <span className="flex min-w-0 items-center gap-2">
            <span className="bg-primary size-2 shrink-0 rounded-full" />
            <span className="truncate">
              {selected ? getModelLabel(selected) : placeholder}
            </span>
          </span>
          <ChevronsUpDownIcon className="text-muted-foreground size-4 shrink-0" />
        </Button>
      </PopoverTrigger>
      <PopoverContent
        align="start"
        className={cn("w-[min(420px,calc(100vw-2rem))] p-0", contentClassName)}
      >
        <Command>
          <div className="border-border flex items-center border-b px-3">
            <SearchIcon className="text-muted-foreground mr-2 size-4 shrink-0" />
            <CommandInput
              placeholder={searchPlaceholder}
              className="h-10 border-0 px-0 focus:ring-0"
            />
          </div>
          <CommandList className="max-h-80">
            <CommandEmpty>{emptyText}</CommandEmpty>
            {groups.map(([providerLabel, providerModels]) => (
              <CommandGroup key={providerLabel} heading={providerLabel}>
                {providerModels.map((model) => {
                  const modelValue = {
                    provider: model.provider,
                    model: model.model,
                  }
                  const selectedValue =
                    value &&
                    serializeValue(value) === serializeValue(modelValue)

                  return (
                    <CommandItem
                      key={serializeValue(modelValue)}
                      value={`${providerLabel} ${getModelLabel(model)} ${
                        model.description ?? ""
                      } ${(model.capabilities ?? []).join(" ")}`}
                      onSelect={() => {
                        onValueChange?.(modelValue, model)
                        setOpen(false)
                      }}
                      className="items-start gap-3 py-3"
                    >
                      <CheckIcon
                        className={cn(
                          "mt-0.5 size-4 shrink-0",
                          selectedValue ? "opacity-100" : "opacity-0"
                        )}
                      />
                      <span className="min-w-0 flex-1">
                        <span className="block truncate text-sm font-medium">
                          {getModelLabel(model)}
                        </span>
                        {model.description ? (
                          <span className="text-muted-foreground mt-0.5 line-clamp-2 block text-xs leading-5">
                            {model.description}
                          </span>
                        ) : null}
                        {showCapabilities && model.capabilities?.length ? (
                          <span className="mt-2 flex flex-wrap gap-1">
                            {model.capabilities.map((capability) => (
                              <Badge
                                key={capability}
                                variant="secondary"
                                className="rounded-sm px-1.5 py-0 text-[11px]"
                              >
                                {capability}
                              </Badge>
                            ))}
                          </span>
                        ) : null}
                      </span>
                    </CommandItem>
                  )
                })}
              </CommandGroup>
            ))}
          </CommandList>
        </Command>
      </PopoverContent>
    </Popover>
  )
}