frontend/src/pages/MyDevice.tsx

import { useState, useEffect } from 'react'
import { useNavigate } from 'react-router-dom'
import { Button } from '@/components/ui/button'
import { Card } from '@/components/ui/card'
import { Input } from '@/components/ui/input'
import { Label } from '@/components/ui/label'
import { Textarea } from '@/components/ui/textarea'
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select'
import { Slider } from '@/components/ui/slider'
import { Badge } from '@/components/ui/badge'
import { ASSISTANT_CONFIGS } from '@/config/assistants'
import { useChat } from '@/hooks/useChat'
import { DocumentsTab } from '@/components/playground/DocumentsTab'

interface DeviceConfig {
  assistantId: string
  model: string
  temperature: number
  maxTokens: number
  role: string
  goal: string
  description: string
  document: string
}

interface SelectedAssistant {
  id: string
  name: string
  type: 'user'|'template'|'new'
  originalTemplate?: string
}

export function MyDevice() {
  const navigate = useNavigate()
  const [currentStep, setCurrentStep] = useState(1)
  const [deviceConfig, setDeviceConfig] = useState<DeviceConfig>({
    assistantId: '',
    model: 'Qwen/Qwen3-30B-A3B',
    temperature: 0.7,
    maxTokens: 1024,
    role: '',
    goal: '',
    description: '',
    document: ''
  })
  
  // RAG document state
  const [ragEnabled, setRagEnabled] = useState(false)
  const [retrievalCount, setRetrievalCount] = useState(3)
  const [currentAssistant, setCurrentAssistant] = useState<SelectedAssistant | null>(null)

  // Load assistant configuration when selected
  const handleAssistantSelect = (assistantId: string) => {
    if (assistantId === 'create-new') {
      // Reset to blank configuration for new assistant
      setDeviceConfig(prev => ({
        ...prev,
        assistantId: 'create-new',
        model: 'Qwen/Qwen3-30B-A3B',
        temperature: 0.7,
        maxTokens: 1024,
        role: '',
        goal: '',
        description: ''
      }))
      
      // Reset chat settings
      setSelectedModel('Qwen/Qwen3-30B-A3B')
      setSystemPrompt('')
      setTemperature(0.7)
      setMaxTokens(1024)
      
      // Set current assistant for RAG
      setCurrentAssistant({
        id: 'create-new',
        name: 'New Custom Assistant',
        type: 'new'
      })
    } else {
      const config = assistantConfigs.find(c => c.id === assistantId)
      if (config) {
        setDeviceConfig(prev => ({
          ...prev,
          assistantId,
          model: config.model,
          temperature: config.temperature,
          maxTokens: config.maxTokens,
          role: config.description,
          goal: config.category,
          description: config.systemPrompt
        }))
        
        // Also update chat settings for Step 2
        setSelectedModel(config.model)
        setSystemPrompt(config.systemPrompt)
        setTemperature(config.temperature)
        setMaxTokens(config.maxTokens)
        
        // Set current assistant for RAG
        setCurrentAssistant({
          id: assistantId,
          name: config.name,
          type: 'template'
        })
      }
    }
  }

  // Check for pre-selected use case from Use Cases page
  useEffect(() => {
    const selectedUseCase = localStorage.getItem('selectedUseCase')
    if (selectedUseCase) {
      handleAssistantSelect(selectedUseCase)
      localStorage.removeItem('selectedUseCase') // Clear after use
    }
  }, [])

  // Chat functionality for Step 2
  const {
    currentSession,
    createNewSession,
    sendMessage,
    isLoading,
    setSystemPrompt,
    setTemperature,
    setMaxTokens,
    setSelectedModel,
    input,
    setInput
  } = useChat()

  const assistantConfigs = ASSISTANT_CONFIGS
  const messages = currentSession?.messages || []

  // Sync deviceConfig changes to chat settings
  useEffect(() => {
    setSelectedModel(deviceConfig.model)
    setSystemPrompt(deviceConfig.description)
    setTemperature(deviceConfig.temperature)
    setMaxTokens(deviceConfig.maxTokens)
  }, [deviceConfig, setSelectedModel, setSystemPrompt, setTemperature, setMaxTokens])

  const renderStepIndicator = () => (
    <div className="flex items-center justify-between mb-8 max-w-4xl mx-auto">
      <div className="flex items-center">
        <div className={`flex items-center ${currentStep >= 1 ? 'text-purple-600' : 'text-gray-400'}`}>
          <div className={`w-8 h-8 rounded-full border-2 flex items-center justify-center text-sm font-medium ${
            currentStep >= 1 ? 'bg-purple-600 text-white border-purple-600' : 'border-gray-300'
          }`}>
            1
          </div>
          <span className="ml-2 font-medium">Configuration</span>
        </div>
        <div className={`w-32 h-0.5 mx-4 ${currentStep >= 2 ? 'bg-purple-600' : 'bg-gray-300'}`} />
      </div>
      
      <div className="flex items-center">
        <div className={`flex items-center ${currentStep >= 2 ? 'text-purple-600' : 'text-gray-400'}`}>
          <div className={`w-8 h-8 rounded-full border-2 flex items-center justify-center text-sm font-medium ${
            currentStep >= 2 ? 'bg-purple-600 text-white border-purple-600' : 'border-gray-300'
          }`}>
            2
          </div>
          <span className="ml-2 font-medium">Test performance</span>
        </div>
        <div className={`w-32 h-0.5 mx-4 ${currentStep >= 3 ? 'bg-purple-600' : 'bg-gray-300'}`} />
      </div>

      <div className="flex items-center">
        <div className={`flex items-center ${currentStep >= 3 ? 'text-purple-600' : 'text-gray-400'}`}>
          <div className={`w-8 h-8 rounded-full border-2 flex items-center justify-center text-sm font-medium ${
            currentStep >= 3 ? 'bg-purple-600 text-white border-purple-600' : 'border-gray-300'
          }`}>
            3
          </div>
          <span className="ml-2 font-medium">Deploy to device</span>
        </div>
        <Button 
          className="ml-8 bg-purple-600 hover:bg-purple-700 text-white px-8"
          onClick={() => setCurrentStep(3)}
          disabled={currentStep < 2}
        >
          Connect
        </Button>
      </div>
    </div>
  )

  const renderStep1 = () => (
    <div className="max-w-2xl mx-auto">
      <Card className="p-8">
        <div className="space-y-6">
          {/* Load Assistant */}
          <div>
            <Label className="text-sm font-medium text-gray-700 mb-2 block">Load Assistant</Label>
            <Select onValueChange={handleAssistantSelect} value={deviceConfig.assistantId}>
              <SelectTrigger>
                <SelectValue placeholder="Select an assistant or create new" />
              </SelectTrigger>
              <SelectContent>
                <SelectItem value="create-new">
                  <div className="flex items-center gap-2">
                    <span className="text-lg">✨</span>
                    <span>Create New Assistant</span>
                  </div>
                </SelectItem>
                <div className="border-b my-1"></div>
                {assistantConfigs.map((config) => (
                  <SelectItem key={config.id} value={config.id}>
                    <div className="flex items-center gap-2">
                      <span>{config.icon}</span>
                      <span>{config.name.replace(/🏔️|🏥|🧸/g, '').trim()}</span>
                    </div>
                  </SelectItem>
                ))}
              </SelectContent>
            </Select>
          </div>

          {/* Parameters */}
          <div>
            <h3 className="text-lg font-semibold mb-4">Parameters</h3>
            
            {/* Model Selection */}
            <div className="mb-4">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">Select a model</Label>
              <Select 
                value={deviceConfig.model} 
                onValueChange={(value) => setDeviceConfig(prev => ({ ...prev, model: value }))}
              >
                <SelectTrigger>
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  <SelectItem value="Qwen/Qwen3-30B-A3B">Qwen3-30B-A3B</SelectItem>
                  <SelectItem value="Qwen/Qwen2.5-32B-Instruct">Qwen2.5-32B-Instruct</SelectItem>
                  <SelectItem value="deepseek-ai/DeepSeek-V2.5">DeepSeek-V2.5</SelectItem>
                </SelectContent>
              </Select>
            </div>

            {/* Temperature */}
            <div className="mb-4">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">
                Temperature: {deviceConfig.temperature}
              </Label>
              <Slider
                value={[deviceConfig.temperature]}
                onValueChange={(value) => setDeviceConfig(prev => ({ ...prev, temperature: value[0] }))}
                max={1}
                min={0}
                step={0.1}
                className="w-full"
              />
            </div>

            {/* Max Tokens */}
            <div className="mb-6">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">Max Tokens</Label>
              <Input
                type="number"
                value={deviceConfig.maxTokens}
                onChange={(e) => setDeviceConfig(prev => ({ ...prev, maxTokens: parseInt(e.target.value) || 1024 }))}
                min={1}
                max={4096}
              />
            </div>
          </div>

          {/* Description */}
          <div>
            <h3 className="text-lg font-semibold mb-4">Description</h3>
            
            <div className="mb-4">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">
                Which role do you want your assistant to be? (Senior data researcher, journalist, lawyer, etc.)
              </Label>
              <Input
                value={deviceConfig.role}
                onChange={(e) => setDeviceConfig(prev => ({ ...prev, role: e.target.value }))}
                placeholder="e.g., Senior data researcher, journalist, lawyer"
              />
            </div>

            <div className="mb-4">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">
                Which goal would you expect your assistant to achieve
              </Label>
              <Input
                value={deviceConfig.goal}
                onChange={(e) => setDeviceConfig(prev => ({ ...prev, goal: e.target.value }))}
                placeholder="Describe the main goal or objective"
              />
            </div>

            <div className="mb-6">
              <Label className="text-sm font-medium text-gray-700 mb-2 block">
                Description (Tasks, outputs)
              </Label>
              <Textarea
                value={deviceConfig.description}
                onChange={(e) => setDeviceConfig(prev => ({ ...prev, description: e.target.value }))}
                placeholder="Describe the tasks and expected outputs..."
                rows={4}
              />
            </div>
          </div>

          {/* Document - RAG Integration */}
          <div>
            <Label className="text-sm font-medium text-gray-700 mb-2 block">Documents & Knowledge Base</Label>
            <DocumentsTab
              isLoading={isLoading}
              ragEnabled={ragEnabled}
              setRagEnabled={setRagEnabled}
              retrievalCount={retrievalCount}
              setRetrievalCount={setRetrievalCount}
              currentAssistant={currentAssistant}
            />
          </div>

          <Button 
            onClick={() => setCurrentStep(2)} 
            className="w-full bg-purple-600 hover:bg-purple-700 text-white"
            disabled={!deviceConfig.assistantId || (deviceConfig.assistantId === 'create-new' && !deviceConfig.description.trim())}
          >
            Continue to Test Performance
          </Button>
        </div>
      </Card>
    </div>
  )

  const renderStep2 = () => (
    <div className="max-w-6xl mx-auto">
      <div className="grid grid-cols-1 lg:grid-cols-4 gap-6 h-[600px]">
        {/* Chat Sessions Sidebar */}
        <div className="lg:col-span-1">
          <Card className="h-full p-4">
            <h3 className="font-semibold mb-4">Chat sessions</h3>
            <div className="space-y-2">
              <Button
                variant="ghost"
                className="w-full justify-start text-left"
                onClick={createNewSession}
              >
                New Chat
              </Button>
              {messages.length > 0 && (
                <div className="p-2 bg-gray-50 rounded text-sm">
                  Current Session ({messages.length} messages)
                </div>
              )}
            </div>
          </Card>
        </div>

        {/* Chat Interface */}
        <div className="lg:col-span-3">
          <Card className="h-full flex flex-col">
            <div className="p-4 border-b">
              <div className="flex items-center justify-between">
                <h3 className="font-semibold">Test Performance</h3>
                <Badge variant="secondary">
                  {deviceConfig.assistantId === 'create-new' ? 'Custom Assistant' : 
                   deviceConfig.assistantId ? assistantConfigs.find(c => c.id === deviceConfig.assistantId)?.name.replace(/🏔️|🏥|🧸/g, '').trim() : 'No Assistant'}
                </Badge>
              </div>
            </div>
            
            <div className="flex-1 overflow-auto p-4 space-y-4">
              {messages.length === 0 ? (
                <div className="flex items-center justify-center h-full text-gray-500">
                  <div className="text-center">
                    <p className="mb-2">Start testing your assistant configuration</p>
                    <p className="text-sm">Send a message to see how your assistant responds</p>
                  </div>
                </div>
              ) : (
                messages.map((message, index) => (
                  <div
                    key={index}
                    className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
                  >
                    <div
                      className={`max-w-[80%] p-3 rounded-lg ${
                        message.role === 'user'
                          ? 'bg-purple-600 text-white'
                          : 'bg-gray-100 text-gray-900'
                      }`}
                    >
                      <p className="whitespace-pre-wrap">{message.content}</p>
                    </div>
                  </div>
                ))
              )}
              {isLoading && (
                <div className="flex justify-start">
                  <div className="bg-gray-100 p-3 rounded-lg">
                    <div className="flex items-center space-x-2">
                      <div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce"></div>
                      <div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce" style={{ animationDelay: '0.1s' }}></div>
                      <div className="w-2 h-2 bg-gray-400 rounded-full animate-bounce" style={{ animationDelay: '0.2s' }}></div>
                    </div>
                  </div>
                </div>
              )}
            </div>

            <div className="p-4 border-t">
              <div className="flex space-x-2">
                <Input
                  placeholder="Type your message..."
                  value={input}
                  onChange={(e) => setInput(e.target.value)}
                  onKeyPress={(e) => {
                    if (e.key === 'Enter' && !e.shiftKey) {
                      e.preventDefault()
                      if (input.trim()) {
                        sendMessage(undefined, { useRag: ragEnabled, retrievalCount })
                      }
                    }
                  }}
                />
                <Button 
                  onClick={() => {
                    if (input.trim()) {
                      sendMessage(undefined, { useRag: ragEnabled, retrievalCount })
                    }
                  }}
                  disabled={isLoading}
                  className="bg-purple-600 hover:bg-purple-700"
                >
                  Send
                </Button>
              </div>
            </div>
          </Card>
        </div>
      </div>

      <div className="mt-6 flex justify-between max-w-6xl mx-auto">
        <Button 
          variant="outline" 
          onClick={() => setCurrentStep(1)}
        >
          Back to Configuration
        </Button>
        <Button 
          onClick={() => setCurrentStep(3)} 
          className="bg-purple-600 hover:bg-purple-700 text-white"
        >
          Continue to Deploy Device
        </Button>
      </div>
    </div>
  )

  const renderStep3 = () => (
    <div className="max-w-2xl mx-auto">
      <Card className="p-8 text-center">
        <div className="mb-6">
          <div className="w-16 h-16 bg-purple-100 rounded-full flex items-center justify-center mx-auto mb-4">
            <svg className="w-8 h-8 text-purple-600" fill="none" stroke="currentColor" viewBox="0 0 24 24">
              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M8.111 16.404a5.5 5.5 0 017.778 0M12 20h.01m-7.08-7.071c3.904-3.905 10.236-3.905 14.141 0M1.394 9.393c5.857-5.857 15.355-5.857 21.212 0" />
            </svg>
          </div>
          <h2 className="text-2xl font-bold text-gray-900 mb-2">Deploy to Device</h2>
          <p className="text-gray-600">
            Your assistant configuration is ready to be deployed to your edge device.
          </p>
        </div>

        <div className="space-y-4 text-left mb-8">
          <div className="bg-gray-50 p-4 rounded-lg">
            <h3 className="font-semibold mb-2">Configuration Summary</h3>
            <div className="space-y-1 text-sm">
              <p><span className="font-medium">Assistant:</span> {
                deviceConfig.assistantId === 'create-new' ? 'Custom Assistant' :
                assistantConfigs.find(c => c.id === deviceConfig.assistantId)?.name.replace(/🏔️|🏥|🧸/g, '').trim() || 'None'
              }</p>
              <p><span className="font-medium">Model:</span> {deviceConfig.model}</p>
              <p><span className="font-medium">Temperature:</span> {deviceConfig.temperature}</p>
              <p><span className="font-medium">Max Tokens:</span> {deviceConfig.maxTokens}</p>
              {deviceConfig.role && <p><span className="font-medium">Role:</span> {deviceConfig.role}</p>}
              <p><span className="font-medium">RAG Enabled:</span> {ragEnabled ? 'Yes' : 'No'}</p>
              {ragEnabled && <p><span className="font-medium">Retrieval Count:</span> {retrievalCount}</p>}
            </div>
          </div>
          
          <div className="bg-blue-50 p-4 rounded-lg border border-blue-200">
            <div className="flex items-center">
              <svg className="w-5 h-5 text-blue-600 mr-2" fill="none" stroke="currentColor" viewBox="0 0 24 24">
                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" />
              </svg>
              <p className="text-blue-800 text-sm">
                <span className="font-medium">Ready to deploy:</span> Your assistant configuration has been optimized for edge deployment.
              </p>
            </div>
          </div>
        </div>

        <div className="space-y-3">
          <Button className="w-full bg-purple-600 hover:bg-purple-700 text-white py-3">
            Deploy to Device
          </Button>
          <Button variant="outline" className="w-full" onClick={() => setCurrentStep(2)}>
            Back to Testing
          </Button>
        </div>
      </Card>
    </div>
  )

  return (
    <div className="min-h-screen bg-gray-50">
      {/* Header */}
      <header className="bg-white border-b border-gray-200">
        <div className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8">
          <div className="flex items-center justify-between h-20">
            <div className="flex items-center space-x-2">
              <img 
                src="/assets/logo.png" 
                alt="EdgeLLM Logo" 
                className="h-16 w-16"
                onError={(e) => {
                  console.error('Logo failed to load')
                  e.currentTarget.style.display = 'none'
                }}
              />
            </div>
            <nav className="flex space-x-8">
              <button 
                onClick={() => navigate('/')}
                className="text-gray-600 hover:text-gray-900 px-3 py-2 text-sm font-medium"
              >
                Home
              </button>
              <button 
                onClick={() => navigate('/technology')}
                className="text-gray-600 hover:text-gray-900 px-3 py-2 text-sm font-medium"
              >
                Technology
              </button>
              <button 
                onClick={() => navigate('/usecases')}
                className="text-gray-600 hover:text-gray-900 px-3 py-2 text-sm font-medium"
              >
                Use Case
              </button>
              <button 
                className="bg-purple-600 text-white px-4 py-2 text-sm font-medium rounded"
              >
                My Device
              </button>
            </nav>
          </div>
        </div>
      </header>

      {/* Main Content */}
      <main className="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-8">
        {renderStepIndicator()}
        
        {currentStep === 1 && renderStep1()}
        {currentStep === 2 && renderStep2()}
        {currentStep === 3 && renderStep3()}
      </main>
    </div>
  )
}