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,
  Database,
  Loader2,
} from "lucide-react";
import KnowledgeManagement from "./KnowledgeManagement";
import {
  getRooms,
  createRoom,
  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;
}

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

  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,
      }));
      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 currentChat = chats.find((chat) => chat.id === currentChatId);

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

  const deleteChat = (chatId: string) => {
    const updatedChats = chats.filter((chat) => chat.id !== chatId);
    setChats(updatedChats);
    if (currentChatId === chatId) {
      setCurrentChatId(updatedChats.length > 0 ? updatedChats[0].id : null);
    }
  };

  const deleteAllChats = () => {
    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,
        };
        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();

    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("\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:", "").trim();

            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
                )
              );
            }
          }
        }
      }
    } 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:
                            "Error: Failed to get response. Please try again.",
                        }
                      : m
                  ),
                }
              : chat
          )
        );
      }
    } finally {
      setIsStreaming(false);
      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-br from-[#00C853] to-[#00A843] 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">
        {/* 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">
            {currentChat?.title || "Chatbot"}
          </h1>
          <button
            onClick={() => setKnowledgeOpen(true)}
            className="flex items-center gap-2 bg-gradient-to-r from-[#FF8F00] to-[#FF6F00] text-white px-3 py-2 rounded-lg hover:from-[#FF6F00] hover:to-[#F57C00] transition text-sm"
          >
            <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"
              }`}
            >
              <div
                className={`max-w-2xl px-4 py-2.5 rounded-xl ${
                  message.role === "user"
                    ? "bg-gradient-to-r from-[#4FC3F7] to-[#29B6F6] text-white"
                    : "bg-white border border-slate-200 text-slate-900"
                }`}
              >
                <p className="whitespace-pre-wrap break-words text-sm">
                  {message.content}
                </p>
                {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">
                        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-1.5 py-0.5 rounded"
                            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">
          <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="Type your message..."
                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-[#4FC3F7] focus:border-transparent resize-none max-h-32"
                disabled={isStreaming}
              />
              <button
                onClick={handleSend}
                disabled={!input.trim() || isStreaming}
                className="bg-gradient-to-r from-[#4FC3F7] to-[#29B6F6] text-white p-2.5 rounded-lg hover:from-[#29B6F6] hover:to-[#039BE5] transition disabled:opacity-50 disabled:cursor-not-allowed"
              >
                <Send className="w-4 h-4" />
              </button>
            </div>
          </div>
        </div>
      </div>

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