next-ai-draw-io/components/model-selector.tsx

"use client"

import {
    AlertTriangle,
    Bot,
    Check,
    ChevronDown,
    Server,
    Settings2,
} from "lucide-react"
import { useEffect, useMemo, useRef, useState } from "react"
import {
    ModelSelectorContent,
    ModelSelectorEmpty,
    ModelSelectorGroup,
    ModelSelectorInput,
    ModelSelectorItem,
    ModelSelectorList,
    ModelSelectorLogo,
    ModelSelectorName,
    ModelSelector as ModelSelectorRoot,
    ModelSelectorSeparator,
    ModelSelectorTrigger,
} from "@/components/ai-elements/model-selector"
import { ButtonWithTooltip } from "@/components/button-with-tooltip"
import { useDictionary } from "@/hooks/use-dictionary"
import type { FlattenedModel } from "@/lib/types/model-config"
import { cn } from "@/lib/utils"

interface ModelSelectorProps {
    models: FlattenedModel[]
    selectedModelId: string | undefined
    onSelect: (modelId: string | undefined) => void
    onConfigure: () => void
    disabled?: boolean
    showUnvalidatedModels?: boolean
}

// Map our provider names to models.dev logo names
const PROVIDER_LOGO_MAP: Record<string, string> = {
    openai: "openai",
    anthropic: "anthropic",
    google: "google",
    azure: "azure",
    bedrock: "amazon-bedrock",
    openrouter: "openrouter",
    deepseek: "deepseek",
    siliconflow: "siliconflow",
    sglang: "openai", // SGLang is OpenAI-compatible, use OpenAI logo
    gateway: "vercel",
    edgeone: "tencent-cloud",
    doubao: "bytedance",
}

// Group models by providerLabel (handles duplicate providers)
function groupModelsByProvider(
    models: FlattenedModel[],
): Map<string, { provider: string; models: FlattenedModel[] }> {
    const groups = new Map<
        string,
        { provider: string; models: FlattenedModel[] }
    >()
    for (const model of models) {
        const key = model.providerLabel
        const existing = groups.get(key)
        if (existing) {
            existing.models.push(model)
        } else {
            groups.set(key, { provider: model.provider, models: [model] })
        }
    }
    return groups
}

