src/app/components/Main.tsx

import { useState, useEffect, useRef } from "react";
import { useNavigate } from "react-router";
import {
  Send,
  Plus,
  Trash2,
  LogOut,
  Menu,
  X,
  MessageSquare,
  User,
  Bot,
  Loader2,
  Database,
} from "lucide-react";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";
import remarkMath from "remark-math";
import rehypeKatex from "rehype-katex";
import type { Components } from "react-markdown";
import KnowledgeManagement from "./KnowledgeManagement";
import {
  getRooms,
  getRoom,
  createRoom,
  deleteRoom,
  streamChat,
  type ChatSource,
} from "../../services/api";

interface StoredUser {
  user_id: string;
  email: string;
  name: string;
  loginTime: string;
}

interface Message {
  id: string;
  role: "user" | "assistant";
  content: string;
  timestamp: number;
  sources?: ChatSource[];
}

interface ChatRoom {
  id: string;
  title: string;
  messages: Message[];
  createdAt: string;
  updatedAt: string | null;
  messagesLoaded: boolean;
}

// Markdown component overrides for clean rendering inside chat bubbles
const markdownComponents: Components = {
  p: ({ children }) => (
    <p className="text-sm mb-2 last:mb-0 leading-relaxed">{children}</p>
  ),
  h1: ({ children }) => (
    <h1 className="text-lg font-bold mb-3 mt-4 first:mt-0">{children}</h1>
  ),
  h2: ({ children }) => (
    <h2 className="text-base font-bold mb-2 mt-3 first:mt-0">{children}</h2>
  ),
  h3: ({ children }) => (
    <h3 className="text-sm font-semibold mb-2 mt-2 first:mt-0">{children}</h3>
  ),
  ul: ({ children }) => (
    <ul className="list-disc pl-5 mb-2 space-y-1 text-sm">{children}</ul>
  ),
  ol: ({ children }) => (
    <ol className="list-decimal pl-5 mb-2 space-y-1 text-sm">{children}</ol>
  ),
  li: ({ children }) => <li className="text-sm leading-relaxed">{children}</li>,
  code: ({ children, className }) => {
    const isBlock = className?.startsWith("language-");
    if (isBlock) {
      return (
        <code className="block text-xs font-mono text-slate-100 leading-relaxed">
          {children}
        </code>
      );
    }
    return (
      <code className="bg-slate-100 text-pink-600 px-1.5 py-0.5 rounded text-xs font-mono">
        {children}
      </code>
    );
  },
  pre: ({ children }) => (
    <pre className="bg-slate-900 rounded-lg p-3 mb-2 mt-1 overflow-x-auto text-xs">
      {children}
    </pre>
  ),
  blockquote: ({ children }) => (
    <blockquote className="border-l-4 border-slate-300 pl-3 text-slate-500 italic mb-2 text-sm">
      {children}
    </blockquote>
  ),
  table: ({ children }) => (
    <div className="overflow-x-auto mb-2">
      <table className="w-full text-sm border-collapse">{children}</table>
    </div>
  ),
  th: ({ children }) => (
    <th className="border border-slate-200 px-3 py-1.5 bg-slate-100 font-medium text-left text-xs">
      {children}
    </th>
  ),
  td: ({ children }) => (
    <td className="border border-slate-200 px-3 py-1.5 text-xs">{children}</td>
  ),
  a: ({ children, href }) => (
    <a
      href={href}
      target="_blank"
      rel="noopener noreferrer"
      className="text-blue-600 underline hover:text-blue-800"
    >
      {children}
    </a>
  ),
  strong: ({ children }) => (
    <strong className="font-semibold">{children}</strong>
  ),
  hr: () => <hr className="border-slate-200 my-3" />,
};

// Typing indicator — three bouncing dots
function useLoadingMessages() {
  const [messages, setMessages] = useState<string[]>([]);

  useEffect(() => {
    fetch("/loading-messages.yaml")
      .then((r) => r.text())
      .then((text) => {
        const parsed = text
          .split("\n")
          .filter((l) => l.trimStart().startsWith("- "))
          .map((l) => l.replace(/^\s*- /, "").trim())
          .filter(Boolean);
        if (parsed.length > 0) setMessages(parsed);
      })
      .catch(() => {});
  }, []);

  return messages;
}

