DropdownModelProvider.tsx

/* eslint-disable @typescript-eslint/no-explicit-any */
import { useEffect, useState, useRef, useMemo, useCallback, memo } from 'react'
import {
  Popover,
  PopoverContent,
  PopoverTrigger,
} from '@/components/ui/popover'
import { useModelProvider } from '@/hooks/useModelProvider'
import { cn, getProviderTitle, getModelDisplayName } from '@/lib/utils'
import { highlightFzfMatch } from '@/utils/highlight'
import Capabilities from './Capabilities'
import { ModelSourceBadge, MissingModelBadge } from '@/components/ModelSourceBadge'
import { IconSettings, IconX, IconDownload } from '@tabler/icons-react'
import { useNavigate } from '@tanstack/react-router'
import { route } from '@/constants/routes'
import { ModelSetting } from '@/containers/ModelSetting'
import { SamplerPopover } from '@/containers/SamplerPopover'
import ProvidersAvatar from '@/containers/ProvidersAvatar'
import { ModelSupportStatus } from '@/containers/ModelSupportStatus'
import { Fzf } from 'fzf'
import { localStorageKey } from '@/constants/localStorage'
import { useTranslation } from '@/i18n/react-i18next-compat'
import { useFavoriteModel } from '@/hooks/useFavoriteModel'
import { isKnownProvider } from '@/stores/provider-registry-store'
import { EMBEDDING_MODEL_ID } from '@/constants/models'
import { useServiceHub } from '@/hooks/useServiceHub'
import { getLastUsedModel } from '@/utils/getModelToStart'
import { switchToModel } from '@/utils/switchModel'
import { ChevronsUpDown } from 'lucide-react'

interface SearchableModel {
  provider: ModelProvider
  model: Model
  searchStr: string
  value: string
  highlightedId?: string
}

// Helper functions for localStorage
const setLastUsedModel = (provider: string, model: string) => {
  try {
    localStorage.setItem(
      localStorageKey.lastUsedModel,
      JSON.stringify({ provider, model })
    )
  } catch (error) {
    console.debug('Failed to set last used model in localStorage:', error)
  }
}

interface DropdownModelProviderProps {
  /** Show the Sampling popover trigger inside the model bar (hidden for projects). */
  showSampler?: boolean
}