export function ModelSelector({
    models,
    selectedModelId,
    onSelect,
    onConfigure,
    disabled = false,
    showUnvalidatedModels = false,
}: ModelSelectorProps) {
    const dict = useDictionary()
    const [open, setOpen] = useState(false)
    // Filter models based on showUnvalidatedModels setting
    const displayModels = useMemo(() => {
        if (showUnvalidatedModels) {
            return models
        }
        return models.filter((m) => m.validated === true)
    }, [models, showUnvalidatedModels])
    const groupedModels = useMemo(
        () => groupModelsByProvider(displayModels),
        [displayModels],
    )

    // Find selected model for display
    const selectedModel = useMemo(
        () => models.find((m) => m.id === selectedModelId),
        [models, selectedModelId],
    )

    const handleSelect = (value: string) => {
        if (value === "__configure__") {
            onConfigure()
        } else if (value === "__server_default__") {
            onSelect(undefined)
        } else {
            onSelect(value)
        }
        setOpen(false)
    }

    const tooltipContent = selectedModel
        ? `${selectedModel.modelId} ${dict.modelConfig.clickToChange}`
        : `${dict.modelConfig.usingServerDefault} ${dict.modelConfig.clickToChange}`

    const wrapperRef = useRef<HTMLDivElement | null>(null)
    const [showLabel, setShowLabel] = useState(true)

    // Threshold (px) under which we hide the label (tweak as needed)
    const HIDE_THRESHOLD = 240
    const SHOW_THRESHOLD = 260
    useEffect(() => {
        const el = wrapperRef.current
        if (!el) return

        const target = el.parentElement ?? el

        const ro = new ResizeObserver((entries) => {
            for (const entry of entries) {
                const width = entry.contentRect.width
                setShowLabel((prev) => {
                    // if currently showing and width dropped below hide threshold -> hide
                    if (prev && width <= HIDE_THRESHOLD) return false
                    // if currently hidden and width rose above show threshold -> show
                    if (!prev && width >= SHOW_THRESHOLD) return true
                    // otherwise keep previous state (hysteresis)
                    return prev
                })
            }
        })

        ro.observe(target)

        const initialWidth = target.getBoundingClientRect().width
        setShowLabel(initialWidth >= SHOW_THRESHOLD)

        return () => ro.disconnect()
    }, [])

    return (
        <div ref={wrapperRef} className="inline-block">
            <ModelSelectorRoot open={open} onOpenChange={setOpen}>
                <ModelSelectorTrigger asChild>
                    <ButtonWithTooltip
                        tooltipContent={tooltipContent}
                        variant="ghost"
                        size="sm"
                        disabled={disabled}
                        className={cn(
                            "hover:bg-accent gap-1.5 h-8 px-2 transition-all duration-150 ease-in-out",
                            !showLabel && "px-1.5 justify-center",
                        )}
                        // accessibility: expose label to screen readers
                        aria-label={tooltipContent}
                    >
                        <Bot className="h-4 w-4 flex-shrink-0 text-muted-foreground" />
                        {/* show/hide visible label based on measured width */}
                        {showLabel ? (
                            <span className="text-xs truncate">
                                {selectedModel
                                    ? selectedModel.modelId
                                    : dict.modelConfig.default}
                            </span>
                        ) : (
                            // Keep an sr-only label for screen readers when hidden
                            <span className="sr-only">
                                {selectedModel
                                    ? selectedModel.modelId
                                    : dict.modelConfig.default}
                            </span>
                        )}
                        <ChevronDown className="h-3 w-3 flex-shrink-0 text-muted-foreground" />
                    </ButtonWithTooltip>
                </ModelSelectorTrigger>

                <ModelSelectorContent title={dict.modelConfig.selectModel}>
                    <ModelSelectorInput
                        placeholder={dict.modelConfig.searchModels}
                    />
                    <ModelSelectorList className="[&::-webkit-scrollbar]:hidden [-ms-overflow-style:none] [scrollbar-width:none]">
                        <ModelSelectorEmpty>
                            {displayModels.length === 0 && models.length > 0
                                ? dict.modelConfig.noVerifiedModels
                                : dict.modelConfig.noModelsFound}
                        </ModelSelectorEmpty>

                        {/* Server Default Option */}
                        <ModelSelectorGroup heading={dict.modelConfig.default}>
                            <ModelSelectorItem
                                value="__server_default__"
                                onSelect={handleSelect}
                                className={cn(
                                    "cursor-pointer",
                                    !selectedModelId && "bg-accent",
                                )}
                            >
                                <Check
                                    className={cn(
                                        "mr-2 h-4 w-4",
                                        !selectedModelId
                                            ? "opacity-100"
                                            : "opacity-0",
                                    )}
                                />
                                <Server className="mr-2 h-4 w-4 text-muted-foreground" />
                                <ModelSelectorName>
                                    {dict.modelConfig.serverDefault}
                                </ModelSelectorName>
                            </ModelSelectorItem>
                        </ModelSelectorGroup>

                        {/* Configured Models by Provider */}
                        {Array.from(groupedModels.entries()).map(
                            ([
                                providerLabel,
                                { provider, models: providerModels },
                            ]) => (
                                <ModelSelectorGroup
                                    key={providerLabel}
                                    heading={providerLabel}
                                >
                                    {providerModels.map((model) => (
                                        <ModelSelectorItem
                                            key={model.id}
                                            value={model.modelId}
                                            onSelect={() =>
                                                handleSelect(model.id)
                                            }
                                            className="cursor-pointer"
                                        >
                                            <Check
                                                className={cn(
                                                    "mr-2 h-4 w-4",
                                                    selectedModelId === model.id
                                                        ? "opacity-100"
                                                        : "opacity-0",
                                                )}
                                            />
                                            <ModelSelectorLogo
                                                provider={
                                                    PROVIDER_LOGO_MAP[
                                                        provider
                                                    ] || provider
                                                }
                                                className="mr-2"
                                            />
                                            <ModelSelectorName>
                                                {model.modelId}
                                            </ModelSelectorName>
                                            {model.validated !== true && (
                                                <span
                                                    title={
                                                        dict.modelConfig
                                                            .unvalidatedModelWarning
                                                    }
                                                >
                                                    <AlertTriangle className="ml-auto h-3 w-3 text-warning" />
                                                </span>
                                            )}
                                        </ModelSelectorItem>
                                    ))}
                                </ModelSelectorGroup>
                            ),
                        )}

                        {/* Configure Option */}
                        <ModelSelectorSeparator />
                        <ModelSelectorGroup>
                            <ModelSelectorItem
                                value="__configure__"
                                onSelect={handleSelect}
                                className="cursor-pointer"
                            >
                                <Settings2 className="mr-2 h-4 w-4" />
                                <ModelSelectorName>
                                    {dict.modelConfig.configureModels}
                                </ModelSelectorName>
                            </ModelSelectorItem>
                        </ModelSelectorGroup>
                        {/* Info text */}
                        <div className="px-3 py-2 text-xs text-muted-foreground border-t">
                            {showUnvalidatedModels
                                ? dict.modelConfig.allModelsShown
                                : dict.modelConfig.onlyVerifiedShown}
                        </div>
                    </ModelSelectorList>
                </ModelSelectorContent>
            </ModelSelectorRoot>
        </div>
    )
}