function TypingIndicator() {
  const messages = useLoadingMessages();
  const [index, setIndex] = useState(0);

  useEffect(() => {
    if (messages.length === 0) return;
    setIndex(Math.floor(Math.random() * messages.length));
    const id = setInterval(() => {
      setIndex((prev) => {
        let next: number;
        do { next = Math.floor(Math.random() * messages.length); } while (messages.length > 1 && next === prev);
        return next;
      });
    }, 300);
    return () => clearInterval(id);
  }, [messages]);

  if (messages.length === 0) {
    return (
      <div className="flex gap-1.5 items-center py-1 px-0.5">
        <span className="w-2 h-2 bg-slate-400 rounded-full animate-bounce" style={{ animationDelay: "0ms", animationDuration: "1s" }} />
        <span className="w-2 h-2 bg-slate-400 rounded-full animate-bounce" style={{ animationDelay: "200ms", animationDuration: "1s" }} />
        <span className="w-2 h-2 bg-slate-400 rounded-full animate-bounce" style={{ animationDelay: "400ms", animationDuration: "1s" }} />
      </div>
    );
  }

  return (
    <div className="flex items-center gap-2 py-1 px-0.5 text-sm text-slate-400 italic">
      <span className="inline-block w-2 h-2 rounded-full bg-slate-400 animate-pulse" />
      <span>{messages[index]}…</span>
    </div>
  );
}

