src/components/model-selector.tsx

"use client"

import { useOllama } from "@/contexts/ollama-context"
import { AI_MODELS, OLLAMA_MODELS, isOllamaModel } from "@/types/prompt"
import { Badge } from "@/components/ui/badge"
import { Button } from "@/components/ui/button"
import {
    Cloud,
    Server,
    ChevronDown,
    Check
} from "lucide-react"
import { useState, useRef, useEffect } from "react"
import { cn } from "@/lib/utils"

interface ModelSelectorProps {
    value: string
    onChange: (model: string) => void
    disabled?: boolean
    showOllama?: boolean
    compact?: boolean
}

export function ModelSelector({
    value,
    onChange,
    disabled = false,
    showOllama = true,
    compact = false
}: ModelSelectorProps) {
    const { settings } = useOllama()
    const [open, setOpen] = useState(false)
    const dropdownRef = useRef<HTMLDivElement>(null)

    // Close on click outside
    useEffect(() => {
        const handleClickOutside = (event: MouseEvent) => {
            if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) {
                setOpen(false)
            }
        }
        document.addEventListener("mousedown", handleClickOutside)
        return () => document.removeEventListener("mousedown", handleClickOutside)
    }, [])

    // Get current model info
    const isOllama = isOllamaModel(value)
    const currentModel = isOllama
        ? (OLLAMA_MODELS as Record<string, any>)[value] || { name: value }
        : (AI_MODELS as Record<string, any>)[value] || { name: value }

    // Build cloud models list
    const cloudModels = Object.entries(AI_MODELS).map(([id, config]) => ({
        id,
        ...config,
        isOllama: false,
    }))

    // Build Ollama models list (from detected + predefined)
    const ollamaModels = settings.enabled
        ? settings.availableModels.map(name => ({
            id: name,
            name: name,
            description: "Local Ollama model",
            isOllama: true,
        }))
        : []

    return (
        <div className="relative" ref={dropdownRef}>
            <Button
                variant="outline"
                size={compact ? "sm" : "default"}
                onClick={() => !disabled && setOpen(!open)}
                disabled={disabled}
                className={cn(
                    "justify-between gap-2",
                    compact ? "w-auto" : "w-full"
                )}
            >
                <div className="flex items-center gap-2">
                    {isOllama ? (
                        <Server className="h-4 w-4 text-green-500" />
                    ) : (
                        <Cloud className="h-4 w-4 text-blue-500" />
                    )}
                    <span className="truncate">{currentModel.name || value}</span>
                </div>
                <ChevronDown className={cn("h-4 w-4 transition-transform", open && "rotate-180")} />
            </Button>

            {open && (
                <div className="absolute z-50 mt-1 w-72 max-h-80 overflow-auto rounded-lg border bg-popover shadow-lg">
                    {/* Cloud Models */}
                    <div className="p-2 border-b">
                        <p className="text-xs font-medium text-muted-foreground px-2 py-1 flex items-center gap-1">
                            <Cloud className="h-3 w-3" />
                            Cloud Models
                        </p>
                        {cloudModels.map((model) => (
                            <button
                                key={model.id}
                                onClick={() => {
                                    onChange(model.id)
                                    setOpen(false)
                                }}
                                className={cn(
                                    "w-full px-2 py-2 text-left rounded-md hover:bg-muted flex items-center justify-between",
                                    value === model.id && "bg-muted"
                                )}
                            >
                                <div>
                                    <p className="text-sm font-medium">{model.name}</p>
                                    <p className="text-xs text-muted-foreground">{model.description}</p>
                                </div>
                                {value === model.id && <Check className="h-4 w-4 text-primary" />}
                            </button>
                        ))}
                    </div>

                    {/* Ollama Models */}
                    {showOllama && (
                        <div className="p-2">
                            <p className="text-xs font-medium text-muted-foreground px-2 py-1 flex items-center gap-1">
                                <Server className="h-3 w-3" />
                                Ollama Models
                                {!settings.enabled && (
                                    <Badge variant="outline" className="ml-1 text-xs">
                                        Not connected
                                    </Badge>
                                )}
                            </p>
                            {ollamaModels.length > 0 ? (
                                ollamaModels.map((model) => (
                                    <button
                                        key={model.id}
                                        onClick={() => {
                                            onChange(model.id)
                                            setOpen(false)
                                        }}
                                        className={cn(
                                            "w-full px-2 py-2 text-left rounded-md hover:bg-muted flex items-center justify-between",
                                            value === model.id && "bg-muted"
                                        )}
                                    >
                                        <div>
                                            <p className="text-sm font-medium">{model.name}</p>
                                            <p className="text-xs text-muted-foreground">{model.description}</p>
                                        </div>
                                        {value === model.id && <Check className="h-4 w-4 text-green-500" />}
                                    </button>
                                ))
                            ) : (
                                <p className="px-2 py-2 text-xs text-muted-foreground">
                                    {settings.enabled
                                        ? "No models found. Run: ollama pull llama3.2"
                                        : "Click 'Ollama' in header to connect"
                                    }
                                </p>
                            )}
                        </div>
                    )}
                </div>
            )}
        </div>
    )
}

// Simple inline selector for compact spaces
export function ModelBadge({ model }: { model: string }) {
    const isOllama = isOllamaModel(model)
    const modelInfo = isOllama
        ? (OLLAMA_MODELS as Record<string, any>)[model]
        : (AI_MODELS as Record<string, any>)[model]

    return (
        <Badge
            variant="outline"
            className={cn(
                "gap-1",
                isOllama ? "border-green-500/50" : "border-blue-500/50"
            )}
        >
            {isOllama ? (
                <Server className="h-3 w-3 text-green-500" />
            ) : (
                <Cloud className="h-3 w-3 text-blue-500" />
            )}
            {modelInfo?.name || model}
        </Badge>
    )
}