frontend/src/pages/MyDevice.tsx

import { useState, useEffect, useRef } 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);

  // Ref for auto-scrolling
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // Auto-scroll to bottom when messages change
  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  // 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 || [];

  // Auto-scroll to bottom when messages change
  useEffect(() => {
    scrollToBottom();
  }, [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-[500px]">
        {/* 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-y-auto p-4 space-y-4 max-h-96 min-h-0">
              {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 ref={messagesEndRef} />
            </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-4 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>
  );
}