frontend/src/components/chat/ChatInterface.tsx

import { useState, useRef, useEffect } from "react";
import { Send, Plus, PanelLeft, Bot } from "lucide-react";
import { Button } from "@/components/ui/button";
import { apiService, Message as APIMessage } from "@/services/apiService";
import { toast } from "@/components/ui/sonner";
import { ChatBubble } from "@/components/chat/ChatBubble";
import { Separator } from "@/components/ui/separator";
import { cn } from "@/lib/utils";
import { storage, STORAGE_KEYS } from "@/lib/storage";
import { Message, Chat, MessageVariation } from "@/types/chat";
import { ProfileModal } from "../modals/ProfileModal";
import { ChatSidebar } from "./ChatSidebar";
import { ChatInputArea } from "./ChatInputArea";
import { WelcomeScreen } from "./WelcomeScreen";
import { DeleteChatDialog } from "./DeleteChatDialog";

interface ChatInterfaceProps {
  onOpenSettings: () => void;
  onOpenSources: () => void;
}

const WELCOME_MESSAGE = "Hello! I'm Insight AI, How can I help you today?";

const generateId = () => Math.random().toString(36).substring(2, 11);

export const ChatInterface = ({ onOpenSettings, onOpenSources }: ChatInterfaceProps) => {
  const [chats, setChats] = useState<Chat[]>(() => {
    const savedChats = storage.get<Chat[]>(STORAGE_KEYS.CHATS);
    if (savedChats) {
      try {
        return savedChats.map((chat: any) => ({
          ...chat,
          messages: chat.messages.map((msg: any) => ({
            ...msg,
            timestamp: new Date(msg.timestamp),
            sender: msg.sender as "user" | "system"
          })),
          createdAt: new Date(chat.createdAt),
          updatedAt: new Date(chat.updatedAt)
        }));
      } catch (error) {
        console.error("Failed to parse saved chats:", error);
        return [];
      }
    }
    return [];
  });

  const [activeChat, setActiveChat] = useState<Chat | null>(() => {
    if (chats.length > 0) {
      return chats[0];
    }
  });

  const [inputValue, setInputValue] = useState("");
  const [isLoading, setIsLoading] = useState(false);
  const [isSidebarOpen, setIsSidebarOpen] = useState(false);
  const [isGeneratingTitle, setIsGeneratingTitle] = useState(false);
  const [isProfileModalOpen, setIsProfileModalOpen] = useState(false);
  const [chatToDelete, setChatToDelete] = useState<string | null>(null);

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

  useEffect(() => {
    // Save chats to storage whenever they change
    if (activeChat && !chats.find(chat => chat.id === activeChat.id)) {
      setChats([activeChat, ...chats]);
    }

    const allChats = activeChat
      ? [
        activeChat,
        ...chats.filter(chat => chat.id !== activeChat.id)
      ]
      : chats;

    storage.set(STORAGE_KEYS.CHATS, allChats);
  }, [chats, activeChat]);

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  useEffect(() => {
    scrollToBottom();
  }, [activeChat?.messages]);

  useEffect(() => {
    // Focus input when component mounts or when loading ends
    if (!isLoading) {
      setTimeout(() => {
        inputRef.current?.focus();
      }, 100);
    }
  }, [isLoading]);

  // Handle new chat and chat selection events
  useEffect(() => {
    const handleNewChat = () => createNewChat();
    const handleSelectChat = (e: Event) => {
      const customEvent = e as CustomEvent;
      const chatId = customEvent.detail?.chatId;
      if (chatId) {
        selectChat(chatId);
      }
    };
    const handleDeleteChat = (e: Event) => {
      const customEvent = e as CustomEvent;
      const chatId = customEvent.detail?.chatId;
      if (chatId) {
        setChatToDelete(chatId);
      }
    };

    document.addEventListener("insight:new-chat", handleNewChat);
    document.addEventListener("insight:select-chat", handleSelectChat);
    document.addEventListener("insight:delete-chat", handleDeleteChat);

    return () => {
      document.removeEventListener("insight:new-chat", handleNewChat);
      document.removeEventListener("insight:select-chat", handleSelectChat);
      document.removeEventListener("insight:delete-chat", handleDeleteChat);
    };
  }, [chats, activeChat]);

  const generateChatTitle = async (query: string) => {
    if (!activeChat || activeChat.title !== "New Chat") return;

    setIsGeneratingTitle(true);

    try {
      const response = await apiService.generateTitle(query);

      if (response.title) {
        // Update the active chat with the new title
        setActiveChat(prevChat => {
          if (!prevChat) return null;

          const updatedChat = { ...prevChat, title: response.title };

          // Update chats list
          setChats(prevChats =>
            prevChats.map(chat =>
              chat.id === updatedChat.id ? updatedChat : chat
            )
          );

          return updatedChat;
        });
      }
    } catch (error) {
      console.error("Error generating chat title:", error);
      // Fallback to using query as title
      if (activeChat.title === "New Chat") {
        setActiveChat(prevChat => {
          if (!prevChat) return null;

          const updatedChat = {
            ...prevChat,
            title: query.slice(0, 30) + (query.length > 30 ? '...' : '')
          };

          setChats(prevChats =>
            prevChats.map(chat =>
              chat.id === updatedChat.id ? updatedChat : chat
            )
          );

          return updatedChat;
        });
      }
    } finally {
      setIsGeneratingTitle(false);
    }
  };

  const handleSendMessage = async (e?: React.FormEvent) => {
    if (e) e.preventDefault();

    if (!inputValue.trim() || !activeChat) return;

    const userMessage: Message = {
      id: generateId(),
      content: inputValue,
      sender: "user",
      timestamp: new Date()
    };

    const loadingMessage: Message = {
      id: generateId(),
      content: "",
      sender: "system",
      timestamp: new Date(),
      isLoading: true
    };

    // Find the active variation if we're replying to a message that has variations
    let parentVariationId: string | undefined;
    let parentMessageId: string | undefined;
    
    // Find the last system message that might have variations
    const lastMessages = [...activeChat.messages].reverse();
    for (const msg of lastMessages) {
      if (msg.sender === "system" && msg.variations && msg.variations.length > 0 && msg.activeVariation) {
        parentMessageId = msg.id;
        parentVariationId = msg.activeVariation;
        break;
      }
    }
    
    // If responding to a variation, link these messages to that variation
    if (parentMessageId && parentVariationId) {
      userMessage.parentMessageId = parentMessageId;
      userMessage.variationId = parentVariationId;
      loadingMessage.parentMessageId = parentMessageId;
      loadingMessage.variationId = parentVariationId;
    }

    // Update active chat with new messages
    const updatedChat = {
      ...activeChat,
      messages: [...activeChat.messages, userMessage, loadingMessage],
      updatedAt: new Date()
    };

    setActiveChat(updatedChat);
    setInputValue("");
    setIsLoading(true);

    try {
      // Prepare chat history for the API based on the active variation path
      const chatHistory: APIMessage[] = [];
      
      // Build chat history based on the active variation path
      const getMessagesForHistory = (messages: Message[]): Message[] => {
        const result: Message[] = [];
        
        for (const msg of messages) {
          if (msg.isLoading) continue;
          
          // If this is a system message with variations, use the active variation content
          if (msg.sender === "system" && msg.variations && msg.variations.length > 0 && msg.activeVariation) {
            const activeVar = msg.variations.find(v => v.id === msg.activeVariation);
            if (activeVar) {
              // Add the message with the active variation content
              result.push({
                ...msg,
                content: activeVar.content
              });
            }
          } else {
            // Add regular messages
            result.push(msg);
          }
        }
        
        return result;
      };
      
      // Get messages following the active variation path
      const historyMessages = getMessagesForHistory(updatedChat.messages.slice(0, -2));
      
      // Convert to API format
      for (const msg of historyMessages) {
        if (!msg.content) continue;
        
        // Strip thinking content from messages before sending to API
        const cleanedContent = msg.content.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
        if (!cleanedContent) continue;
        
        chatHistory.push({
          role: msg.sender === "user" ? "user" : "assistant",
          content: cleanedContent
        });
      }

      const response = await apiService.queryRulings({
        query: userMessage.content,
        chat_history: chatHistory
      });

      // Replace loading message with actual response
      const updatedMessages = updatedChat.messages.map(msg =>
        msg.id === loadingMessage.id
          ? {
            ...msg,
            content: response.answer,
            isLoading: false,
            result: response.retrieved_sources,
            variations: [{ id: loadingMessage.id + "-original", content: response.answer, timestamp: new Date() }],
            activeVariation: loadingMessage.id + "-original"
          }
          : msg
      );

      const finalChat = {
        ...updatedChat,
        messages: updatedMessages,
      };

      setActiveChat(finalChat);

      // Update chats list
      setChats(prevChats =>
        prevChats.map(chat =>
          chat.id === finalChat.id ? finalChat : chat
        )
      );

      // Generate title if this is a new chat
      if (updatedChat.title === "New Chat" && updatedChat.messages.length <= 3) {
        generateChatTitle(userMessage.content);
      }

    } catch (error) {
      console.error("Error querying AI:", error);

      // Replace loading message with error
      const updatedMessages = updatedChat.messages.map(msg =>
        msg.id === loadingMessage.id
          ? {
            ...msg,
            content: "I'm sorry, I couldn't process your request. Please try again.",
            isLoading: false,
            error: true
          }
          : msg
      );

      setActiveChat({
        ...updatedChat,
        messages: updatedMessages
      });

      toast.error("Failed to process your request");
    } finally {
      setIsLoading(false);
      // Refocus the input after sending message
      setTimeout(() => {
        inputRef.current?.focus();
      }, 100);
    }
  };

  const createNewChat = () => {
    const newChat: Chat = {
      id: generateId(),
      title: "New Chat",
      messages: [
        {
          id: "welcomems",
          content: WELCOME_MESSAGE,
          sender: "system" as const,
          timestamp: new Date()
        }
      ],
      createdAt: new Date(),
      updatedAt: new Date()
    };

    setActiveChat(newChat);
    setChats(prev => [newChat, ...prev]);
    setTimeout(() => {
      inputRef.current?.focus();
    }, 100);
    setIsSidebarOpen(false);
  };

  const selectChat = (chatId: string) => {
    const selectedChat = chats.find(chat => chat.id === chatId);
    if (selectedChat) {
      setActiveChat(selectedChat);
      setIsSidebarOpen(false);
      setTimeout(() => {
        inputRef.current?.focus();
      }, 100);
    }
  };

  const deleteChat = (chatId: string) => {
    const updatedChats = chats.filter(chat => chat.id !== chatId);
    setChats(updatedChats);

    // If we're deleting the active chat, switch to another one
    if (activeChat?.id === chatId) {
      setActiveChat(updatedChats.length > 0 ? updatedChats[0] : null);

      // If no chats left, create a new one
      if (updatedChats.length === 0) {
        createNewChat();
      }
    }
    
    // Clear the chat being deleted
    setChatToDelete(null);
  };

  const handleDeleteMessage = (messageId: string) => {
    if (!activeChat) return;
    
    // Find the index of the message to delete
    const messageIndex = activeChat.messages.findIndex(msg => msg.id === messageId);
    if (messageIndex === -1) return;
    
    // Determine if we need to delete a pair (user message + assistant response)
    const isUserMessage = activeChat.messages[messageIndex].sender === "user";
    const updatedMessages = [...activeChat.messages];
    
    if (isUserMessage && messageIndex + 1 < updatedMessages.length && 
        updatedMessages[messageIndex + 1].sender === "system") {
      // Remove both the user message and the following assistant response
      updatedMessages.splice(messageIndex, 2);
    } else if (!isUserMessage && messageIndex > 0 && 
               updatedMessages[messageIndex - 1].sender === "user") {
      // Remove both the assistant message and the preceding user message
      updatedMessages.splice(messageIndex - 1, 2);
    } else {
      // Just remove the single message
      updatedMessages.splice(messageIndex, 1);
    }
    
    const updatedChat = {
      ...activeChat,
      messages: updatedMessages,
      updatedAt: new Date()
    };
    
    setActiveChat(updatedChat);
    
    // Update chats list
    setChats(prevChats =>
      prevChats.map(chat =>
        chat.id === updatedChat.id ? updatedChat : chat
      )
    );
  };

  const handleRegenerateMessage = (messageId: string) => {
    if (!activeChat) return;
    
    // Find the system message that needs to be regenerated
    const messageIndex = activeChat.messages.findIndex(
      msg => msg.id === messageId && msg.sender === "system"
    );
    
    if (messageIndex < 0) return;
    
    const message = activeChat.messages[messageIndex];
    
    // Find the last user message before this system message
    let userMessageContent = "";
    let userMessageIndex = -1;
    
    for (let i = messageIndex - 1; i >= 0; i--) {
      if (activeChat.messages[i].sender === "user") {
        userMessageContent = activeChat.messages[i].content;
        userMessageIndex = i;
        break;
      }
    }
    
    if (!userMessageContent) return;
    
    // Create a new variation
    const variationId = generateId();
    const now = new Date();
    
    // Get existing variations or initialize if none exist
    let existingVariations = message.variations || [];
    
    // If there are no variations yet, add the original content as the first variation
    if (existingVariations.length === 0) {
      existingVariations = [
        { id: message.id + "-original", content: message.content, timestamp: message.timestamp }
      ];
    }
    
    // Create a new variations array with the loading state
    const updatedVariations = [
      ...existingVariations,
      { id: variationId, content: "", timestamp: now }
    ];
    
    // Update the message with loading state
    const updatedMessages = [...activeChat.messages];
    updatedMessages[messageIndex] = {
      ...updatedMessages[messageIndex],
      isLoading: true,
      variations: updatedVariations,
      activeVariation: variationId
    };
    
    // Find and remove any messages that were children of the previous variation
    // We'll preserve the core tree but remove messages specific to other variations
    const messagesToKeep = updatedMessages.slice(0, messageIndex + 1);
    
    const updatedChat = {
      ...activeChat,
      messages: messagesToKeep
    };
    
    setActiveChat(updatedChat);
    setIsLoading(true);
    
    // Prepare chat history for the API - strip thinking content
    const chatHistory: APIMessage[] = activeChat.messages
      .slice(0, userMessageIndex)
      .filter(msg => !msg.isLoading && msg.content)
      .map(msg => {
        // Strip thinking content from messages
        const cleanedContent = msg.content.replace(/<think>[\s\S]*?<\/think>/g, '').trim();
        return {
          role: msg.sender === "user" ? "user" : "assistant",
          content: cleanedContent
        };
      });
    
    // Send the API request
    apiService.queryRulings({
      query: userMessageContent,
      chat_history: chatHistory
    })
      .then(response => {
        // Update the variation with the actual response
        const finalVariations = updatedMessages[messageIndex].variations!.map(v => 
          v.id === variationId 
            ? { ...v, content: response.answer, timestamp: new Date() } 
            : v
        );
        
        const finalMessages = [...messagesToKeep];
        finalMessages[messageIndex] = {
          ...finalMessages[messageIndex],
          variations: finalVariations,
          isLoading: false,
          error: false,
          result: response.retrieved_sources,
          activeVariation: variationId
        };
        
        const finalChat = {
          ...updatedChat,
          messages: finalMessages
        };
        
        setActiveChat(finalChat);
        setChats(prevChats =>
          prevChats.map(chat =>
            chat.id === finalChat.id ? finalChat : chat
          )
        );
      })
      .catch(error => {
        console.error("Error regenerating response:", error);
        
        // Remove the failed variation
        const finalVariations = updatedMessages[messageIndex].variations!.filter(v => 
          v.id !== variationId
        );
        
        const finalMessages = [...updatedMessages];
        finalMessages[messageIndex] = {
          ...finalMessages[messageIndex],
          variations: finalVariations,
          isLoading: false,
          activeVariation: finalVariations.length > 0 ? finalVariations[0].id : undefined
        };
        
        setActiveChat({
          ...updatedChat,
          messages: finalMessages
        });
        
        toast.error("Failed to generate variation");
      })
      .finally(() => {
        setIsLoading(false);
        setTimeout(() => {
          inputRef.current?.focus();
        }, 100);
      });
  };

  const handleSelectVariation = (messageId: string, variationId: string) => {
    if (!activeChat) return;

    // Find the message
    const messageIndex = activeChat.messages.findIndex(msg => msg.id === messageId);
    if (messageIndex < 0) return;
    
    // Update the active variation for this message
    const updatedMessages = [...activeChat.messages];
    updatedMessages[messageIndex] = {
      ...updatedMessages[messageIndex],
      activeVariation: variationId
    };
    
    // Keep messages up to and including the varied message
    const baseMessages = updatedMessages.slice(0, messageIndex + 1);
    
    // Find any existing messages that belong to this variation
    const childMessages = activeChat.messages
      .filter(msg => msg.parentMessageId === messageId && msg.variationId === variationId);
    
    // Combine to get the complete message list
    const finalMessages = [...baseMessages, ...childMessages];
    
    const updatedChat = {
      ...activeChat,
      messages: finalMessages
    };
    
    setActiveChat(updatedChat);
    
    // Update chats list
    setChats(prevChats =>
      prevChats.map(chat =>
        chat.id === updatedChat.id ? updatedChat : chat
      )
    );
  };

  const handleRetryMessage = (messageId: string) => {
    if (!activeChat) return;

    // Find the failed message
    const failedMessageIndex = activeChat.messages.findIndex(
      msg => msg.id === messageId && msg.error
    );

    if (failedMessageIndex < 0) return;

    // Get the last user message before this failed message
    let userMessageContent = "";
    for (let i = failedMessageIndex - 1; i >= 0; i--) {
      if (activeChat.messages[i].sender === "user") {
        userMessageContent = activeChat.messages[i].content;
        break;
      }
    }

    if (!userMessageContent) return;

    // Remove the failed message
    const updatedMessages = [...activeChat.messages];
    updatedMessages[failedMessageIndex] = {
      ...updatedMessages[failedMessageIndex],
      isLoading: true,
      error: false,
      content: ""
    };

    const updatedChat = {
      ...activeChat,
      messages: updatedMessages
    };

    setActiveChat(updatedChat);
    setIsLoading(true);

    // Prepare chat history for the API
    const chatHistory: APIMessage[] = updatedChat.messages
      .filter(msg => !msg.isLoading && msg.content && updatedChat.messages.indexOf(msg) < failedMessageIndex - 1)
      .map(msg => ({
        role: msg.sender === "user" ? "user" : "assistant",
        content: msg.content
      }));

    // Retry the query
    apiService.queryRulings({
      query: userMessageContent,
      chat_history: chatHistory
    })
      .then(response => {
        const finalMessages = [...updatedMessages];
        finalMessages[failedMessageIndex] = {
          ...finalMessages[failedMessageIndex],
          content: response.answer,
          isLoading: false,
          error: false,
          result: response.retrieved_sources
        };

        const finalChat = {
          ...updatedChat,
          messages: finalMessages
        };

        setActiveChat(finalChat);
        setChats(prevChats =>
          prevChats.map(chat =>
            chat.id === finalChat.id ? finalChat : chat
          )
        );
      })
      .catch(error => {
        console.error("Error retrying query:", error);

        const finalMessages = [...updatedMessages];
        finalMessages[failedMessageIndex] = {
          ...finalMessages[failedMessageIndex],
          content: "I'm sorry, I couldn't process your request. Please try again.",
          isLoading: false,
          error: true
        };

        setActiveChat({
          ...updatedChat,
          messages: finalMessages
        });

        toast.error("Failed to process your request");
      })
      .finally(() => {
        setIsLoading(false);
        setTimeout(() => {
          inputRef.current?.focus();
        }, 100);
      });
  };

  const openProfileModal = () => {
    setIsProfileModalOpen(true);
  };

  return (
    <>
      <div className="flex h-full w-full">
        {/* Chat Sidebar */}
        <ChatSidebar 
          chats={chats}
          activeChat={activeChat}
          isGeneratingTitle={isGeneratingTitle}
          createNewChat={createNewChat}
          selectChat={selectChat}
          onRequestDelete={setChatToDelete}
          onOpenSettings={onOpenSettings}
          onOpenSources={onOpenSources}
          openProfileModal={openProfileModal}
          isSidebarOpen={isSidebarOpen}
          setIsSidebarOpen={setIsSidebarOpen}
        />

        {/* Chat Main Area */}
        <div className="flex-1 flex flex-col overflow-hidden">
          {/* Mobile Header */}
          <div className="md:hidden p-2 flex items-center border-b">
            <Button variant="ghost" size="icon" onClick={() => setIsSidebarOpen(true)}>
              <PanelLeft className="h-5 w-5" />
            </Button>
            <div className="mx-auto font-medium flex items-center">
              <Bot className="h-4 w-4 mr-1.5" />
              {activeChat?.title || "New Chat"}
            </div>
            <Button variant="ghost" size="icon" onClick={createNewChat}>
              <Plus className="h-5 w-5" />
            </Button>
          </div>

          {/* Chat Area */}
          <div className="relative flex-1 overflow-hidden">
            {!activeChat ? (
              <WelcomeScreen onCreateNewChat={createNewChat} />
            ) : (
              <>
                <div className="h-full overflow-y-auto px-4 pb-32 pt-4">
                  <div className="max-w-3xl mx-auto space-y-4">
                    {activeChat.messages.map((message) => (
                      <ChatBubble
                        key={message.id}
                        message={message}
                        onViewSearchResults={onOpenSources}
                        onRetry={handleRetryMessage}
                        onRegenerate={handleRegenerateMessage}
                        onDelete={handleDeleteMessage}
                        onSelectVariation={handleSelectVariation}
                      />
                    ))}
                    <div ref={messagesEndRef} />
                  </div>
                </div>

                {/* Input Area */}
                <ChatInputArea
                  inputRef={inputRef}
                  inputValue={inputValue}
                  setInputValue={setInputValue}
                  handleSendMessage={handleSendMessage}
                  isLoading={isLoading}
                />
              </>
            )}
          </div>
        </div>
      </div>
      
      {/* Profile Modal */}
      <ProfileModal isOpen={isProfileModalOpen} onClose={() => setIsProfileModalOpen(false)} />
      
      {/* Delete Chat Confirmation Dialog */}
      <DeleteChatDialog
        isOpen={chatToDelete !== null}
        onOpenChange={() => setChatToDelete(null)}
        onDelete={() => chatToDelete && deleteChat(chatToDelete)}
      />
    </>
  );
};

export default ChatInterface;