ForgeLM/tauri-app/src/components/ModelSelectorPill.tsx

import { useState, useRef, useEffect } from "react";
import { Bot, Brain, ChevronDown, Cpu, Gem, Loader2, Search, Sparkles, Wind } from "lucide-react";
import type { ModelInfo, ServerStatus } from "../lib/types";
import type { RuntimeModelFile } from "../features/runtime/types";
import { getDisplayName } from "../lib/modelAlias";

interface ModelSelectorPillProps {
  models: ModelInfo[];
  localModels: RuntimeModelFile[];
  selectedModel: string;
  onModelChange: (id: string) => void;
  onModelSwap?: (path: string) => void;
  isSwapping?: boolean;
  status: ServerStatus;
}

export default function ModelSelectorPill({
  models,
  localModels,
  selectedModel,
  onModelChange,
  onModelSwap,
  isSwapping,
  status,
}: ModelSelectorPillProps) {
  const [dropdownOpen, setDropdownOpen] = useState(false);
  const [openDirection, setOpenDirection] = useState<"up" | "down">("down");
  const dropdownRef = useRef<HTMLDivElement>(null);

  useEffect(() => {
    if (!dropdownOpen) return;
    const handleClickOutside = (e: MouseEvent) => {
      if (dropdownRef.current && !dropdownRef.current.contains(e.target as Node)) {
        setDropdownOpen(false);
      }
    };
    const el = dropdownRef.current;
    if (el) {
      const rect = el.getBoundingClientRect();
      const estimatedMenuHeight = 260;
      const spaceBelow = window.innerHeight - rect.bottom;
      const spaceAbove = rect.top;
      setOpenDirection(spaceBelow >= estimatedMenuHeight || spaceBelow >= spaceAbove ? "down" : "up");
    }
    document.addEventListener("mousedown", handleClickOutside);
    return () => document.removeEventListener("mousedown", handleClickOutside);
  }, [dropdownOpen]);

  const handleSelectModel = (modelId: string) => {
    setDropdownOpen(false);
    onModelChange(modelId);
    // Only call onModelSwap for local models (which have file paths)
    // API models don't have paths and shouldn't trigger a model swap
    const isLocalModel = localModels.some((m) => m.path === modelId);
    if (onModelSwap && isLocalModel) {
      onModelSwap(modelId);
    }
  };

  const getModelIcon = (label: string) => {
    const l = label.toLowerCase();
    if (l.includes("qwen")) return <Brain size={13} strokeWidth={2} />;
    if (l.includes("deepseek")) return <Search size={13} strokeWidth={2} />;
    if (l.includes("mistral")) return <Wind size={13} strokeWidth={2} />;
    if (l.includes("gemma") || l.includes("gemini")) return <Gem size={13} strokeWidth={2} />;
    if (l.includes("phi")) return <Sparkles size={13} strokeWidth={2} />;
    if (l.includes("llama")) return <Bot size={13} strokeWidth={2} />;
    return <Cpu size={13} strokeWidth={2} />;
  };

  // Unified model list: merge API models + local scanned files
  const allAvailableModels = [
    ...models.map((m) => ({ id: m.id, label: getDisplayName(m.id) })),
    ...localModels.map((m) => ({ id: m.path, label: getDisplayName(m.path) })),
  ];
  // Deduplicate by label (display name) - same name means same model
  const seenLabels = new Set<string>();
  const dedupedModels = allAvailableModels.filter((m) => {
    if (seenLabels.has(m.label)) return false;
    seenLabels.add(m.label);
    return true;
  });

  // Show all available models in the selector
  const selectedLabel = getDisplayName(selectedModel) || "No model";

  return (
    <div className="model-selector-pill-wrapper" ref={dropdownRef}>
      <button
        type="button"
        className="model-selector-pill"
        onClick={() => setDropdownOpen(!dropdownOpen)}
        disabled={dedupedModels.length === 0 || isSwapping}
        title={dedupedModels.length === 0 ? "Start server to load models" : ""}
        aria-label="Choose active model"
        aria-expanded={dropdownOpen}
      >
        <span className={`model-status-dot status-${status}`} aria-hidden="true" />
        {isSwapping ? (
          <Loader2 size={12} strokeWidth={2} className="animate-spin" />
        ) : (
          <>
            <span className="model-selector-icon" aria-hidden="true">{getModelIcon(selectedLabel)}</span>
            <span className="model-selector-label">{selectedLabel}</span>
          </>
        )}
        <ChevronDown size={12} strokeWidth={2} className="model-selector-chevron" />
      </button>
      {dropdownOpen && (
        <ul className={`model-selector-menu model-selector-menu--${openDirection}`} role="listbox">
          {dedupedModels.map((m) => (
            <li
              key={m.id}
              role="option"
              aria-selected={m.id === selectedModel}
              className={`model-selector-item${m.id === selectedModel ? " model-selector-item--selected" : ""}`}
              onClick={() => handleSelectModel(m.id)}
            >
              <span className="model-selector-item-icon" aria-hidden="true">{getModelIcon(m.label)}</span>
              <span className="model-selector-item-label">{m.label}</span>
            </li>
          ))}
        </ul>
      )}
    </div>
  );
}