frontend/src/pages/Playground.tsx

import { useState, useEffect } from 'react'
import { Button } from '@/components/ui/button'
import { Card, CardHeader, CardTitle, CardContent } from '@/components/ui/card'
import { Slider } from '@/components/ui/slider'
import { Label } from '@/components/ui/label'
import { Badge } from '@/components/ui/badge'
import { 
  Select,
  SelectContent,
  SelectGroup,
  SelectItem,
  SelectLabel,
  SelectTrigger,
  SelectValue,
} from '@/components/ui/select'

import { 
  Collapsible, 
  CollapsibleContent, 
  CollapsibleTrigger 
} from '@/components/ui/collapsible'
import {
  AlertDialog,
  AlertDialogAction,
  AlertDialogCancel,
  AlertDialogContent,
  AlertDialogDescription,
  AlertDialogFooter,
  AlertDialogHeader,
  AlertDialogTitle,
} from '@/components/ui/alert-dialog'
import { Chat } from '@/components/ui/chat'
import { useChat } from '@/hooks/useChat'
import { 
  Brain, 
  Zap, 
  ChevronDown, 
  MessageSquare, 
  RotateCcw,
  Code,
  Upload,
  Share,
  History,
  Settings,
  PanelLeftOpen,
  PanelLeftClose,
  Cloud,
  BookOpen,
  Download,
  AlertTriangle,
  Plus,
  Trash2
} 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 Playground() {
  // Chat functionality
  const {
    sessions,
    currentSession,
    currentSessionId,
    createNewSession,
    selectSession,
    deleteSession,

    messages,
    input,
    setInput,
    sendMessage,
    stopGeneration,
    isLoading,
    selectedModel,
    setSelectedModel,
    systemPrompt,
    setSystemPrompt,
    temperature,
    setTemperature,
    maxTokens,
    setMaxTokens
  } = useChat()

  // UI state
  const [showSessions, setShowSessions] = useState(false)
  const [isSystemPromptOpen, setIsSystemPromptOpen] = useState(false)
  const [autoLoadingModel, setAutoLoadingModel] = useState<string | null>(null)
  const [showLoadConfirm, setShowLoadConfirm] = useState(false)
  const [pendingModelToLoad, setPendingModelToLoad] = useState<ModelInfo | null>(null)
  
  // Model management state
  const [models, setModels] = useState<ModelInfo[]>([])

  // Preset system prompts
  const systemPromptPresets = [
    {
      name: "Default Assistant",
      prompt: "You are a helpful, harmless, and honest AI assistant. Provide clear, accurate, and well-structured responses."
    },
    {
      name: "Code Expert",
      prompt: "You are an expert software developer. Provide clean, efficient code with clear explanations. Always follow best practices and include comments where helpful."
    },
    {
      name: "Technical Writer",
      prompt: "You are a technical writer. Create clear, comprehensive documentation and explanations. Use proper formatting and structure your responses logically."
    },
    {
      name: "Creative Writer",
      prompt: "You are a creative writer. Use vivid language, engaging storytelling, and imaginative descriptions. Be expressive and artistic in your responses."
    },
    {
      name: "Research Assistant",
      prompt: "You are a research assistant. Provide detailed, well-researched responses with clear reasoning. Cite sources when relevant and present information objectively."
    },
    {
      name: "Teacher",
      prompt: "You are an experienced teacher. Explain concepts clearly, use examples, and break down complex topics into understandable parts. Be encouraging and patient."
    }
  ]

  // Sample prompts for quick start
  const samplePrompts = [
    {
      title: "Marketing Slogan",
      description: "Create a catchy marketing slogan for a new eco-friendly product.",
      prompt: "Create a catchy marketing slogan for a new eco-friendly water bottle that keeps drinks cold for 24 hours. The target audience is environmentally conscious millennials and Gen Z consumers."
    },
    {
      title: "Creative Storytelling",
      description: "Write a short story about a time traveler.",
      prompt: "Write a 300-word short story about a time traveler who accidentally changes a major historical event while trying to observe ancient Rome."
    },
    {
      title: "Technical Explanation",
      description: "Explain a complex technical concept simply.",
      prompt: "Explain how blockchain technology works in simple terms that a 12-year-old could understand, using analogies and examples."
    },
    {
      title: "Code Generation",
      description: "Generate code with explanations.",
      prompt: "Write a Python function that takes a list of numbers and returns the second largest number. Include error handling and detailed comments explaining each step."
    }
  ]

  // Load available models on startup
  useEffect(() => {
    fetchModels()
  }, [])

  // Update selected model when models change
  useEffect(() => {
    // Only reset if the selected model no longer exists in the models list
    if (selectedModel && !models.find(m => m.model_name === selectedModel)) {
      const firstModel = models[0]
      if (firstModel) {
        setSelectedModel(firstModel.model_name)
      }
    }
  }, [models, selectedModel, setSelectedModel])

  // Auto-load/unload local models when selection changes
  useEffect(() => {
    const handleModelChange = async () => {
      if (!selectedModel || !models.length) return

      const selectedModelInfo = models.find(m => m.model_name === selectedModel)
      if (!selectedModelInfo) return

      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''

      // If selected model is a local model and not loaded, show confirmation
      if (selectedModelInfo.type === 'local' && !selectedModelInfo.is_loaded) {
        setPendingModelToLoad(selectedModelInfo)
        setShowLoadConfirm(true)
        return // Don't auto-load, wait for user confirmation
      }

      // Unload other local models that are loaded but not selected
      const loadedLocalModels = models.filter(m => 
        m.type === 'local' && 
        m.is_loaded && 
        m.model_name !== selectedModel
      )

      for (const model of loadedLocalModels) {
        try {
          const response = await fetch(`${baseUrl}/unload-model`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ model_name: model.model_name })
          })
          
          if (response.ok) {
            console.log(`✅ Auto-unloaded local model: ${model.model_name}`)
          }
        } catch (error) {
          console.error(`Error auto-unloading model ${model.model_name}:`, error)
        }
      }
      
      // Refresh models after any unloading
      if (loadedLocalModels.length > 0) {
        fetchModels()
      }
    }

    handleModelChange()
  }, [selectedModel, models])

  const handleLoadModelConfirm = async () => {
    if (!pendingModelToLoad) return
    
    setShowLoadConfirm(false)
    setAutoLoadingModel(pendingModelToLoad.model_name)
    
    try {
      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''
      const response = await fetch(`${baseUrl}/load-model`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ model_name: pendingModelToLoad.model_name })
      })
      
      if (response.ok) {
        console.log(`✅ User confirmed and loaded: ${pendingModelToLoad.model_name}`)
        fetchModels() // Refresh model states
      } else {
        console.error(`❌ Failed to load model: ${pendingModelToLoad.model_name}`)
        // Revert to an API model if load failed
        const apiModel = models.find(m => m.type === 'api')
        if (apiModel) {
          setSelectedModel(apiModel.model_name)
        }
      }
    } catch (error) {
      console.error('Error loading model:', error)
      // Revert to an API model if error
      const apiModel = models.find(m => m.type === 'api')
      if (apiModel) {
        setSelectedModel(apiModel.model_name)
      }
    } finally {
      setAutoLoadingModel(null)
      setPendingModelToLoad(null)
    }
  }

  const handleLoadModelCancel = () => {
    setShowLoadConfirm(false)
    setPendingModelToLoad(null)
    
    // Revert to an API model
    const apiModel = models.find(m => m.type === 'api')
    if (apiModel) {
      setSelectedModel(apiModel.model_name)
    }
  }

  // Cleanup: unload all local models when component unmounts or user leaves
  useEffect(() => {
    const handlePageUnload = async () => {
      const baseUrl = window.location.hostname === 'localhost' ? `${window.location.protocol}//${window.location.host}` : ''
      const loadedLocalModels = models.filter(m => m.type === 'local' && m.is_loaded)
      
      for (const model of loadedLocalModels) {
        try {
          await fetch(`${baseUrl}/unload-model`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify({ model_name: model.model_name })
          })
          console.log(`✅ Cleanup: unloaded ${model.model_name}`)
        } catch (error) {
          console.error(`Error cleaning up model ${model.model_name}:`, error)
        }
      }
    }

    // Cleanup on component unmount
    return () => {
      handlePageUnload()
    }
  }, [models])

  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)
        
        // Set selected model to current model if available, otherwise first API model
        if (data.current_model && selectedModel !== data.current_model) {
          setSelectedModel(data.current_model)
        } else if (!selectedModel && data.models.length > 0) {
          // Prefer API models as default
          const apiModel = data.models.find(m => m.type === 'api')
          const defaultModel = apiModel || data.models[0]
          setSelectedModel(defaultModel.model_name)
        }
      }
    } catch (err) {
      console.error('Failed to fetch models:', err)
    }
  }







  const handleSamplePromptClick = (samplePrompt: string) => {
    setInput(samplePrompt)
  }

  return (
    <div className="min-h-screen bg-background flex">
      {/* Chat Sessions Sidebar */}
      <div className={`
        ${showSessions ? 'translate-x-0' : '-translate-x-full'}
        fixed inset-y-0 left-0 z-50 w-80 bg-background border-r transition-transform duration-300 ease-in-out
        lg:translate-x-0 lg:static lg:inset-0
      `}>
        <div className="p-4 space-y-4">
          <div className="flex items-center justify-between">
            <h2 className="font-semibold">Chat Sessions</h2>
            <Button onClick={createNewSession} size="sm">
              <Plus className="h-4 w-4 mr-1" />
              New
            </Button>
          </div>
          <div className="space-y-2">
            {sessions.map((session) => (
              <Card 
                key={session.id}
                className={`p-3 cursor-pointer transition-colors hover:bg-accent ${
                  currentSessionId === session.id ? 'bg-accent border-primary' : ''
                }`}
                onClick={() => selectSession(session.id)}
              >
                <div className="flex items-center justify-between">
                  <span className="text-sm font-medium truncate">{session.title}</span>
                  <Button
                    size="sm"
                    variant="ghost"
                    onClick={(e) => {
                      e.stopPropagation()
                      deleteSession(session.id)
                    }}
                    className="h-6 w-6 p-0"
                  >
                    <Trash2 className="h-3 w-3" />
                  </Button>
                </div>
                <div className="text-xs text-muted-foreground">
                  {session.messages.length} messages
                </div>
              </Card>
            ))}
          </div>
        </div>
      </div>

      {/* Overlay for mobile */}
      {showSessions && (
        <div 
          className="fixed inset-0 z-40 bg-black/50 lg:hidden"
          onClick={() => setShowSessions(false)}
        />
      )}

      {/* Main Content */}
      <div className="flex-1 flex flex-col overflow-hidden">
        {/* Header */}
        <div className="border-b bg-background/95 backdrop-blur supports-[backdrop-filter]:bg-background/60">
          <div className="flex h-14 items-center px-6">
            <div className="flex items-center gap-2">
              <Button
                variant="ghost"
                size="sm"
                onClick={() => setShowSessions(!showSessions)}
                className="lg:hidden"
              >
                {showSessions ? <PanelLeftClose className="h-4 w-4" /> : <PanelLeftOpen className="h-4 w-4" />}
              </Button>
              <MessageSquare className="h-5 w-5" />
              <h1 className="text-lg font-semibold">Chat Playground</h1>
              {currentSession && (
                <Badge variant="outline" className="text-xs">
                  {currentSession.title.slice(0, 20)}...
                </Badge>
              )}
            </div>
            <div className="ml-auto flex items-center gap-2 overflow-x-auto">
              <Button
                variant="outline"
                size="sm"
                onClick={() => setShowSessions(!showSessions)}
                className="hidden lg:flex flex-shrink-0"
              >
                <History className="h-4 w-4 mr-2" />
                <span className="hidden sm:inline">Sessions</span>
              </Button>
              <Button variant="outline" size="sm" className="flex-shrink-0">
                <Code className="h-4 w-4 mr-2" />
                <span className="hidden sm:inline">View code</span>
                <span className="sm:hidden">Code</span>
              </Button>
              <Button variant="outline" size="sm" className="flex-shrink-0">
                <Upload className="h-4 w-4 mr-2" />
                <span className="hidden sm:inline">Import</span>
              </Button>
              <Button variant="outline" size="sm" className="flex-shrink-0">
                <Share className="h-4 w-4 mr-2" />
                <span className="hidden sm:inline">Export</span>
              </Button>
            </div>
          </div>
        </div>

        {/* Content Area */}
        <div className="flex-1 flex overflow-hidden">
          {/* Chat Area */}
          <div className="flex-1 flex flex-col">
            {/* Sample Prompts */}
            {messages.length === 0 && (
              <div className="p-6 border-b">
                <Card>
                  <CardHeader>
                    <CardTitle className="text-base">Start with a sample prompt</CardTitle>
                  </CardHeader>
                  <CardContent>
                    <div className="grid grid-cols-1 sm:grid-cols-2 xl:grid-cols-4 gap-4">
                      {samplePrompts.map((sample, index) => (
                        <Button
                          key={index}
                          variant="outline"
                          className="h-auto p-4 text-left justify-start min-w-0"
                          onClick={() => handleSamplePromptClick(sample.prompt)}
                          disabled={isLoading}
                        >
                          <div className="min-w-0">
                            <div className="font-medium text-sm mb-1 truncate">{sample.title}</div>
                            <div className="text-xs text-muted-foreground line-clamp-2">{sample.description}</div>
                          </div>
                        </Button>
                      ))}
                    </div>
                  </CardContent>
                </Card>
              </div>
            )}

            {/* Chat Messages and Input */}
                      <Chat
              messages={messages.map(msg => ({
                id: msg.id,
                role: msg.role as 'user' | 'assistant' | 'system',
                content: msg.content,
                createdAt: new Date(msg.timestamp)
              }))}
              input={input}
              handleInputChange={(e) => setInput(e.target.value)}
              handleSubmit={async (e) => {
                e.preventDefault()
                if (!selectedModel || !models.find(m => m.model_name === selectedModel)) return
                await sendMessage()
              }}
              isGenerating={isLoading}
              stop={stopGeneration}
              className="flex-1"
            />
          </div>

          {/* Settings Panel */}
          <div className="w-80 border-l bg-muted/30 overflow-y-auto">
            <div className="p-4 space-y-6">
              <div className="flex items-center gap-2">
                <Settings className="h-4 w-4" />
                <h2 className="font-semibold text-sm">Configuration</h2>
              </div>

              {/* Model Selection */}
              <Card>
                <CardHeader>
                  <CardTitle className="text-sm">Model Selection</CardTitle>
                </CardHeader>
                <CardContent className="space-y-3">
                  {/* Simple Model Dropdown */}
                  <div>
                    <Label className="text-xs font-medium mb-2">Active Model</Label>
                    <Select value={selectedModel || ""} onValueChange={setSelectedModel}>
                      <SelectTrigger className="w-full">
                        <SelectValue placeholder="Select a model...">
                          {selectedModel && (() => {
                            const model = models.find(m => m.model_name === selectedModel)
                            if (!model) return selectedModel
                            const isApiModel = model.type === 'api'
                            return (
                              <div className="flex items-center gap-2">
                                {isApiModel ? (
                                  <Cloud className="h-4 w-4 text-blue-500" />
                                ) : model.supports_thinking ? (
                                  <Brain className="h-4 w-4 text-purple-500" />
                                ) : (
                                  <Zap className="h-4 w-4 text-green-500" />
                                )}
                                <span className="truncate">{model.name}</span>
                                {autoLoadingModel === selectedModel ? (
                                  <Badge variant="outline" className="text-xs">
                                    Loading...
                                  </Badge>
                                ) : (
                                  <Badge variant="outline" className="text-xs">
                                    {isApiModel ? "API" : model.is_loaded ? "Loaded" : "Available"}
                                  </Badge>
                                )}
                              </div>
                            )
                          })()}
                        </SelectValue>
                      </SelectTrigger>
                      <SelectContent>
                        <SelectGroup>
                          <SelectLabel>🌐 API Models</SelectLabel>
                          {models.filter(m => m.type === 'api').map((model) => (
                            <SelectItem key={model.model_name} value={model.model_name}>
                              <div className="flex items-center gap-2">
                                <Cloud className="h-4 w-4 text-blue-500" />
                                <span>{model.name}</span>
                                <Badge variant="outline" className="text-xs bg-blue-50">API</Badge>
                              </div>
                            </SelectItem>
                          ))}
                        </SelectGroup>
                        <SelectGroup>
                          <SelectLabel>💻 Local Models</SelectLabel>
                          {models.filter(m => m.type === 'local').map((model) => (
                            <SelectItem key={model.model_name} value={model.model_name}>
                              <div className="flex items-center gap-2">
                                {model.supports_thinking ? (
                                  <Brain className="h-4 w-4 text-purple-500" />
                                ) : (
                                  <Zap className="h-4 w-4 text-green-500" />
                                )}
                                <span>{model.name}</span>
                                {autoLoadingModel === model.model_name ? (
                                  <Badge variant="outline" className="text-xs bg-yellow-50">Loading...</Badge>
                                ) : model.is_loaded ? (
                                  <Badge variant="outline" className="text-xs bg-green-50">Loaded</Badge>
                                ) : (
                                  <Badge variant="outline" className="text-xs bg-gray-50">Available</Badge>
                                )}
                              </div>
                            </SelectItem>
                          ))}
                        </SelectGroup>
                      </SelectContent>
                    </Select>
                  </div>

                  {/* Model Catalog Link */}
                  <div className="pt-2 border-t">
                    <Button variant="outline" size="sm" className="w-full" asChild>
                      <a href="/models" className="flex items-center gap-2">
                        <BookOpen className="h-4 w-4" />
                        View Model Catalog
                      </a>
                    </Button>
                  </div>
                </CardContent>
              </Card>

              {/* Parameters */}
              <Card>
                <CardHeader>
                  <CardTitle className="text-sm">Parameters</CardTitle>
                </CardHeader>
                <CardContent className="space-y-4">
                  {/* Temperature */}
                  <div>
                    <Label className="text-xs font-medium">
                      Temperature: {temperature.toFixed(2)}
                    </Label>
                    <Slider
                      value={[temperature]}
                      onValueChange={(value: number[]) => setTemperature(value[0])}
                      min={0}
                      max={2}
                      step={0.01}
                      className="mt-2"
                      disabled={isLoading}
                    />
                    <p className="text-xs text-muted-foreground mt-1">
                      Lower = more focused, Higher = more creative
                    </p>
                  </div>

                  {/* Max Tokens */}
                  <div>
                    <Label className="text-xs font-medium">
                      Max Tokens: {maxTokens}
                    </Label>
                    <Slider
                      value={[maxTokens]}
                      onValueChange={(value: number[]) => setMaxTokens(value[0])}
                      min={100}
                      max={4096}
                      step={100}
                      className="mt-2"
                      disabled={isLoading}
                    />
                  </div>
                </CardContent>
              </Card>

              {/* System Prompt */}
              <Card>
                <Collapsible 
                  open={isSystemPromptOpen} 
                  onOpenChange={setIsSystemPromptOpen}
                >
                  <CardHeader>
                    <CollapsibleTrigger asChild>
                      <Button variant="ghost" className="w-full justify-between p-0" disabled={isLoading}>
                        <div className="flex items-center gap-2">
                          <MessageSquare className="h-4 w-4" />
                          <span className="text-sm font-medium">System Prompt</span>
                          {systemPrompt && <Badge variant="secondary" className="text-xs">Custom</Badge>}
                        </div>
                        <ChevronDown className={`h-4 w-4 transition-transform ${isSystemPromptOpen ? 'transform rotate-180' : ''}`} />
                      </Button>
                    </CollapsibleTrigger>
                  </CardHeader>
                  <CollapsibleContent>
                    <CardContent className="space-y-3">
                      {/* Preset System Prompts */}
                      <div>
                        <Label className="text-xs font-medium text-muted-foreground">Quick Presets</Label>
                        <div className="grid grid-cols-1 gap-1 mt-1">
                          {systemPromptPresets.map((preset) => (
                            <Button
                              key={preset.name}
                              variant="outline"
                              size="sm"
                              className="h-auto p-2 text-xs justify-start"
                              onClick={() => setSystemPrompt(preset.prompt)}
                              disabled={isLoading}
                            >
                              {preset.name}
                            </Button>
                          ))}
                        </div>
                      </div>

                      {/* Custom System Prompt */}
                      <div>
                        <div className="flex items-center justify-between mb-2">
                          <Label htmlFor="system-prompt" className="text-xs font-medium">
                            Custom System Prompt
                          </Label>
                          {systemPrompt && (
                            <Button
                              variant="ghost"
                              size="sm"
                              onClick={() => setSystemPrompt('')}
                              className="h-6 px-2 text-xs"
                              disabled={isLoading}
                            >
                              <RotateCcw className="h-3 w-3 mr-1" />
                              Clear
                            </Button>
                          )}
                        </div>
                        <textarea
                          id="system-prompt"
                          value={systemPrompt}
                          onChange={(e: React.ChangeEvent<HTMLTextAreaElement>) => setSystemPrompt(e.target.value)}
                          placeholder="Enter custom system prompt to define how the model should behave..."
                          className="w-full min-h-[80px] text-xs p-2 border rounded-md bg-background"
                          disabled={isLoading}
                        />
                        <p className="text-xs text-muted-foreground mt-1">
                          System prompts define the model's role and behavior.
                        </p>
                      </div>
                    </CardContent>
                  </CollapsibleContent>
                </Collapsible>
              </Card>
            </div>
          </div>
        </div>
      </div>

      {/* Model Load Confirmation Dialog */}
      <AlertDialog open={showLoadConfirm} onOpenChange={setShowLoadConfirm}>
        <AlertDialogContent>
          <AlertDialogHeader>
            <AlertDialogTitle className="flex items-center gap-2">
              <Download className="h-5 w-5 text-blue-500" />
              Load Local Model
            </AlertDialogTitle>
            <AlertDialogDescription asChild>
              <div className="space-y-3">
                <p>
                  You're about to load <strong>{pendingModelToLoad?.name}</strong> locally.
                </p>
                
                <div className="bg-yellow-50 border border-yellow-200 rounded-lg p-3">
                  <div className="flex items-start gap-2">
                    <AlertTriangle className="h-4 w-4 text-yellow-600 mt-0.5" />
                    <div className="text-sm">
                      <p className="font-medium text-yellow-800">Resource Requirements:</p>
                      <ul className="mt-1 text-yellow-700 space-y-1">
                        <li>• <strong>Storage:</strong> {pendingModelToLoad?.size_gb}</li>
                        <li>• <strong>RAM:</strong> ~{pendingModelToLoad?.size_gb} (while running)</li>
                        <li>• <strong>Download:</strong> First-time loading will download the model</li>
                      </ul>
                    </div>
                  </div>
                </div>

                <div className="bg-blue-50 border border-blue-200 rounded-lg p-3">
                  <div className="text-sm text-blue-700">
                    <p className="font-medium text-blue-800">Model Features:</p>
                    <p className="mt-1">{pendingModelToLoad?.description}</p>
                    {pendingModelToLoad?.supports_thinking && (
                      <p className="mt-1 flex items-center gap-1">
                        <Brain className="h-3 w-3" />
                        Supports thinking process
                      </p>
                    )}
                  </div>
                </div>

                <p className="text-sm text-muted-foreground">
                  The model will be cached locally for faster future access. You can unload it anytime to free up memory.
                </p>
              </div>
            </AlertDialogDescription>
          </AlertDialogHeader>
          <AlertDialogFooter>
            <AlertDialogCancel onClick={handleLoadModelCancel}>
              Cancel
            </AlertDialogCancel>
            <AlertDialogAction onClick={handleLoadModelConfirm}>
              Load Model
            </AlertDialogAction>
          </AlertDialogFooter>
        </AlertDialogContent>
      </AlertDialog>
    </div>
  )
}