frontend/src/pages/Models.tsx

import { useState, useEffect } from 'react'
import { Card, CardHeader, CardTitle, CardContent } from '@/components/ui/card'
import { Button } from '@/components/ui/button'
import { Badge } from '@/components/ui/badge'
import { 
  BookOpen, 
  Brain, 
  Zap, 
  Download, 
  Trash2, 
  Loader2,
  Info,
  CheckCircle,
  Cloud,
  HardDrive
} from 'lucide-react'

interface ModelInfo {
  model_name: string
  name: string
  supports_thinking: boolean
  description: string
  size_gb: string
  is_loaded: boolean
  type: 'local' | 'api'
}

interface ModelsResponse {
  models: ModelInfo[]
  current_model: string
}

export function Models() {
  const [models, setModels] = useState<ModelInfo[]>([])
  const [loading, setLoading] = useState(true)
  const [modelLoading, setModelLoading] = useState<string | null>(null)

  useEffect(() => {
    fetchModels()
  }, [])

  const fetchModels = async () => {
    try {
      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''
      const res = await fetch(`${baseUrl}/models`)
      if (res.ok) {
        const data: ModelsResponse = await res.json()
        setModels(data.models)
      }
    } catch (err) {
      console.error('Failed to fetch models:', err)
    } finally {
      setLoading(false)
    }
  }

  const handleLoadModel = async (modelName: string) => {
    setModelLoading(modelName)
    try {
      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''
      const res = await fetch(`${baseUrl}/load-model`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ model_name: modelName }),
      })
      
      if (res.ok) {
        await fetchModels()
      }
    } catch (err) {
      console.error('Failed to load model:', err)
    } finally {
      setModelLoading(null)
    }
  }

  const handleUnloadModel = async (modelName: string) => {
    try {
      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''
      const res = await fetch(`${baseUrl}/unload-model`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ model_name: modelName }),
      })
      
      if (res.ok) {
        await fetchModels()
      }
    } catch (err) {
      console.error('Failed to unload model:', err)
    }
  }

  if (loading) {
    return (
      <div className="min-h-screen bg-background">
        <div className="border-b">
          <div className="flex h-14 items-center px-6">
            <div className="flex items-center gap-2">
              <BookOpen className="h-5 w-5" />
              <h1 className="text-lg font-semibold">Model Catalog</h1>
            </div>
          </div>
        </div>
        <div className="flex items-center justify-center h-64">
          <Loader2 className="h-8 w-8 animate-spin" />
        </div>
      </div>
    )
  }

  return (
    <div className="min-h-screen bg-background">
      {/* Header */}
      <div className="border-b">
        <div className="flex h-14 items-center px-6">
          <div className="flex items-center gap-2">
            <BookOpen className="h-5 w-5" />
            <h1 className="text-lg font-semibold">Model Catalog</h1>
          </div>
          <div className="ml-auto">
            <Button variant="outline" size="sm" onClick={fetchModels}>
              Refresh
            </Button>
          </div>
        </div>
      </div>

      <div className="flex-1 p-6">
        <div className="max-w-4xl mx-auto space-y-6">
          
          {/* Info Card */}
          <Card className="bg-blue-50 border-blue-200">
            <CardContent className="pt-6">
              <div className="flex items-start gap-3">
                <Info className="h-5 w-5 text-blue-600 mt-0.5" />
                <div>
                  <h3 className="font-medium text-blue-900">Model Management</h3>
                  <p className="text-sm text-blue-700 mt-1">
                    Load models to use them in the playground. Models are cached locally for faster access.
                    Each model requires significant storage space and initial download time.
                  </p>
                </div>
              </div>
            </CardContent>
          </Card>

          {/* API Models Section */}
          <div>
            <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
              <Cloud className="h-5 w-5" />
              API Models
              <Badge variant="outline" className="text-xs">Cloud-Powered</Badge>
            </h2>
            <div className="grid grid-cols-1 md:grid-cols-2 gap-6 mb-8">
              {models.filter(m => m.type === 'api').map((model) => (
                <ModelCard 
                  key={model.model_name} 
                  model={model} 
                  modelLoading={modelLoading}
                  onLoad={handleLoadModel}
                  onUnload={handleUnloadModel}
                />
              ))}
            </div>
          </div>

          {/* Local Models Section */}
          <div>
            <h2 className="text-xl font-semibold mb-4 flex items-center gap-2">
              <HardDrive className="h-5 w-5" />
              Local Models
              <Badge variant="outline" className="text-xs">Privacy-First</Badge>
            </h2>
            <div className="grid grid-cols-1 md:grid-cols-2 gap-6">
              {models.filter(m => m.type === 'local').map((model) => (
                <ModelCard 
                  key={model.model_name} 
                  model={model} 
                  modelLoading={modelLoading}
                  onLoad={handleLoadModel}
                  onUnload={handleUnloadModel}
                />
              ))}
            </div>
          </div>

          {/* Stats Card */}
          <Card>
            <CardHeader>
              <CardTitle>Model Statistics</CardTitle>
            </CardHeader>
            <CardContent>
              <div className="grid grid-cols-1 md:grid-cols-3 gap-4">
                <div className="text-center">
                  <div className="text-2xl font-bold text-blue-600">{models.length}</div>
                  <div className="text-sm text-muted-foreground">Available Models</div>
                </div>
                <div className="text-center">
                  <div className="text-2xl font-bold text-green-600">
                    {models.filter(m => m.is_loaded).length}
                  </div>
                  <div className="text-sm text-muted-foreground">Loaded Models</div>
                </div>
                <div className="text-center">
                  <div className="text-2xl font-bold text-purple-600">
                    {models.filter(m => m.supports_thinking).length}
                  </div>
                  <div className="text-sm text-muted-foreground">Thinking Models</div>
                </div>
              </div>
            </CardContent>
          </Card>
        </div>
      </div>
    </div>
  )
}

