src/app/components/Main.tsx

import { useState, useEffect, useRef, useCallback } from "react";
import { useNavigate } from "react-router";
import { Database, Menu } from "lucide-react";
import { AnimatePresence } from "motion/react";
import KnowledgeManagement from "./KnowledgeManagement";
import {
  getRooms,
  getRoom,
  createRoom,
  deleteRoom,
  streamChat,
  type ChatSource,
} from "../../services/api";
import { textToSpeech } from "../../services/voiceApi";
import { replayAudio } from "../../audio/AudioPlayer";
import ChatLayout from "./chat/ChatLayout";
import Sidebar from "./chat/Sidebar";
import ChatWindow from "./chat/ChatWindow";
import ChatInput from "./chat/ChatInput";
import VoiceStatusBar from "./chat/VoiceStatusBar";
import { useVoiceSession } from "../../hooks/useVoiceSession";
import type { VoiceState } from "../../hooks/useVoiceSession";
import type { Message, ChatSession, StoredUser } from "./chat/types";

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

export default function Main() {
  const navigate = useNavigate();
  const [chats, setChats] = useState<ChatRoom[]>([]);
  const [currentChatId, setCurrentChatId] = useState<string | null>(null);
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingMsgId, setStreamingMsgId] = useState<string | null>(null);
  const [roomsLoading, setRoomsLoading] = useState(false);
  const [user, setUser] = useState<StoredUser | null>(null);
  const [knowledgeOpen, setKnowledgeOpen] = useState(false);
  const [mobileSidebarOpen, setMobileSidebarOpen] = useState(false);
  const abortControllerRef = useRef<AbortController | null>(null);
  const isVoiceActiveRef = useRef(false);
  const setVoiceStateRef = useRef<((s: VoiceState) => void) | null>(null);

  // Stable refs so voice callbacks always see the latest values
  const currentChatIdRef = useRef<string | null>(null);
  useEffect(() => { currentChatIdRef.current = currentChatId; }, [currentChatId]);

  const userRef = useRef<StoredUser | null>(null);
  useEffect(() => { userRef.current = user; }, [user]);

  useEffect(() => {
    if (currentChatId) {
      localStorage.setItem("chatbot_last_room_id", currentChatId);
    }
  }, [currentChatId]);

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

  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);
    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) {
        const lastRoomId = localStorage.getItem("chatbot_last_room_id");
        const restoredId = lastRoomId && mapped.find((r) => r.id === lastRoomId)
          ? lastRoomId
          : mapped[0].id;
        setCurrentChatId(restoredId);
      }
    } catch {
      // silently fail — UI shows empty state
    } finally {
      setRoomsLoading(false);
    }
  };

  const loadRoomMessages = async (roomId: string): Promise<Message[]> => {
    try {
      const detail = await getRoom(roomId);
      const messages: Message[] = detail.messages.map((m) => ({
        id: m.id,
        role: m.role,
        content: m.content,
        audioText: m.audio_text ?? undefined,
        timestamp: new Date(m.created_at).getTime(),
        sources: m.sources ?? [],
      }));
      const asstMsgs = messages.filter(m => m.role === "assistant");
      const lastAsst = asstMsgs[asstMsgs.length - 1];
      if (lastAsst) {
        const hasCR = lastAsst.content.includes("\r");
        // console.log("[loadRoomMessages] last assistant contentLen=", lastAsst.content.length, "hasCR=", hasCR);
        // console.log("[loadRoomMessages] CONTENT JSON →", JSON.stringify(lastAsst.content.slice(0, 600)));
      }
      // console.log("[loadRoomMessages] room", roomId, "→", messages.length, "messages from server");
      setChats((prev) =>
        prev.map((chat) =>
          chat.id === roomId ? { ...chat, messages, messagesLoaded: true } : chat
        )
      );
      return messages;
    } catch (err) {
      console.error("[loadRoomMessages] failed:", err);
      setChats((prev) =>
        prev.map((chat) =>
          chat.id === roomId ? { ...chat, messagesLoaded: true } : chat
        )
      );
      return [];
    }
  };

  // Ensures a chat room exists; creates one for voice if needed.
  const ensureRoom = useCallback(async (title?: string): Promise<string | null> => {
    if (currentChatIdRef.current) return currentChatIdRef.current;
    const currentUser = userRef.current;
    if (!currentUser) return null;
    try {
      const res = await createRoom(currentUser.user_id, title ?? "Voice Session");
      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]);
      setCurrentChatId(newRoom.id);
      return newRoom.id;
    } catch {
      return null;
    }
  }, []);

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

  const handleNewChat = () => setCurrentChatId(null);

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

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

  const handleStop = () => {
    abortControllerRef.current?.abort();
  };

  const handleSend = useCallback(async (text: string, skipReload = false) => {
    // console.log("[handleSend] called, user:", user?.user_id ?? "null", "text:", text.slice(0, 40));
    if (!user) {
      console.warn("[handleSend] early return: no user");
      return;
    }

    let roomId = await ensureRoom(text.slice(0, 50));
    // console.log("[handleSend] roomId:", roomId);
    if (!roomId) {
      // console.warn("[handleSend] early return: no roomId");
      return;
    }

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

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

    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: [] as ChatSource[],
                },
              ],
            }
          : chat
      )
    );

    abortControllerRef.current = new AbortController();

    let audioText = "";

    try {
      const response = await streamChat(user.user_id, roomId, text);
      if (!response.body) throw new Error("No response body");

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = "";
      let currentEvent = "";
      let currentDataLines: string[] = [];
      let streamDone = false;

      const dispatchEvent = (eventType: string, data: string) => {
        if (eventType === "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 (eventType === "chunk") {
          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 (eventType === "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 (eventType === "audio_text" && data) {
          audioText = data;
        } else if (eventType === "done") {
          streamDone = true;
        }
      };

      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 === "") {
            // Blank line = end of SSE event — dispatch with all accumulated data lines joined by \n
            if (currentEvent) {
              dispatchEvent(currentEvent, currentDataLines.join("\n"));
            }
            currentEvent = "";
            currentDataLines = [];
          } else if (line.startsWith("event:")) {
            currentEvent = line.replace("event:", "").trim();
          } else if (line.startsWith("data:")) {
            currentDataLines.push(line.replace(/^data: ?/, ""));
          }
        }
        if (streamDone) break;
      }

      if (skipReload) {
        // Voice mode: skip GET /room, return immediately so TTS can start without delay.
        // audioText and audioChunks will be stored on the temp assistantMsgId.
        if (audioText) {
          setChats((prev) =>
            prev.map((chat) =>
              chat.id === roomId
                ? {
                    ...chat,
                    messages: chat.messages.map((m) =>
                      m.id === assistantMsgId ? { ...m, audioText } : m
                    ),
                  }
                : chat
            )
          );
        }
        return { audioText, assistantMsgId };
      }

      if (audioText) {
        setChats((prev) =>
          prev.map((chat) =>
            chat.id === roomId
              ? {
                  ...chat,
                  messages: chat.messages.map((m) =>
                    m.id === assistantMsgId ? { ...m, audioText } : m
                  ),
                }
              : chat
          )
        );
      }
      return { audioText, assistantMsgId };
    } catch (err: unknown) {
      if ((err as Error).name !== "AbortError") {
        console.error("[handleSend] streamChat error:", err);
        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
          )
        );
        audioText = "";
      }
    } finally {
      // console.log("[handleSend] finally: clearing streamingMsgId →", assistantMsgId, "| skipReload:", skipReload);
      setIsStreaming(false);
      setStreamingMsgId(null);
      abortControllerRef.current = null;
    }
  }, [user, ensureRoom]);

  const cancelTtsRef = useRef<(() => void) | null>(null);

  const playTtsAudio = useCallback(async (
    ttsText: string,
    onStarted?: () => void,
  ): Promise<{ chunks: ArrayBuffer[]; sampleRate: number }> => {
    // Stop any in-flight TTS before starting a new one (e.g. rapid re-query).
    cancelTtsRef.current?.();
    cancelTtsRef.current = null;
    try {
      const { pcm, sampleRate } = await textToSpeech(ttsText);
      const durationMs = (pcm.byteLength / 2 / sampleRate) * 1000;
      // Use the same proven playback path as the speaker button: schedule the
      // full response at once via replayAudio, avoiding all streaming complexity.
      const cancel = replayAudio([pcm], sampleRate);
      cancelTtsRef.current = cancel;
      onStarted?.();
      await new Promise<void>((resolve) => setTimeout(resolve, durationMs + 150));
      cancel();
      cancelTtsRef.current = null;
      return { chunks: [pcm], sampleRate };
    } catch {
      // TTS failure is non-fatal
      return { chunks: [], sampleRate: 24000 };
    }
  }, []);

  const { voiceState, start, stop, stopRecording, setStateExternal, isActive: isVoiceActive } = useVoiceSession({
    onTranscript: async (text: string) => {
      // console.log("[onTranscript] received:", text);
      // Pass skipReload=true so handleSend returns immediately after streaming
      // without waiting for GET /room — TTS starts with zero extra delay.
      const result = await handleSend(text, true);
      const { audioText = "", assistantMsgId } = result ?? {};
      // console.log("[onTranscript] handleSend done, audioText:", audioText ? audioText.slice(0, 40) : "(empty)");
      if (audioText && isVoiceActiveRef.current) {
        const { chunks, sampleRate } = await playTtsAudio(audioText, () => {
          if (isVoiceActiveRef.current) setVoiceStateRef.current?.("SPEAKING");
        });
        if (assistantMsgId && chunks.length > 0) {
          setChats((prev) =>
            prev.map((chat) =>
              chat.id === currentChatIdRef.current
                ? {
                    ...chat,
                    messages: chat.messages.map((m) =>
                      m.id === assistantMsgId
                        ? { ...m, audioChunks: chunks, audioSampleRate: sampleRate }
                        : m
                    ),
                  }
                : chat
            )
          );
        }
      }
      if (isVoiceActiveRef.current) setVoiceStateRef.current?.("IDLE");
    },
    sessionParams: {},
  });

  // Keep refs in sync with latest values
  useEffect(() => { isVoiceActiveRef.current = isVoiceActive; }, [isVoiceActive]);
  useEffect(() => { setVoiceStateRef.current = setStateExternal; }, [setStateExternal]);

  const handleVoiceToggle = useCallback(() => {
    if (!isVoiceActive) {
      start();
    } else if (voiceState === "LISTENING") {
      stopRecording();
    } else {
      stop();
    }
  }, [isVoiceActive, voiceState, start, stop, stopRecording]);

  const sessions: ChatSession[] = chats.map((c) => ({
    id: c.id,
    title: c.title,
    createdAt: c.createdAt,
    updatedAt: c.updatedAt,
  }));

  return (
    <ChatLayout
      sidebar={
        <Sidebar
          sessions={sessions}
          activeSessionId={currentChatId}
          isLoadingSessions={roomsLoading}
          user={user}
          onNewChat={handleNewChat}
          onSelectSession={setCurrentChatId}
          onDeleteSession={handleDeleteSession}
          onLogout={handleLogout}
          mobileOpen={mobileSidebarOpen}
          onMobileClose={() => setMobileSidebarOpen(false)}
        />
      }
    >
      {/* Background decoration */}
      <div className="absolute inset-0 pointer-events-none z-0 overflow-hidden">
        <div
          className="absolute inset-0 opacity-[0.13]"
          style={{
            backgroundImage: "radial-gradient(circle, #94a3b8 1px, transparent 1px)",
            backgroundSize: "28px 28px",
          }}
        />
        <div
          className="absolute -top-28 -right-28 w-80 h-80 rounded-full blur-3xl opacity-[0.08]"
          style={{ background: "#0ea5e9" }}
        />
        <div
          className="absolute -bottom-28 -left-28 w-96 h-96 rounded-full blur-3xl opacity-[0.07]"
          style={{ background: "#f97316" }}
        />
        <div
          className="absolute top-1/2 -left-32 w-64 h-64 rounded-full blur-3xl opacity-[0.05]"
          style={{ background: "#10b981" }}
        />
      </div>

      {/* Header */}
      <div className="relative z-10 bg-white border-b border-neutral-100 px-4 py-3 flex items-center gap-3">
        {/* Hamburger — mobile only */}
        <button
          className="md:hidden flex-shrink-0 p-1.5 rounded-lg text-neutral-500 hover:text-brand-green hover:bg-brand-green-50 transition-all"
          onClick={() => setMobileSidebarOpen(true)}
          aria-label="Open menu"
        >
          <Menu className="w-5 h-5" />
        </button>

        <h1 className="text-base font-bold text-neutral-900 flex-1 truncate">
          {currentChat?.title ?? "New Chat"}
        </h1>

        <button
          onClick={() => setKnowledgeOpen(true)}
          className="flex items-center gap-2 bg-brand-amber text-white px-3 py-2 rounded-lg text-sm font-semibold hover:brightness-105 hover:scale-105 transition-all duration-200 flex-shrink-0"
        >
          <Database className="w-4 h-4" />
          <span className="hidden sm:inline">Knowledge</span>
        </button>
      </div>

      {/* Chat window */}
      <div className="relative z-10 flex-1 flex flex-col min-h-0">
        <ChatWindow
          messages={currentChat?.messages ?? []}
          isLoading={isStreaming}
          streamingMsgId={streamingMsgId}
          userName={user?.name}
        />
      </div>

      {/* Input area */}
      <div className="relative z-10 border-t border-neutral-100 bg-white/80 backdrop-blur-sm">
        <div className="w-full max-w-4xl xl:max-w-5xl mx-auto flex flex-col gap-2 px-3 sm:px-4 pt-3 pb-4">
          <AnimatePresence>
            {isVoiceActive && (
              <VoiceStatusBar voiceState={voiceState} onStop={stop} />
            )}
          </AnimatePresence>
          <ChatInput
            onSend={handleSend}
            onStop={handleStop}
            isLoading={isStreaming}
            voiceState={voiceState}
            onVoiceToggle={handleVoiceToggle}
          />
        </div>
      </div>

      <KnowledgeManagement
        open={knowledgeOpen}
        onClose={() => setKnowledgeOpen(false)}
      />
    </ChatLayout>
  );
}