const DropdownModelProvider = memo(function DropdownModelProvider({
  showSampler = true,
}: DropdownModelProviderProps) {
  const {
    providers,
    getProviderByName,
    selectModelProvider,
    getModelBy,
    selectedProvider,
    selectedModel,
    updateProvider,
  } = useModelProvider()
  const [displayModel, setDisplayModel] = useState<string>('')
  const navigate = useNavigate()
  const { t } = useTranslation()
  const { favoriteModels } = useFavoriteModel()
  const serviceHub = useServiceHub()

  // Search state
  const [open, setOpen] = useState(false)
  const [searchValue, setSearchValue] = useState('')
  const searchInputRef = useRef<HTMLInputElement>(null)

  // Helper function to check if a model exists in providers
  const checkModelExists = useCallback(
    (providerName: string, modelId: string) => {
      const provider = providers.find(
        (p) => p.provider === providerName && p.active
      )
      return provider?.models.find((m) => m.id === modelId)
    },
    [providers]
  )

  // Helper function to get context size from model settings
  const getContextSize = useCallback((): number => {
    if (!selectedModel?.settings?.ctx_len?.controller_props?.value) {
      return 16384 // Default context size
    }
    return selectedModel.settings.ctx_len.controller_props.value as number
  }, [selectedModel?.settings?.ctx_len?.controller_props?.value])

  // Function to check if a llamacpp model has vision capabilities and update model capabilities
  const checkAndUpdateModelVisionCapability = useCallback(
    async (modelId: string) => {
      try {
        const hasVision = await serviceHub.models().checkMmprojExists(modelId)
        if (hasVision) {
          // Update the model capabilities to include 'vision'
          const provider = getProviderByName('llamacpp')
          if (provider) {
            const modelIndex = provider.models.findIndex(
              (m) => m.id === modelId
            )
            if (modelIndex !== -1) {
              const model = provider.models[modelIndex]
              const capabilities = model.capabilities || []

              // Add 'vision' capability if not already present AND if user hasn't manually configured capabilities
              // Check if model has a custom capabilities config flag

              const hasUserConfiguredCapabilities =
                (model as any)._userConfiguredCapabilities === true

              if (
                !capabilities.includes('vision') &&
                !hasUserConfiguredCapabilities
              ) {
                const updatedModels = [...provider.models]
                updatedModels[modelIndex] = {
                  ...model,
                  capabilities: [...capabilities, 'vision'],
                  // Mark this as auto-detected, not user-configured
                  _autoDetectedVision: true,
                } as any

                updateProvider('llamacpp', { models: updatedModels })
              }
            }
          }
        }
      } catch (error) {
        console.debug('Error checking mmproj for model:', modelId, error)
      }
    },
    [getProviderByName, updateProvider, serviceHub]
  )

  // Initialize model provider on first mount (no model selected yet)
  useEffect(() => {
    const initializeModel = async () => {
      if (selectedProvider && selectedModel) {
        return
      }

      const lastUsed = getLastUsedModel()
      if (lastUsed && checkModelExists(lastUsed.provider, lastUsed.model)) {
        selectModelProvider(lastUsed.provider, lastUsed.model)
        if (lastUsed.provider === 'llamacpp') {
          await serviceHub
            .models()
            .checkMmprojExistsAndUpdateOffloadMMprojSetting(
              lastUsed.model,
              updateProvider,
              getProviderByName
            )
          await checkAndUpdateModelVisionCapability(lastUsed.model)
        }
      } else {
        const localProvider = providers.find(
          (p) =>
            (p.provider === 'llamacpp-upstream' ||
              p.provider === 'llamacpp' ||
              p.provider === 'mlx') &&
            p.active &&
            p.models.length > 0
        )
        if (localProvider && localProvider.models.length > 0) {
          const firstModel = localProvider.models.find(
            (m) => m.id !== EMBEDDING_MODEL_ID
          )
          if (!firstModel) {
            selectModelProvider('', '')
            return
          }
          selectModelProvider(localProvider.provider, firstModel.id)
          setLastUsedModel(localProvider.provider, firstModel.id)
        } else {
          selectModelProvider('', '')
        }
      }
    }

    initializeModel()
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [
    providers,
    selectModelProvider,
    checkModelExists,
    updateProvider,
    getProviderByName,
    checkAndUpdateModelVisionCapability,
  ])

  // Update display model when selection changes
  useEffect(() => {
    if (selectedProvider && selectedModel) {
      setDisplayModel(getModelDisplayName(selectedModel))
    } else {
      setDisplayModel(t('common:selectAModel'))
    }
  }, [selectedProvider, selectedModel, t])

  // Check vision capabilities for all llamacpp models
  useEffect(() => {
    const checkAllLlamacppModelsForVision = async () => {
      const llamacppProvider = providers.find(
        (p) => p.provider === 'llamacpp' && p.active
      )
      if (llamacppProvider) {
        const checkPromises = llamacppProvider.models.map((model) =>
          checkAndUpdateModelVisionCapability(model.id)
        )
        await Promise.allSettled(checkPromises)
      }
    }

    if (open) {
      checkAllLlamacppModelsForVision()
    }
  }, [open, providers, checkAndUpdateModelVisionCapability])

  // Reset search value when dropdown closes
  const onOpenChange = useCallback((open: boolean) => {
    setOpen(open)
    if (!open) {
      requestAnimationFrame(() => setSearchValue(''))
    } else {
      // Focus search input when opening
      setTimeout(() => {
        searchInputRef.current?.focus()
      }, 100)
    }
  }, [])

  // Clear search and focus input
  const onClearSearch = useCallback(() => {
    setSearchValue('')
    searchInputRef.current?.focus()
  }, [])

  // Jump to the Hub to download a local model. Carry over whatever the user
  // already typed so the Hub search is prefilled instead of starting blank.
  const onDownloadModel = useCallback(() => {
    setOpen(false)
    navigate({
      to: route.hub.index,
      search: searchValue.trim() ? { q: searchValue.trim() } : {},
    })
  }, [navigate, searchValue])

  // Create searchable items from all models
  const searchableItems = useMemo(() => {
    const items: SearchableModel[] = []

    providers.forEach((provider) => {
      if (!provider.active) return

      provider.models.forEach((modelItem) => {
        // Skip embedding models - they can't be used for chat
        if (modelItem.embedding || modelItem.id === EMBEDDING_MODEL_ID) return

        // Skip models that require API key but don't have one (except llamacpp)
        // For custom providers, allow if they have at least one model loaded
        const isPredefined = isKnownProvider(provider.provider)
        if (
          provider &&
          provider.provider !== 'llamacpp' &&
          !provider.api_key?.length &&
          (isPredefined || provider.models.length === 0)
        )
          return

        const capabilities = modelItem.capabilities || []
        const capabilitiesString = capabilities.join(' ')
        const providerTitle = getProviderTitle(provider.provider)

        // Create search string with model id, provider, and capabilities
        const searchStr =
          `${modelItem.id} ${providerTitle} ${provider.provider} ${capabilitiesString}`.toLowerCase()

        items.push({
          provider,
          model: modelItem,
          searchStr,
          value: `${provider.provider}:${modelItem.id}`,
        })
      })
    })

    return items
  }, [providers])

  // Create Fzf instance for fuzzy search
  const fzfInstance = useMemo(() => {
    return new Fzf(searchableItems, {
      selector: (item) =>
        `${getModelDisplayName(item.model)} ${item.model.id}`.toLowerCase(),
    })
  }, [searchableItems])

  // Get favorite models that are currently available
  const favoriteItems = useMemo(() => {
    const matched = searchableItems.filter((item) =>
      favoriteModels.some((fav) => fav.id === item.model.id)
    )
    // A model id can appear under more than one provider (e.g. llamacpp +
    // llamacpp-upstream). Favorites are keyed by model id, so collapse to a
    // single entry per id — otherwise nicknaming one copy makes the favorite
    // show twice (once as the nickname, once as the raw id). Prefer the copy
    // that carries a user nickname (model.displayName); keep first match
    // otherwise. Map.set on an existing key preserves insertion order.
    const byId = new Map<string, SearchableModel>()
    for (const item of matched) {
      const existing = byId.get(item.model.id)
      if (!existing || (!existing.model.displayName && item.model.displayName)) {
        byId.set(item.model.id, item)
      }
    }
    return Array.from(byId.values())
  }, [searchableItems, favoriteModels])

  // Filter models based on search value
  const filteredItems = useMemo(() => {
    if (!searchValue) return searchableItems

    return fzfInstance.find(searchValue.toLowerCase()).map((result) => {
      const item = result.item
      const positions = Array.from(result.positions) || []
      const highlightedId = highlightFzfMatch(
        item.model.id,
        positions,
        'text-accent'
      )

      return {
        ...item,
        highlightedId,
      }
    })
  }, [searchableItems, searchValue, fzfInstance])

  // Group filtered items by provider, excluding favorites when not searching
  const groupedItems = useMemo(() => {
    const groups: Record<string, SearchableModel[]> = {}

    if (!searchValue) {
      const isLocalProvider = (name: string) =>
        name === 'mlx' ||
        name === 'llamacpp' ||
        name === 'llamacpp-upstream' ||
        name === 'foundation-models'

      const activeProviders = providers
        .filter((p) => p.active)
        .sort((a, b) => {
          // Local providers first, regardless of whether they have models
          const aIsLocal = isLocalProvider(a.provider)
          const bIsLocal = isLocalProvider(b.provider)
          if (aIsLocal !== bIsLocal) return aIsLocal ? -1 : 1

          // Within the same group, non-empty providers first
          const aHasModels = a.models.length > 0
          const bHasModels = b.models.length > 0
          if (aHasModels !== bHasModels) return aHasModels ? -1 : 1

          // Custom providers without API key but with models should be treated like "have API key"
          const aIsPredefined = isKnownProvider(a.provider)
          const bIsPredefined = isKnownProvider(b.provider)
          const aHasApiKeyOrCustomModel =
            (a.api_key?.length ?? 0) > 0 ||
            (!aIsPredefined && a.models.length > 0)
          const bHasApiKeyOrCustomModel =
            (b.api_key?.length ?? 0) > 0 ||
            (!bIsPredefined && b.models.length > 0)
          // Providers with API keys or custom with models filled second
          if (aHasApiKeyOrCustomModel && !bHasApiKeyOrCustomModel) return -1
          if (!aHasApiKeyOrCustomModel && bHasApiKeyOrCustomModel) return 1

          // Sort remaining by provider name
          return a.provider.localeCompare(b.provider)
        })

      activeProviders.forEach((provider) => {
        groups[provider.provider] = []
      })
    }

    // Add the filtered items to their respective groups
    filteredItems.forEach((item) => {
      const providerKey = item.provider.provider
      if (!groups[providerKey]) {
        groups[providerKey] = []
      }

      // When not searching, exclude favorite models from regular provider sections
      const isFavorite = favoriteModels.some((fav) => fav.id === item.model.id)
      if (!searchValue && isFavorite) return // Skip adding this item to regular provider section

      groups[providerKey].push(item)
    })

    return groups
  }, [filteredItems, providers, searchValue, favoriteModels])

  const handleSelect = useCallback(
    async (searchableModel: SearchableModel) => {
      // Immediately update display to prevent double-click issues
      setDisplayModel(getModelDisplayName(searchableModel.model))
      setSearchValue('')
      setOpen(false)

      // Optimistically update the global model-provider selection so the
      // provider avatar, capabilities and support-status icons re-render
      // instantly — without waiting for stopAllModels / server restart /
      // registerRemoteProvider to complete inside switchToModel. switchToModel
      // will call this again at the end (idempotent) once the switch is done.
      selectModelProvider(
        searchableModel.provider.provider,
        searchableModel.model.id
      )

      // Fire-and-forget llamacpp mmproj / vision capability checks. These must
      // not block the switch itself.
      if (searchableModel.provider.provider === 'llamacpp') {
        serviceHub
          .models()
          .checkMmprojExistsAndUpdateOffloadMMprojSetting(
            searchableModel.model.id,
            updateProvider,
            getProviderByName
          )
          .catch((error) => {
            console.debug(
              'Error checking mmproj for model:',
              searchableModel.model.id,
              error
            )
          })

        checkAndUpdateModelVisionCapability(searchableModel.model.id).catch(
          (error) => {
            console.debug(
              'Error checking vision capability for model:',
              searchableModel.model.id,
              error
            )
          }
        )
      }

      // Unified switch: stops current engine, (re)starts the Local API Server,
      // registers cloud providers, and synchronises global state.
      switchToModel({
        modelId: searchableModel.model.id,
        providerName: searchableModel.provider.provider,
        serviceHub,
      }).catch((error) => {
        console.error('[DropdownModelProvider] switchToModel failed:', error)
      })
    },
    [
      updateProvider,
      getProviderByName,
      checkAndUpdateModelVisionCapability,
      selectModelProvider,
      serviceHub,
    ]
  )

  const currentModel = selectedModel?.id
    ? getModelBy(selectedModel?.id)
    : undefined

  if (!providers.length) return null

  const provider = getProviderByName(selectedProvider)

  return (
    <Popover open={open} onOpenChange={onOpenChange}>
      <PopoverTrigger asChild>
        <div
          className="border relative z-20 px-4 py-1.5 flex items-center gap-1.5 rounded-full"
          // Prevent the mousedown event from bubbling up to the HeaderPage's
          // data-tauri-drag-region so that clicking or hovering on the model
          // selector pill does not accidentally start a window-drag on macOS.
          // The Popover uses pointerdown / click events internally so this does
          // not affect its open/close behaviour.
          onMouseDown={(e) => e.stopPropagation()}
        >
          <button
            type="button"
            className="font-medium cursor-pointer flex items-center gap-1.5 relative z-20 max-w-50"
          >
            {provider && (
              <div className="shrink-0">
                <ProvidersAvatar provider={provider} />
              </div>
            )}
            <span
              className={cn(
                'text-foreground truncate leading-normal',
                !selectedModel?.id && 'text-muted-foreground'
              )}
            >
              {displayModel}
            </span>
            <ChevronsUpDown className="size-4 shrink-0 text-muted-foreground" />
          </button>
          {currentModel?.settings &&
            provider &&
            (provider.provider === 'llamacpp' ||
              provider.provider === 'llamacpp-upstream' ||
              provider.provider === 'mlx') && (
              <div onClick={(e) => e.stopPropagation()}>
                <ModelSetting
                  model={currentModel as Model}
                  provider={provider}
                />
              </div>
            )}
          <ModelSupportStatus
            modelId={selectedModel?.id}
            provider={selectedProvider}
            contextSize={getContextSize()}
            className="ml-0.5 shrink-0"
          />
          {showSampler && (
            <div className="shrink-0" onClick={(e) => e.stopPropagation()}>
              <SamplerPopover />
            </div>
          )}
        </div>
      </PopoverTrigger>

      <PopoverContent
        className={cn(
          'w-70 p-0 backdrop-blur-2xl bg-background/95 border',
          searchValue.length === 0 && 'h-80'
        )}
        align="start"
        // sideOffset={16}
        // alignOffset={-10}
        side="bottom"
        avoidCollisions={searchValue.length === 0 ? true : false}
      >
        <div className="flex flex-col size-full">
          {/* Search input */}
          <div className="relative p-2 border-b">
            <input
              ref={searchInputRef}
              value={searchValue}
              onChange={(e) => setSearchValue(e.target.value)}
              placeholder={t('common:searchModels')}
              className="text-sm font-normal outline-0"
            />
            {searchValue.length > 0 && (
              <div className="absolute right-2 top-0 bottom-0 flex items-center justify-center">
                <IconX
                  size={16}
                  className="text-muted-foreground cursor-pointer"
                  onClick={onClearSearch}
                />
              </div>
            )}
          </div>

          {/* Model list */}
          <div className="max-h-80 overflow-y-auto">
            {Object.keys(groupedItems).length === 0 && searchValue ? (
              <div className="py-3 px-4 text-sm ">
                {t('common:noModelsFoundFor', { searchValue })}
              </div>
            ) : (
              <div className="py-1">
                {/* Favorites section - only show when not searching */}
                {!searchValue && favoriteItems.length > 0 && (
                  <div className="bg-secondary/30 rounded-sm m-2 py-1">
                    {/* Favorites header */}
                    <div className="flex items-center gap-1.5 px-2 py-1">
                      <span className="text-sm font-medium text-muted-foreground">
                        {t('common:favorites')}
                      </span>
                    </div>

                    {/* Favorite models */}
                    {favoriteItems.map((searchableModel) => {
                      const isSelected =
                        selectedModel?.id === searchableModel.model.id &&
                        selectedProvider === searchableModel.provider.provider
                      const capabilities =
                        searchableModel.model.capabilities || []

                      return (
                        <div
                          key={`fav-${searchableModel.value}`}
                          title={searchableModel.model.id}
                          onClick={() => handleSelect(searchableModel)}
                          className={cn(
                            'mx-1 mb-1 px-2 py-1.5 rounded-sm cursor-pointer flex items-center gap-2 transition-all duration-200',
                            'hover:bg-secondary/40',
                            isSelected && 'bg-secondary/50'
                          )}
                        >
                          <div className="flex items-center gap-1 flex-1 min-w-0">
                            <div className="shrink-0 -ml-1">
                              <ProvidersAvatar
                                provider={searchableModel.provider}
                              />
                            </div>
                            <span className="text-sm truncate">
                              {getModelDisplayName(searchableModel.model)}
                            </span>
                            {searchableModel.model.source && (
                              <ModelSourceBadge
                                source={searchableModel.model.source}
                                className="shrink-0"
                              />
                            )}
                            {searchableModel.model.missing && (
                              <MissingModelBadge
                                source={searchableModel.model.source}
                                className="shrink-0"
                              />
                            )}
                            <div className="flex-1"></div>
                            {capabilities.length > 0 && (
                              <div className="shrink-0 -mr-1.5">
                                <Capabilities capabilities={capabilities} />
                              </div>
                            )}
                          </div>
                        </div>
                      )
                    })}
                  </div>
                )}

                {/* Divider between favorites and regular providers */}
                {favoriteItems.length > 0 && (
                  <div className="border-b mx-2"></div>
                )}

                {/* Regular provider sections */}
                {Object.entries(groupedItems).map(([providerKey, models]) => {
                  const providerInfo = providers.find(
                    (p) => p.provider === providerKey
                  )

                  if (!providerInfo) return null

                  return (
                    <div
                      key={providerKey}
                      className="bg-secondary/30 first:mt-0 rounded-sm my-1.5 mx-1.5 first:mb-0 py-1"
                    >
                      {/* Provider header */}
                      <div className="flex items-center justify-between px-2 py-1">
                        <div className="flex items-center gap-1.5">
                          <ProvidersAvatar provider={providerInfo} />
                          <span className="capitalize text-sm font-medium text-muted-foreground">
                            {getProviderTitle(providerInfo.provider)}
                          </span>
                        </div>

                        <div
                          className="size-6 cursor-pointer flex items-center justify-center rounded-sm bg-secondary-foreground/8 transition-all duration-200 ease-in-out"
                          onClick={(e) => {
                            e.stopPropagation()
                            navigate({
                              to: route.settings.providers,
                              params: { providerName: providerInfo.provider },
                            })
                            setOpen(false)
                          }}
                        >
                          <IconSettings
                            size={16}
                            className="text-muted-foreground"
                          />
                        </div>
                      </div>

                      {/* Models for this provider */}
                      {models.length === 0 ? (
                        // Show message when provider has no available models
                        <></>
                      ) : (
                        models.map((searchableModel) => {
                          const isSelected =
                            selectedModel?.id === searchableModel.model.id &&
                            selectedProvider ===
                              searchableModel.provider.provider
                          const capabilities =
                            searchableModel.model.capabilities || []

                          return (
                            <div
                              key={searchableModel.value}
                              title={searchableModel.model.id}
                              onClick={() => handleSelect(searchableModel)}
                              className={cn(
                                'mx-1 mb-1 px-2 py-1.5 rounded-sm cursor-pointer flex items-center gap-2 transition-all duration-200',
                                'hover:bg-secondary/40',
                                isSelected &&
                                  'bg-secondary/60 hover:bg-secondary/60'
                              )}
                            >
                              <div className="flex items-center gap-2 flex-1 min-w-0">
                                <span
                                  className="text-sm truncate"
                                  title={searchableModel.model.id}
                                >
                                  {getModelDisplayName(searchableModel.model)}
                                </span>
                                {searchableModel.model.source && (
                                  <ModelSourceBadge
                                    source={searchableModel.model.source}
                                    className="shrink-0"
                                  />
                                )}
                                {searchableModel.model.missing && (
                                  <MissingModelBadge
                                    source={searchableModel.model.source}
                                    className="shrink-0"
                                  />
                                )}
                                <div className="flex-1"></div>
                                {capabilities.length > 0 && (
                                  <div className="shrink-0 -mr-1.5">
                                    <Capabilities capabilities={capabilities} />
                                  </div>
                                )}
                              </div>
                            </div>
                          )
                        })
                      )}
                    </div>
                  )
                })}
              </div>
            )}
          </div>

          {/* Download CTA — shortcut into the Hub so users can grab a local
              model without leaving the selector first. */}
          <div className="border-t p-1.5 mt-auto">
            <button
              type="button"
              onClick={onDownloadModel}
              className="w-full flex items-center gap-2 px-2 py-1.5 rounded-sm cursor-pointer text-sm text-muted-foreground transition-colors duration-200 hover:bg-secondary/40 hover:text-foreground"
            >
              <IconDownload size={16} className="shrink-0" />
              <span>{t('common:downloadModel')}</span>
            </button>
          </div>
        </div>
      </PopoverContent>
    </Popover>
  )
})

export default DropdownModelProvider