components/model-selector.tsx

'use client';

import { useState } from 'react';
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
import { Badge } from '@/components/ui/badge';
import { Settings, Cpu, Zap, Brain } from 'lucide-react';

interface ModelSelectorProps {
  selectedModel: string;
  onModelChange: (model: string) => void;
  models: string[];
}

const modelIcons: Record<string, React.ReactNode> = {
  'deepseek-ai': <Brain className="w-3 h-3" />,
  'Qwen': <Cpu className="w-3 h-3" />,
  'moonshotai': <Zap className="w-3 h-3" />,
  'zai-org': <Settings className="w-3 h-3" />,
  'MiniMaxAI': <Brain className="w-3 h-3" />,
  'meta-llama': <Cpu className="w-3 h-3" />,
  'google': <Brain className="w-3 h-3" />,
};

const getProviderFromModel = (modelName: string): string => {
  if (modelName.includes('deepseek-ai') || modelName.includes('Qwen') || 
      modelName.includes('moonshotai') || modelName.includes('zai-org') || 
      modelName.includes('MiniMaxAI') || modelName.includes('meta-llama') || 
      modelName.includes('google')) {
    return 'Hugging Face';
  }
  return 'Unknown';
};

const getModelCategory = (modelName: string): string => {
  if (modelName.includes('DeepSeek')) return 'Reasoning';
  if (modelName.includes('Coder')) return 'Code';
  if (modelName.includes('VL')) return 'Vision';
  if (modelName.includes('Kimi') && modelName.includes('Thinking')) return 'Reasoning';
  if (modelName.includes('Kimi')) return 'General';
  if (modelName.includes('GLM')) return 'General';
  return 'Language';
};

const getDisplayName = (modelName: string): string => {
  const parts = modelName.split('/');
  const model = parts[1] || modelName;
  
  // Clean up model names for better display
  const cleanName = model
    .replace('-0324', '')
    .replace('-0528', '')
    .replace('-0905', '')
    .replace('-Instruct', '')
    .replace('-A22B', '')
    .replace('-A35B', '')
    .replace('-VL-7B', '')
    .replace('-Exp', '')
    .replace('-Distill-', '-')
    .replace('-Terminus', '')
    .replace('DeepSeek-V3-0324', 'DeepSeek V3')
    .replace('DeepSeek-R1-0528', 'DeepSeek R1')
    .replace('DeepSeek-V3.1', 'DeepSeek V3.1')
    .replace('DeepSeek-V3.1-Terminus', 'DeepSeek V3.1 Terminus')
    .replace('DeepSeek-V3.2-Exp', 'DeepSeek V3.2 Exp')
    .replace('Qwen3-Coder-480B-A35B-Instruct', 'Qwen3 Coder 480B')
    .replace('Qwen2.5-VL-7B-Instruct', 'Qwen2.5 VL 7B')
    .replace('Kimi-K2-Instruct', 'Kimi K2')
    .replace('Kimi-K2-Instruct-0905', 'Kimi K2 0905')
    .replace('Kimi-K2-Thinking', 'Kimi K2 Thinking')
    .replace('GLM-4.6', 'GLM-4.6')
    .replace('MiniMax-M2', 'MiniMax M2')
    .replace('Llama-3.1-8B-Instruct', 'Llama 3.1 8B')
    .replace('Llama-3.1-70B-Instruct', 'Llama 3.1 70B')
    .replace('Llama-3.3-70B-Instruct', 'Llama 3.3 70B')
    .replace('Llama-4-Scout-17B-16E-Instruct', 'Llama 4 Scout 17B')
    .replace('gemma-3-27b-it', 'Gemma 3 27B');
  
  return cleanName;
};

export function ModelSelector({ selectedModel, onModelChange, models }: ModelSelectorProps) {
  const [open, setOpen] = useState(false);

  const getIconForModel = (modelName: string) => {
    for (const [key, icon] of Object.entries(modelIcons)) {
      if (modelName.includes(key)) {
        return icon;
      }
    }
    return <Brain className="w-3 h-3" />;
  };

  const groupModelsByProvider = () => {
    const groups: Record<string, string[]> = {};
    
    models.forEach(model => {
      const provider = getProviderFromModel(model);
      if (!groups[provider]) {
        groups[provider] = [];
      }
      groups[provider].push(model);
    });
    
    return groups;
  };

  const groupedModels = groupModelsByProvider();

  return (
    <Select
      value={selectedModel}
      onValueChange={onModelChange}
      onOpenChange={setOpen}
    >
      <SelectTrigger className="w-[280px]">
        <div className="flex items-center gap-2 truncate">
          {getIconForModel(selectedModel)}
          <div className="flex-1 min-w-0">
            <div className="flex items-center gap-2">
              <span className="font-medium truncate">
                {getDisplayName(selectedModel)}
              </span>
              <Badge variant="outline" className="text-xs">
                {getModelCategory(selectedModel)}
              </Badge>
            </div>
            <div className="text-xs text-muted-foreground truncate">
              {getProviderFromModel(selectedModel)}
            </div>
          </div>
        </div>
      </SelectTrigger>
      
      <SelectContent>
        {Object.entries(groupedModels).map(([provider, providerModels]) => (
          <div key={provider}>
            <div className="px-2 py-1.5 text-xs font-semibold text-muted-foreground border-b">
              {provider}
            </div>
            {providerModels.map((model) => (
              <SelectItem key={model} value={model}>
                <div className="flex items-center gap-2 w-full">
                  {getIconForModel(model)}
                  <div className="flex-1 min-w-0">
                    <div className="flex items-center gap-2">
                      <span className="font-medium truncate">
                        {getDisplayName(model)}
                      </span>
                      <Badge variant="outline" className="text-xs flex-shrink-0">
                        {getModelCategory(model)}
                      </Badge>
                    </div>
                    <div className="text-xs text-muted-foreground truncate">
                      {model}
                    </div>
                  </div>
                </div>
              </SelectItem>
            ))}
          </div>
        ))}
      </SelectContent>
    </Select>
  );
}