Chat.tsx

import { useState, useRef, useEffect } from "react";
import { ChevronUp, ChevronDown, Send, Search, Zap, Upload, X } from "lucide-react";
import { Streamdown } from "streamdown";
import ChatSidebar, { ChatSession } from "@/components/ChatSidebar";

/**
 * Chat Page with Sidebar and Local Storage
 * 
 * Features:
 * - Ask | Imagine mode switcher at top
 * - Sidebar with chat history (local storage)
 * - Advanced prompt input with Search/Think toggles
 * - DeepSeek reasoning panel (collapsible)
 * - Rich response formatting
 * - File upload with OCR
 * - Image gallery for Imagine mode
 */

type ChatMode = "ask" | "imagine";

interface Message {
  id: string;
  role: "user" | "assistant";
  content: string;
  reasoning?: string;
  timestamp: Date;
}

interface GeneratedImage {
  id: string;
  prompt: string;
  url: string;
  timestamp: Date;
}

export default function Chat() {
  // ========================================================================
  // State Management
  // ========================================================================

  const [mode, setMode] = useState<ChatMode>("ask");
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState("");
  const [isLoading, setIsLoading] = useState(false);
  const [enableSearch, setEnableSearch] = useState(false);
  const [enableThinking, setEnableThinking] = useState(false);
  const [showReasoning, setShowReasoning] = useState(false);
  const [uploadedFiles, setUploadedFiles] = useState<File[]>([]);
  const [generatedImages, setGeneratedImages] = useState<GeneratedImage[]>([]);
  const [galleryOpen, setGalleryOpen] = useState(false);
  const [sidebarOpen, setSidebarOpen] = useState(true);
  const [currentChatId, setCurrentChatId] = useState<string | null>(null);

  const messagesEndRef = useRef<HTMLDivElement>(null);
  const fileInputRef = useRef<HTMLInputElement>(null);

  // ========================================================================
  // API Configuration - UPDATE THIS WITH YOUR HUGGING FACE SPACE URL
  // ========================================================================

  const API_BASE_URL = process.env.REACT_APP_API_URL || "http://localhost:3000";

  // ========================================================================
  // Local Storage Functions
  // ========================================================================

  const generateChatId = () => `chat_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;

  const saveChatToStorage = (chatId: string, msgs: Message[], title: string) => {
    // Save chat messages
    localStorage.setItem(`domify_chat_${chatId}`, JSON.stringify(msgs));

    // Update chat metadata
    const savedChats = localStorage.getItem("domify_chats");
    let chats: ChatSession[] = savedChats ? JSON.parse(savedChats) : [];

    const existingIndex = chats.findIndex((c) => c.id === chatId);
    const chatSession: ChatSession = {
      id: chatId,
      title: title || "Untitled Chat",
      timestamp: Date.now(),
      mode,
      messageCount: msgs.length,
    };

    if (existingIndex >= 0) {
      chats[existingIndex] = chatSession;
    } else {
      chats.push(chatSession);
    }

    localStorage.setItem("domify_chats", JSON.stringify(chats));
  };

  const loadChatFromStorage = (chatId: string) => {
    const savedMessages = localStorage.getItem(`domify_chat_${chatId}`);
    if (savedMessages) {
      try {
        const parsed = JSON.parse(savedMessages);
        const msgs = parsed.map((m: any) => ({
          ...m,
          timestamp: new Date(m.timestamp),
        }));
        setMessages(msgs);
        setCurrentChatId(chatId);
      } catch (error) {
        console.error("Error loading chat:", error);
      }
    }
  };

  const createNewChat = () => {
    const chatId = generateChatId();
    setMessages([]);
    setCurrentChatId(chatId);
    setGeneratedImages([]);
    setInput("");
  };

  // ========================================================================
  // Auto-save chat to storage when messages change
  // ========================================================================

  useEffect(() => {
    if (currentChatId && messages.length > 0) {
      const title =
        messages[0]?.content?.substring(0, 50) || "New Chat";
      saveChatToStorage(currentChatId, messages, title);
    }
  }, [messages, currentChatId]);

  // ========================================================================
  // Auto-scroll to latest message
  // ========================================================================

  useEffect(() => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  }, [messages]);

  // ========================================================================
  // Handle file upload
  // ========================================================================

  const handleFileUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
    const files = Array.from(e.target.files || []);
    setUploadedFiles((prev) => [...prev, ...files]);
  };

  const removeFile = (index: number) => {
    setUploadedFiles((prev) => prev.filter((_, i) => i !== index));
  };

  // ========================================================================
  // Handle message send (Ask mode)
  // ========================================================================

  const handleSendMessage = async () => {
    if (!input.trim() && uploadedFiles.length === 0) return;

    // Create new chat if none exists
    if (!currentChatId) {
      createNewChat();
    }

    const userMessage: Message = {
      id: Date.now().toString(),
      role: "user",
      content: input,
      timestamp: new Date(),
    };

    setMessages((prev) => [...prev, userMessage]);
    setInput("");
    setUploadedFiles([]);
    setIsLoading(true);

    try {
      // Call your backend API
      const response = await fetch(`${API_BASE_URL}/api/trpc/chat.send`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
          prompt: input,
          enableSearch,
          enableThinking,
          history: messages.map((m) => ({
            role: m.role,
            content: m.content,
          })),
        }),
      });

      const data = await response.json();

      if (data.result?.data) {
        const assistantMessage: Message = {
          id: (Date.now() + 1).toString(),
          role: "assistant",
          content: data.result.data.response,
          reasoning: data.result.data.reasoning,
          timestamp: new Date(),
        };

        setMessages((prev) => [...prev, assistantMessage]);
        if (data.result.data.reasoning) {
          setShowReasoning(true);
        }
      }
    } catch (error) {
      console.error("Error sending message:", error);
      const errorMessage: Message = {
        id: (Date.now() + 1).toString(),
        role: "assistant",
        content: "Sorry, there was an error processing your request. Please try again.",
        timestamp: new Date(),
      };
      setMessages((prev) => [...prev, errorMessage]);
    } finally {
      setIsLoading(false);
    }
  };

  // ========================================================================
  // Handle image generation (Imagine mode)
  // ========================================================================

  const handleGenerateImage = async () => {
    if (!input.trim()) return;

    // Create new chat if none exists
    if (!currentChatId) {
      createNewChat();
    }

    setIsLoading(true);

    try {
      const response = await fetch(`${API_BASE_URL}/api/trpc/imagine.generate`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ prompt: input }),
      });

      const data = await response.json();

      if (data.result?.data?.imageUrl) {
        const newImage: GeneratedImage = {
          id: Date.now().toString(),
          prompt: input,
          url: data.result.data.imageUrl,
          timestamp: new Date(),
        };

        setGeneratedImages((prev) => [...prev, newImage]);
        setInput("");
        setGalleryOpen(true);

        // Save to storage
        const chatId = currentChatId!;
        const savedChats = localStorage.getItem("domify_chats");
        let chats: ChatSession[] = savedChats ? JSON.parse(savedChats) : [];
        const existingIndex = chats.findIndex((c) => c.id === chatId);
        if (existingIndex >= 0) {
          chats[existingIndex].messageCount = generatedImages.length + 1;
          localStorage.setItem("domify_chats", JSON.stringify(chats));
        }
      }
    } catch (error) {
      console.error("Error generating image:", error);
    } finally {
      setIsLoading(false);
    }
  };

  // ========================================================================
  // Render: Ask Mode
  // ========================================================================

  if (mode === "ask") {
    return (
      <div className="min-h-screen bg-background text-foreground flex">
        {/* Sidebar */}
        <ChatSidebar
          currentChatId={currentChatId}
          onNewChat={createNewChat}
          onSelectChat={loadChatFromStorage}
          onDeleteChat={() => {
            setMessages([]);
            setCurrentChatId(null);
          }}
          isOpen={sidebarOpen}
          onToggle={() => setSidebarOpen(!sidebarOpen)}
        />

        {/* Main Content */}
        <div
          className={`flex-1 flex flex-col transition-all duration-300 ${
            sidebarOpen ? "md:ml-64" : "ml-0"
          }`}
        >
          {/* Header */}
          <div className="glass-panel-lg m-4 p-6">
            <div className="flex items-center justify-between">
              <h1 className="gradient-text text-3xl font-bold">Domify Academy Bot</h1>
              <div className="flex gap-2">
                <button
                  onClick={() => setMode("ask" as ChatMode)}
                  className={`px-6 py-2 rounded-lg font-medium transition-smooth ${
                    mode === ("ask" as ChatMode)
                      ? "bg-primary text-primary-foreground glow-primary"
                      : "btn-ghost"
                  }`}
                >
                  Ask
                </button>
                <button
                  onClick={() => setMode("imagine" as ChatMode)}
                  className={`px-6 py-2 rounded-lg font-medium transition-smooth ${
                    mode === ("imagine" as ChatMode)
                      ? "bg-primary text-primary-foreground glow-primary"
                      : "btn-ghost"
                  }`}
                >
                  Imagine
                </button>
              </div>
            </div>
          </div>

          {/* Messages Area */}
          <div className="flex-1 overflow-y-auto px-4 space-y-4 scrollbar-thin">
            {messages.length === 0 ? (
              <div className="flex items-center justify-center h-full">
                <div className="text-center space-y-4">
                  <p className="text-muted-foreground text-lg">
                    Start a conversation with Domify Academy Bot
                  </p>
                  <p className="text-sm text-muted-foreground">
                    Ask questions, get reasoning, search online, and more
                  </p>
                </div>
              </div>
            ) : (
              messages.map((msg) => (
                <div
                  key={msg.id}
                  className={`animate-slide-up ${
                    msg.role === "user" ? "flex justify-end" : "flex justify-start"
                  }`}
                >
                  <div
                    className={`max-w-2xl glass-panel p-4 ${
                      msg.role === "user"
                        ? "bg-primary/20 border-primary/30"
                        : "bg-secondary/10 border-secondary/30"
                    }`}
                  >
                    {/* Reasoning Panel (if available) */}
                    {msg.reasoning && msg.role === "assistant" && (
                      <div className="mb-4">
                        <button
                          onClick={() => setShowReasoning(!showReasoning)}
                          className="flex items-center gap-2 text-sm text-accent hover:text-primary transition-smooth"
                        >
                          <span>🧠 Reasoning</span>
                          {showReasoning ? (
                            <ChevronUp size={16} />
                          ) : (
                            <ChevronDown size={16} />
                          )}
                        </button>
                        {showReasoning && (
                          <div className="mt-2 p-3 bg-white/5 rounded-lg text-sm text-muted-foreground border border-white/10 animate-slide-up">
                            {msg.reasoning}
                          </div>
                        )}
                      </div>
                    )}

                    {/* Message Content */}
                    <div className="markdown">
                      <Streamdown>{msg.content}</Streamdown>
                    </div>

                    {/* Timestamp */}
                    <p className="text-xs text-muted-foreground mt-2">
                      {msg.timestamp.toLocaleTimeString()}
                    </p>
                  </div>
                </div>
              ))
            )}

            {isLoading && (
              <div className="flex justify-start animate-slide-up">
                <div className="glass-panel p-4 bg-secondary/10">
                  <div className="flex gap-2">
                    <div className="w-2 h-2 bg-primary rounded-full animate-bounce" />
                    <div
                      className="w-2 h-2 bg-primary rounded-full animate-bounce"
                      style={{ animationDelay: "0.2s" }}
                    />
                    <div
                      className="w-2 h-2 bg-primary rounded-full animate-bounce"
                      style={{ animationDelay: "0.4s" }}
                    />
                  </div>
                </div>
              </div>
            )}

            <div ref={messagesEndRef} />
          </div>

          {/* Input Area */}
          <div className="glass-panel-lg m-4 p-6 space-y-4">
            {/* File Preview */}
            {uploadedFiles.length > 0 && (
              <div className="flex flex-wrap gap-2">
                {uploadedFiles.map((file, idx) => (
                  <div
                    key={idx}
                    className="glass-panel px-3 py-2 flex items-center gap-2 text-sm"
                  >
                    <span>{file.name}</span>
                    <button
                      onClick={() => removeFile(idx)}
                      className="hover:text-destructive transition-smooth"
                    >
                      <X size={16} />
                    </button>
                  </div>
                ))}
              </div>
            )}

            {/* Toggle Buttons */}
            <div className="flex gap-2 flex-wrap">
              <button
                onClick={() => setEnableSearch(!enableSearch)}
                className={`flex items-center gap-2 px-4 py-2 rounded-lg transition-smooth ${
                  enableSearch
                    ? "bg-primary/30 text-primary border border-primary/50"
                    : "btn-ghost"
                }`}
              >
                <Search size={18} />
                <span>Search Online</span>
              </button>

              <button
                onClick={() => setEnableThinking(!enableThinking)}
                className={`flex items-center gap-2 px-4 py-2 rounded-lg transition-smooth ${
                  enableThinking
                    ? "bg-accent/30 text-accent border border-accent/50"
                    : "btn-ghost"
                }`}
              >
                <Zap size={18} />
                <span>Think Longer</span>
              </button>

              <button
                onClick={() => fileInputRef.current?.click()}
                className="flex items-center gap-2 px-4 py-2 rounded-lg btn-ghost"
              >
                <Upload size={18} />
                <span>Upload</span>
              </button>

              <input
                ref={fileInputRef}
                type="file"
                multiple
                onChange={handleFileUpload}
                className="hidden"
                accept="image/*,.pdf,.txt,.doc,.docx"
              />
            </div>

            {/* Input Box */}
            <div className="flex gap-2">
              <textarea
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === "Enter" && !e.shiftKey) {
                    e.preventDefault();
                    handleSendMessage();
                  }
                }}
                placeholder="Ask me anything... (Shift+Enter for new line)"
                className="input-glass flex-1 resize-none max-h-32"
                rows={3}
              />
              <button
                onClick={handleSendMessage}
                disabled={isLoading || (!input.trim() && uploadedFiles.length === 0)}
                className="btn-primary self-end"
              >
                <Send size={20} />
              </button>
            </div>
          </div>
        </div>
      </div>
    );
  }

  // ========================================================================
  // Render: Imagine Mode
  // ========================================================================

  return (
    <div className="min-h-screen bg-background text-foreground flex">
      {/* Sidebar */}
      <ChatSidebar
        currentChatId={currentChatId}
        onNewChat={createNewChat}
        onSelectChat={loadChatFromStorage}
        onDeleteChat={() => {
          setGeneratedImages([]);
          setCurrentChatId(null);
        }}
        isOpen={sidebarOpen}
        onToggle={() => setSidebarOpen(!sidebarOpen)}
      />

      {/* Main Content */}
      <div
        className={`flex-1 flex flex-col transition-all duration-300 ${
          sidebarOpen ? "md:ml-64" : "ml-0"
        }`}
      >
        {/* Header */}
        <div className="glass-panel-lg m-4 p-6">
          <div className="flex items-center justify-between">
            <h1 className="gradient-text text-3xl font-bold">Domify Academy Bot</h1>
            <div className="flex gap-2">
              <button
                onClick={() => setMode("ask" as ChatMode)}
                className={`px-6 py-2 rounded-lg font-medium transition-smooth ${
                  mode === ("ask" as ChatMode)
                    ? "bg-primary text-primary-foreground glow-primary"
                    : "btn-ghost"
                }`}
              >
                Ask
              </button>
              <button
                onClick={() => setMode("imagine" as ChatMode)}
                className={`px-6 py-2 rounded-lg font-medium transition-smooth ${
                  mode === ("imagine" as ChatMode)
                    ? "bg-primary text-primary-foreground glow-primary"
                    : "btn-ghost"
                }`}
              >
                Imagine
              </button>
            </div>
          </div>
        </div>

        {/* Gallery */}
        {galleryOpen && generatedImages.length > 0 && (
          <div className="glass-panel-lg m-4 p-6">
            <div className="flex items-center justify-between mb-4">
              <h2 className="text-xl font-semibold">Generated Images</h2>
              <button
                onClick={() => setGalleryOpen(false)}
                className="btn-ghost"
              >
                <ChevronUp size={20} />
              </button>
            </div>

            <div className="overflow-x-auto pb-4">
              <div className="flex gap-4">
                {generatedImages.map((img) => (
                  <div key={img.id} className="flex-shrink-0 glass-panel p-4 space-y-2">
                    <img
                      src={img.url}
                      alt={img.prompt}
                      className="w-48 h-48 object-cover rounded-lg"
                    />
                    <p className="text-sm text-muted-foreground truncate">{img.prompt}</p>
                    <div className="flex gap-2">
                      <a
                        href={img.url}
                        download
                        className="btn-primary text-sm flex-1 text-center"
                      >
                        Download
                      </a>
                      <button
                        className="btn-ghost text-sm flex-1"
                        disabled
                      >
                        Video (Soon)
                      </button>
                    </div>
                  </div>
                ))}
              </div>
            </div>
          </div>
        )}

        {/* Main Content */}
        <div className="flex-1 flex items-center justify-center px-4">
          <div className="w-full max-w-2xl space-y-6">
            <div className="text-center space-y-2">
              <h2 className="text-3xl font-bold gradient-text">Create Images</h2>
              <p className="text-muted-foreground">
                Describe what you want to imagine, and I'll create it for you
              </p>
            </div>

            {/* Input Area */}
            <div className="glass-panel-lg p-6 space-y-4">
              <textarea
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={(e) => {
                  if (e.key === "Enter" && !e.shiftKey) {
                    e.preventDefault();
                    handleGenerateImage();
                  }
                }}
                placeholder="Describe the image you want to generate..."
                className="input-glass w-full resize-none"
                rows={4}
              />

              <button
                onClick={handleGenerateImage}
                disabled={isLoading || !input.trim()}
                className="btn-primary w-full py-3 text-lg"
              >
                {isLoading ? "Generating..." : "Generate Image"}
              </button>
            </div>
          </div>
        </div>
      </div>
    </div>
  );
}