// ModelCard component for reusability
interface ModelCardProps {
  model: ModelInfo
  modelLoading: string | null
  onLoad: (modelName: string) => void
  onUnload: (modelName: string) => void
}

function ModelCard({ model, modelLoading, onLoad, onUnload }: ModelCardProps) {
  const isApiModel = model.type === 'api'
  
  return (
    <Card className="relative">
      <CardHeader>
        <div className="flex items-start justify-between">
          <div className="flex items-center gap-3">
            {isApiModel ? (
              <Cloud className="h-6 w-6 text-blue-500" />
            ) : model.supports_thinking ? (
              <Brain className="h-6 w-6 text-blue-500" />
            ) : (
              <Zap className="h-6 w-6 text-green-500" />
            )}
            <div>
              <CardTitle className="text-lg">{model.name}</CardTitle>
              <div className="flex items-center gap-2 mt-1 flex-wrap">
                {isApiModel ? (
                  <Badge variant="default" className="bg-blue-600">
                    <Cloud className="h-3 w-3 mr-1" />
                    API Model
                  </Badge>
                ) : (
                  <Badge variant={model.supports_thinking ? "default" : "secondary"}>
                    <HardDrive className="h-3 w-3 mr-1" />
                    {model.supports_thinking ? "Thinking Model" : "Instruction Model"}
                  </Badge>
                )}
                {model.is_loaded && (
                  <Badge variant="outline" className="text-green-600 border-green-600">
                    <CheckCircle className="h-3 w-3 mr-1" />
                    {isApiModel ? "Ready" : "Loaded"}
                  </Badge>
                )}
              </div>
            </div>
          </div>
        </div>
      </CardHeader>
      <CardContent className="space-y-4">
        <div>
          <p className="text-sm text-muted-foreground mb-2">{model.description}</p>
          <div className="flex items-center gap-4 text-xs text-muted-foreground">
            <span>Size: {model.size_gb}</span>
            {!isApiModel && <span>Format: Safetensors</span>}
            {isApiModel && <span>Type: Cloud API</span>}
          </div>
        </div>

        <div className="space-y-2">
          <h4 className="text-sm font-medium">Capabilities</h4>
          <div className="flex flex-wrap gap-2">
            <Badge variant="outline" className="text-xs">Text Generation</Badge>
            <Badge variant="outline" className="text-xs">Conversation</Badge>
            <Badge variant="outline" className="text-xs">Code</Badge>
            {model.supports_thinking && (
              <Badge variant="outline" className="text-xs">Reasoning</Badge>
            )}
            {isApiModel && model.model_name.includes('vl') && (
              <Badge variant="outline" className="text-xs">Vision</Badge>
            )}
          </div>
        </div>

        <div className="pt-2 border-t">
          {model.is_loaded ? (
            <div className="flex gap-2">
              {!isApiModel && (
                <Button
                  variant="outline"
                  size="sm"
                  onClick={() => onUnload(model.model_name)}
                  className="flex-1"
                >
                  <Trash2 className="h-4 w-4 mr-2" />
                  Unload
                </Button>
              )}
              <Button size="sm" className="flex-1" asChild>
                <a href="/playground">Use in Playground</a>
              </Button>
            </div>
          ) : (
            <Button
              onClick={() => onLoad(model.model_name)}
              disabled={modelLoading === model.model_name}
              className="w-full"
              size="sm"
            >
              {modelLoading === model.model_name ? (
                <>
                  <Loader2 className="h-4 w-4 mr-2 animate-spin" />
                  {isApiModel ? "Connecting..." : "Loading..."}
                </>
              ) : (
                <>
                  {isApiModel ? (
                    <Cloud className="h-4 w-4 mr-2" />
                  ) : (
                    <Download className="h-4 w-4 mr-2" />
                  )}
                  {isApiModel ? "Connect" : "Load Model"}
                </>
              )}
            </Button>
          )}
        </div>
      </CardContent>
    </Card>
  )
}