export default function Main() {
  const navigate = useNavigate();
  const [sidebarOpen, setSidebarOpen] = useState(true);
  const [chats, setChats] = useState<ChatRoom[]>([]);
  const [currentChatId, setCurrentChatId] = useState<string | null>(null);
  const [input, setInput] = useState("");
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingMsgId, setStreamingMsgId] = useState<string | null>(null);
  const [roomsLoading, setRoomsLoading] = useState(false);
  const [roomsError, setRoomsError] = useState<string | null>(null);
  const messagesEndRef = useRef<HTMLDivElement>(null);
  const [user, setUser] = useState<StoredUser | null>(null);
  const [knowledgeOpen, setKnowledgeOpen] = useState(false);
  const abortControllerRef = useRef<AbortController | null>(null);

  useEffect(() => {
    const storedUser = localStorage.getItem("chatbot_user");
    if (storedUser) {
      const parsedUser: StoredUser = JSON.parse(storedUser);
      setUser(parsedUser);
      loadRooms(parsedUser.user_id);
    }
  }, []);

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

  useEffect(() => {
    if (!currentChatId) return;
    const chat = chats.find((c) => c.id === currentChatId);
    if (chat && !chat.messagesLoaded) {
      loadRoomMessages(currentChatId);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [currentChatId]);

  const loadRooms = async (userId: string) => {
    setRoomsLoading(true);
    setRoomsError(null);
    try {
      const apiRooms = await getRooms(userId);
      const mapped: ChatRoom[] = apiRooms.map((r) => ({
        id: r.id,
        title: r.title,
        messages: [],
        createdAt: r.created_at,
        updatedAt: r.updated_at,
        messagesLoaded: false,
      }));
      setChats(mapped);
      if (mapped.length > 0) {
        setCurrentChatId(mapped[0].id);
      }
    } catch (err) {
      setRoomsError(
        err instanceof Error ? err.message : "Failed to load chats"
      );
    } finally {
      setRoomsLoading(false);
    }
  };

  const loadRoomMessages = async (roomId: string) => {
    try {
      const detail = await getRoom(roomId);
      const messages: Message[] = detail.messages.map((m) => ({
        id: m.id,
        role: m.role,
        content: m.content,
        timestamp: new Date(m.created_at).getTime(),
      }));
      setChats((prev) =>
        prev.map((chat) =>
          chat.id === roomId
            ? { ...chat, messages, messagesLoaded: true }
            : chat
        )
      );
    } catch {
      setChats((prev) =>
        prev.map((chat) =>
          chat.id === roomId ? { ...chat, messagesLoaded: true } : chat
        )
      );
    }
  };

  const currentChat = chats.find((chat) => chat.id === currentChatId);

  const createNewChat = () => {
    setCurrentChatId(null);
  };

  const deleteChat = async (chatId: string) => {
    if (!user) return;
    try {
      await deleteRoom(chatId, user.user_id);
    } catch {
      return;
    }
    const updatedChats = chats.filter((chat) => chat.id !== chatId);
    setChats(updatedChats);
    if (currentChatId === chatId) {
      setCurrentChatId(updatedChats.length > 0 ? updatedChats[0].id : null);
    }
  };

  const deleteAllChats = async () => {
    if (!user) return;
    await Promise.allSettled(
      chats.map((chat) => deleteRoom(chat.id, user.user_id))
    );
    setChats([]);
    setCurrentChatId(null);
  };

  const handleLogout = () => {
    localStorage.removeItem("chatbot_user");
    navigate("/login");
  };

  const handleSend = async () => {
    if (!input.trim() || isStreaming || !user) return;

    let roomId = currentChatId;

    if (!roomId) {
      try {
        const res = await createRoom(user.user_id, input.slice(0, 50));
        const newRoom: ChatRoom = {
          id: res.data.id,
          title: res.data.title,
          messages: [],
          createdAt: res.data.created_at,
          updatedAt: res.data.updated_at,
          messagesLoaded: true,
        };
        setChats((prev) => [newRoom, ...prev]);
        roomId = newRoom.id;
        setCurrentChatId(roomId);
      } catch {
        return;
      }
    }

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

    setChats((prev) =>
      prev.map((chat) =>
        chat.id === roomId
          ? {
              ...chat,
              messages: [...chat.messages, userMessage],
              updatedAt: new Date().toISOString(),
            }
          : chat
      )
    );

    const sentMessage = input;
    setInput("");
    setIsStreaming(true);

    const assistantMsgId = crypto.randomUUID();
    setStreamingMsgId(assistantMsgId);

    setChats((prev) =>
      prev.map((chat) =>
        chat.id === roomId
          ? {
              ...chat,
              messages: [
                ...chat.messages,
                {
                  id: assistantMsgId,
                  role: "assistant",
                  content: "",
                  timestamp: Date.now(),
                  sources: [],
                },
              ],
            }
          : chat
      )
    );

    abortControllerRef.current = new AbortController();

    try {
      const response = await streamChat(user.user_id, roomId, sentMessage);

      if (!response.body) throw new Error("No response body");

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = "";
      let currentEvent = "";

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split(/\r?\n/);
        buffer = lines.pop() ?? "";

        for (const line of lines) {
          if (line.startsWith("event:")) {
            currentEvent = line.replace("event:", "").trim();
          } else if (line.startsWith("data:")) {
            const data = line.replace(/^data: ?/, "");

            if (currentEvent === "sources" && data) {
              const sources: ChatSource[] = JSON.parse(data);
              setChats((prev) =>
                prev.map((chat) =>
                  chat.id === roomId
                    ? {
                        ...chat,
                        messages: chat.messages.map((m) =>
                          m.id === assistantMsgId ? { ...m, sources } : m
                        ),
                      }
                    : chat
                )
              );
            } else if (currentEvent === "chunk" && data) {
              setChats((prev) =>
                prev.map((chat) =>
                  chat.id === roomId
                    ? {
                        ...chat,
                        messages: chat.messages.map((m) =>
                          m.id === assistantMsgId
                            ? { ...m, content: m.content + data }
                            : m
                        ),
                      }
                    : chat
                )
              );
            } else if (currentEvent === "message" && data) {
              setChats((prev) =>
                prev.map((chat) =>
                  chat.id === roomId
                    ? {
                        ...chat,
                        messages: chat.messages.map((m) =>
                          m.id === assistantMsgId
                            ? { ...m, content: data }
                            : m
                        ),
                      }
                    : chat
                )
              );
            } else if (currentEvent === "done") {
              break;
            }
          }
        }
      }
    } catch (err: unknown) {
      if ((err as Error).name !== "AbortError") {
        setChats((prev) =>
          prev.map((chat) =>
            chat.id === roomId
              ? {
                  ...chat,
                  messages: chat.messages.map((m) =>
                    m.id === assistantMsgId
                      ? {
                          ...m,
                          content:
                            "Sorry, I couldn't get a response. Please try again.",
                        }
                      : m
                  ),
                }
              : chat
          )
        );
      }
    } finally {
      setIsStreaming(false);
      setStreamingMsgId(null);
      abortControllerRef.current = null;
    }
  };

  const handleKeyPress = (e: React.KeyboardEvent) => {
    if (e.key === "Enter" && !e.shiftKey) {
      e.preventDefault();
      handleSend();
    }
  };

  return (
    <div className="flex h-screen bg-slate-50">
      {/* Sidebar */}
      <div
        className={`${
          sidebarOpen ? "w-64" : "w-0"
        } bg-gradient-to-b from-[#059669] to-[#047857] text-white transition-all duration-300 flex flex-col overflow-hidden`}
      >
        <div className="p-3 border-b border-white/20">
          <button
            onClick={createNewChat}
            className="w-full flex items-center justify-center gap-2 bg-white/20 hover:bg-white/30 px-3 py-2 rounded-lg transition text-sm"
          >
            <Plus className="w-4 h-4" />
            New Chat
          </button>
        </div>

        <div className="flex-1 overflow-y-auto p-2 space-y-1">
          {roomsLoading ? (
            <div className="flex justify-center py-4">
              <Loader2 className="w-4 h-4 animate-spin text-white/70" />
            </div>
          ) : roomsError ? (
            <p className="text-xs text-red-200 text-center px-2 py-2">
              {roomsError}
            </p>
          ) : (
            chats.map((chat) => (
              <div
                key={chat.id}
                className={`flex items-center gap-2 p-2 rounded-lg cursor-pointer transition group ${
                  currentChatId === chat.id
                    ? "bg-white/25"
                    : "hover:bg-white/15"
                }`}
              >
                <MessageSquare className="w-3.5 h-3.5 flex-shrink-0" />
                <div
                  className="flex-1 truncate text-sm"
                  onClick={() => setCurrentChatId(chat.id)}
                >
                  {chat.title}
                </div>
                <button
                  onClick={(e) => {
                    e.stopPropagation();
                    deleteChat(chat.id);
                  }}
                  className="opacity-0 group-hover:opacity-100 transition"
                >
                  <Trash2 className="w-3.5 h-3.5 text-red-100 hover:text-white" />
                </button>
              </div>
            ))
          )}
        </div>

        <div className="border-t border-white/20 p-3 space-y-2">
          {chats.length > 0 && (
            <button
              onClick={deleteAllChats}
              className="w-full flex items-center justify-center gap-2 text-red-100 hover:text-white px-3 py-2 rounded-lg hover:bg-white/15 transition text-xs"
            >
              <Trash2 className="w-3.5 h-3.5" />
              Clear All Chats
            </button>
          )}

          <div className="flex items-center gap-2 p-2 rounded-lg bg-white/20">
            <div className="w-7 h-7 bg-white/30 rounded-full flex items-center justify-center">
              <User className="w-3.5 h-3.5" />
            </div>
            <div className="flex-1 min-w-0">
              <div className="text-xs truncate">{user?.name}</div>
              <div className="text-[10px] text-white/70 truncate">
                {user?.email}
              </div>
            </div>
            <button
              onClick={handleLogout}
              className="text-white/70 hover:text-white transition"
              title="Logout"
            >
              <LogOut className="w-3.5 h-3.5" />
            </button>
          </div>
        </div>
      </div>

      {/* Main Content */}
      <div className="flex-1 flex flex-col min-w-0">
        {/* Header */}
        <div className="bg-white border-b border-slate-200 p-3 flex items-center gap-3">
          <button
            onClick={() => setSidebarOpen(!sidebarOpen)}
            className="text-slate-600 hover:text-slate-900 transition"
          >
            {sidebarOpen ? (
              <X className="w-5 h-5" />
            ) : (
              <Menu className="w-5 h-5" />
            )}
          </button>
          <h1 className="text-base text-slate-900 flex-1 truncate">
            {currentChat?.title || "Chatbot"}
          </h1>
          <button
            onClick={() => setKnowledgeOpen(true)}
            className="flex items-center gap-2 bg-[#F59E0B] hover:bg-[#D97706] text-white px-3 py-2 rounded-lg transition-all duration-200 hover:scale-105 text-sm flex-shrink-0"
          >
            <Database className="w-4 h-4" />
            Knowledge
          </button>
        </div>

        {/* Messages */}
        <div className="flex-1 overflow-y-auto p-4 space-y-4">
          {currentChat?.messages.length === 0 && (
            <div className="flex items-center justify-center h-full">
              <div className="text-center">
                <MessageSquare className="w-12 h-12 text-slate-300 mx-auto mb-3" />
                <h2 className="text-base text-slate-600 mb-1">
                  Start a conversation
                </h2>
                <p className="text-sm text-slate-400">
                  Send a message to begin chatting
                </p>
              </div>
            </div>
          )}

          {currentChat?.messages.map((message) => (
            <div
              key={message.id}
              className={`flex ${
                message.role === "user" ? "justify-end" : "justify-start"
              }`}
            >
              {/* Avatar for assistant */}
              {message.role === "assistant" && (
                <div className="w-7 h-7 rounded-full bg-gradient-to-br from-[#059669] to-[#047857] flex items-center justify-center flex-shrink-0 mt-0.5 mr-2">
                  <Bot className="w-3.5 h-3.5 text-white" />
                </div>
              )}

              <div
                className={`max-w-2xl px-4 py-3 rounded-2xl ${
                  message.role === "user"
                    ? "bg-[#3B82F6] text-white rounded-tr-sm shadow-sm"
                    : "bg-[#F3F4F6] border-0 text-slate-900 rounded-tl-sm shadow-sm"
                }`}
              >
                {message.role === "user" ? (
                  <p className="text-sm whitespace-pre-wrap break-words leading-relaxed">
                    {message.content}
                  </p>
                ) : message.content === "" && streamingMsgId === message.id ? (
                  // Waiting for first chunk — show typing indicator
                  <TypingIndicator />
                ) : (
                  // Render markdown for assistant messages
                  <div className="text-slate-900">
                    <ReactMarkdown
                      remarkPlugins={[remarkGfm, remarkMath]}
                      rehypePlugins={[rehypeKatex]}
                      components={markdownComponents}
                    >
                      {message.content}
                    </ReactMarkdown>
                  </div>
                )}

                {/* Sources */}
                {message.role === "assistant" &&
                  message.sources &&
                  message.sources.length > 0 && (
                    <div className="mt-2 pt-2 border-t border-slate-100">
                      <p className="text-[10px] text-slate-400 mb-1.5">
                        Sources:
                      </p>
                      <div className="flex flex-wrap gap-1">
                        {message.sources.map((src, i) => (
                          <span
                            key={i}
                            className="text-[10px] bg-slate-100 text-slate-600 px-2 py-0.5 rounded-full border border-slate-200"
                            title={
                              src.page_label
                                ? `Page ${src.page_label}`
                                : undefined
                            }
                          >
                            📄 {src.filename}
                            {src.page_label ? ` p.${src.page_label}` : ""}
                          </span>
                        ))}
                      </div>
                    </div>
                  )}
              </div>
            </div>
          ))}

          {!currentChat && chats.length === 0 && !roomsLoading && (
            <div className="flex items-center justify-center h-full">
              <div className="text-center">
                <MessageSquare className="w-12 h-12 text-slate-300 mx-auto mb-3" />
                <h2 className="text-base text-slate-600 mb-1">
                  Welcome to Chatbot
                </h2>
                <p className="text-sm text-slate-400">
                  Create a new chat to get started
                </p>
              </div>
            </div>
          )}

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

        {/* Input Area */}
        <div className="bg-white border-t border-slate-200 p-3 shadow-[0_-2px_10px_rgba(0,0,0,0.06)]">
          <div className="max-w-4xl mx-auto">
            <div className="flex gap-2 items-end">
              <textarea
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={handleKeyPress}
                placeholder="Ask me anything... (Enter to send, Shift+Enter for newline)"
                rows={1}
                className="flex-1 px-3 py-2 text-sm rounded-lg border border-slate-300 focus:outline-none focus:ring-2 focus:ring-[#3B82F6] focus:border-transparent resize-none max-h-32"
                disabled={isStreaming}
              />
              <button
                onClick={handleSend}
                disabled={!input.trim() || isStreaming}
                className="bg-[#3B82F6] hover:bg-[#2563EB] text-white p-2.5 rounded-lg transition-all duration-200 hover:scale-105 disabled:opacity-50 disabled:cursor-not-allowed flex-shrink-0"
              >
                {isStreaming ? (
                  <Loader2 className="w-4 h-4 animate-spin" />
                ) : (
                  <Send className="w-4 h-4" />
                )}
              </button>
            </div>
          </div>
        </div>
      </div>

      {/* Knowledge Management Modal */}
      <KnowledgeManagement
        open={knowledgeOpen}
        onClose={() => setKnowledgeOpen(false)}
      />
    </div>
  );
}