fusion/packages/dashboard/app/hooks/useChat.ts

import { useState, useEffect, useCallback, useRef } from "react";
import {
  fetchChatSessions,
  createChatSession as apiCreateChatSession,
  fetchChatMessages,
  updateChatSession,
  deleteChatSession,
  attachChatStream,
  streamChatResponse,
  cancelChatResponse,
  fetchAgents,
  type ChatFailureInfo,
  type ChatSessionListResponse,
} from "../api";
import { subscribeSse } from "../sse-bus";
import { getScopedItem, setScopedItem, removeScopedItem } from "../utils/projectStorage";
import type { Agent, ChatInFlightGenerationState, ChatMessage } from "@fusion/core";

const ACTIVE_SESSION_STORAGE_KEY = "kb-chat-active-session";

export interface ChatSessionInfo {
  id: string;
  title?: string | null;
  agentId: string;
  status: string;
  modelProvider?: string | null;
  modelId?: string | null;
  createdAt: string;
  updatedAt: string;
  lastMessagePreview?: string;
  lastMessageAt?: string;
  isGenerating?: boolean;
  inFlightGeneration?: ChatInFlightGenerationState | null;
}

// Re-export shared chat types so existing consumers (`import { ChatMessageInfo } from "../hooks/useChat"`)
// keep working — single source of truth lives in chatTypes.ts.
export type { ChatMessageInfo, FailureInfo, FallbackInfo, ToolCallInfo } from "./chatTypes";
import type { ChatMessageInfo, FailureInfo, FallbackInfo, ToolCallInfo } from "./chatTypes";
import { createChatStreamHandlers } from "./createChatStreamHandlers";
import { isLikelyTabSuspensionError, useTabVisibilitySuspension } from "./visibilitySuspension";

export interface UseChatReturn {
  // Session state
  sessions: ChatSessionInfo[];
  activeSession: ChatSessionInfo | null;
  sessionsLoading: boolean;

  // Message state
  messages: ChatMessageInfo[];
  messagesLoading: boolean;
  isStreaming: boolean;
  streamingText: string;
  streamingThinking: string;
  streamingToolCalls: ToolCallInfo[];
  pendingMessage: string;

  // Session operations
  selectSession: (id: string, sessionOverride?: ChatSessionInfo) => void;
  createSession: (
    input: { agentId: string; title?: string; modelProvider?: string; modelId?: string },
  ) => Promise<ChatSessionInfo>;
  archiveSession: (id: string) => Promise<void>;
  deleteSession: (id: string) => Promise<void>;

  // Message operations
  /** Send a message, optionally with file attachments to upload with the prompt. */
  sendMessage: (content: string, attachments?: File[]) => void;
  stopStreaming: () => void;
  clearPendingMessage: () => void;
  loadMoreMessages: () => Promise<void>;
  hasMoreMessages: boolean;

  // Search/filter
  searchQuery: string;
  setSearchQuery: (query: string) => void;
  filteredSessions: ChatSessionInfo[];

  // Refresh
  refreshSessions: () => Promise<void>;

  // Agent name resolution
  agentsMap: Map<string, Agent>;
}

function parseModelDescriptor(model: string): { modelProvider?: string; modelId?: string } {
  const value = typeof model === "string" ? model.trim() : "";
  const slashIndex = value.indexOf("/");
  if (!value || slashIndex <= 0 || slashIndex >= value.length - 1) {
    return {};
  }
  return {
    modelProvider: value.slice(0, slashIndex),
    modelId: value.slice(slashIndex + 1),
  };
}

function extractCompletedToolCalls(metadata: Record<string, unknown> | null | undefined): ToolCallInfo[] | undefined {
  const rawToolCalls = metadata?.toolCalls;
  if (!Array.isArray(rawToolCalls)) {
    return undefined;
  }

  const parsed = rawToolCalls
    .map((toolCall): ToolCallInfo | null => {
      if (!toolCall || typeof toolCall !== "object") {
        return null;
      }

      const record = toolCall as Record<string, unknown>;
      const toolName = typeof record.toolName === "string" ? record.toolName : "";
      if (!toolName) {
        return null;
      }

      const args = record.args;

      return {
        toolName,
        ...(args && typeof args === "object" ? { args: args as Record<string, unknown> } : {}),
        isError: Boolean(record.isError),
        result: record.result,
        status: "completed" as const,
      };
    })
    .filter((toolCall): toolCall is ToolCallInfo => toolCall !== null);

  return parsed.length > 0 ? parsed : undefined;
}

function extractFallbackInfo(metadata: Record<string, unknown> | null | undefined): FallbackInfo | undefined {
  const rawFallback = metadata?.fallback;
  if (!rawFallback || typeof rawFallback !== "object") {
    return undefined;
  }

  const record = rawFallback as Record<string, unknown>;
  const primaryModel = typeof record.primaryModel === "string" ? record.primaryModel : "";
  const fallbackModel = typeof record.fallbackModel === "string" ? record.fallbackModel : "";
  const triggerPoint = record.triggerPoint;
  if (!primaryModel || !fallbackModel || (triggerPoint !== "session-creation" && triggerPoint !== "prompt-time")) {
    return undefined;
  }

  return {
    primaryModel,
    fallbackModel,
    triggerPoint,
  };
}

function extractFailureInfo(metadata: Record<string, unknown> | null | undefined): FailureInfo | undefined {
  const rawFailure = metadata?.failureInfo;
  if (!rawFailure || typeof rawFailure !== "object") {
    return undefined;
  }

  const record = rawFailure as Record<string, unknown>;
  const summary = typeof record.summary === "string" ? record.summary.trim() : "";
  if (!summary) {
    return undefined;
  }

  const reference = (() => {
    const rawReference = record.reference;
    if (!rawReference || typeof rawReference !== "object") {
      return undefined;
    }
    const referenceRecord = rawReference as Record<string, unknown>;
    const kind = typeof referenceRecord.kind === "string" ? referenceRecord.kind.trim() : "";
    const id = typeof referenceRecord.id === "string" ? referenceRecord.id.trim() : "";
    if (!kind || !id) {
      return undefined;
    }
    return {
      kind,
      id,
      ...(typeof referenceRecord.label === "string" && referenceRecord.label.trim()
        ? { label: referenceRecord.label.trim() }
        : {}),
    };
  })();

  return {
    summary,
    ...(typeof record.errorClass === "string" && record.errorClass.trim()
      ? { errorClass: record.errorClass.trim() }
      : {}),
    ...(typeof record.code === "string" && record.code.trim()
      ? { code: record.code.trim() }
      : {}),
    ...(typeof record.detail === "string" && record.detail.trim()
      ? { detail: record.detail.trim() }
      : {}),
    ...(reference ? { reference } : {}),
  };
}

function normalizeFailureInfo(data: string | ChatFailureInfo): FailureInfo {
  if (typeof data === "string") {
    const summary = data.trim() || "Failed to get response";
    return { summary };
  }

  const summary = typeof data.summary === "string" && data.summary.trim()
    ? data.summary.trim()
    : "Failed to get response";

  return {
    summary,
    ...(typeof data.errorClass === "string" && data.errorClass.trim()
      ? { errorClass: data.errorClass.trim() }
      : {}),
    ...(typeof data.code === "string" && data.code.trim()
      ? { code: data.code.trim() }
      : {}),
    ...(typeof data.detail === "string" && data.detail.trim()
      ? { detail: data.detail.trim() }
      : {}),
    ...(data.reference ? { reference: data.reference } : {}),
  };
}

function mapChatMessageToInfo(message: ChatMessage): ChatMessageInfo {
  return {
    id: message.id,
    sessionId: message.sessionId,
    role: message.role,
    content: message.content,
    thinkingOutput: message.thinkingOutput,
    toolCalls: extractCompletedToolCalls(message.metadata),
    fallbackInfo: extractFallbackInfo(message.metadata),
    failureInfo: extractFailureInfo(message.metadata),
    attachments: message.attachments,
    createdAt: message.createdAt,
  };
}

export function useChat(
  projectId?: string,
  addToast?: (msg: string, type?: "success" | "error" | "warning") => void,
): UseChatReturn {
  // Session state
  const [sessions, setSessions] = useState<ChatSessionInfo[]>([]);
  const [activeSession, setActiveSession] = useState<ChatSessionInfo | null>(null);
  const [sessionsLoading, setSessionsLoading] = useState(true);

  // Message state
  const [messages, setMessages] = useState<ChatMessageInfo[]>([]);
  const [messagesLoading, setMessagesLoading] = useState(false);
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingText, setStreamingText] = useState("");
  const [streamingThinking, setStreamingThinking] = useState("");
  const [streamingToolCalls, setStreamingToolCalls] = useState<ToolCallInfo[]>([]);
  const [pendingMessage, setPendingMessage] = useState("");

  // Search/filter
  const [searchQuery, setSearchQuery] = useState("");

  // Pagination
  const [hasMoreMessages, setHasMoreMessages] = useState(true);

  // Agent name resolution map
  const [agentsMap, setAgentsMap] = useState<Map<string, Agent>>(new Map());

  // Stream connection ref for cleanup
  const streamRef = useRef<{ close: () => void } | null>(null);
  const cancelledByUserRef = useRef(false);
  const pendingMessageRef = useRef("");
  // Cancel any pending requestAnimationFrame flushes from the active stream.
  // Set when sendMessage starts, cleared on done/error. Called from stopStreaming
  // so a clear-then-rAF-fires sequence doesn't flash stale text back in.
  const cancelStreamingFlushesRef = useRef<(() => void) | null>(null);

  // Refs for SSE event handlers to access current state
  const sessionsRef = useRef(sessions);
  const activeSessionRef = useRef(activeSession);
  const isStreamingRef = useRef(isStreaming);
  sessionsRef.current = sessions;
  activeSessionRef.current = activeSession;
  isStreamingRef.current = isStreaming;

  useEffect(() => {
    pendingMessageRef.current = pendingMessage;
  }, [pendingMessage]);

  // Tracks message IDs that were added via streaming completion.
  // Used to prevent duplicate messages when SSE event arrives before streaming state clears.
  const streamingMessageIdsRef = useRef<Set<string>>(new Set());

  // Tracks the project context version to detect stale SSE events after project switches.
  // Incremented whenever projectId changes, invalidating any in-flight SSE handlers.
  const projectContextVersionRef = useRef(0);
  // Track previous projectId to detect changes
  const previousProjectIdRef = useRef<string | undefined>(projectId);

  // Detect project changes and invalidate SSE context
  if (previousProjectIdRef.current !== projectId) {
    previousProjectIdRef.current = projectId;
    projectContextVersionRef.current++;
  }

  // Fetch agents on mount for name resolution (project-scoped with stale-request protection)
  useEffect(() => {
    const contextVersionAtStart = projectContextVersionRef.current;
    fetchAgents(undefined, projectId)
      .then((agents) => {
        // Ignore response if project changed during fetch
        if (projectContextVersionRef.current !== contextVersionAtStart) return;
        const map = new Map<string, Agent>();
        for (const agent of agents) {
          map.set(agent.id, agent);
        }
        setAgentsMap(map);
      })
      .catch(() => {
        // Silently fail - keep empty map
      });
  }, [projectId]);

  // Fetch sessions
  const refreshSessions = useCallback(async () => {
    setSessionsLoading(true);
    try {
      const data: ChatSessionListResponse = await fetchChatSessions(projectId);
      // Sort by updatedAt descending
      const sorted = [...data.sessions].sort(
        (a, b) => new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime(),
      );
      setSessions(sorted);
    } catch {
      // Silently fail on refresh
    } finally {
      setSessionsLoading(false);
    }
  }, [projectId]);

  // Initial load
  useEffect(() => {
    refreshSessions();
  }, [refreshSessions]);

  // Restore active session from localStorage after initial load.
  // Uses refs to avoid circular dependency with selectSession and to avoid
  // re-selecting/resetting the thread on every sessions refresh.
  const selectSessionRef = useRef<(id: string, sessionOverride?: ChatSessionInfo) => void>(() => {
    /* noop - will be replaced after selectSession is defined */
  });
  const hasRestoredActiveSessionRef = useRef(false);

  useEffect(() => {
    hasRestoredActiveSessionRef.current = false;
  }, [projectId]);

  useEffect(() => {
    if (sessionsLoading || hasRestoredActiveSessionRef.current || activeSessionRef.current) return;

    const savedSessionId = getScopedItem(ACTIVE_SESSION_STORAGE_KEY, projectId);
    if (!savedSessionId) {
      hasRestoredActiveSessionRef.current = true;
      return;
    }

    const session = sessions.find((s) => s.id === savedSessionId);
    if (session) {
      hasRestoredActiveSessionRef.current = true;
      selectSessionRef.current(savedSessionId, session);
      return;
    }

    hasRestoredActiveSessionRef.current = true;
  }, [sessionsLoading, sessions, projectId]);

  // Load messages when active session changes
  const loadMessages = useCallback(
    async (sessionId: string, opts?: { offset?: number }) => {
      setMessagesLoading(true);
      try {
        const data = await fetchChatMessages(sessionId, { limit: 50, ...opts }, projectId);
        const mappedMessages = data.messages.map(mapChatMessageToInfo);
        if (opts?.offset && opts.offset > 0) {
          // Prepend older messages
          setMessages((prev) => [...mappedMessages, ...prev]);
        } else {
          setMessages(mappedMessages);
        }
        setHasMoreMessages(data.messages.length >= 50);
      } catch {
        // Silently fail
      } finally {
        setMessagesLoading(false);
      }
    },
    [projectId],
  );

  const resetTransientComposerState = useCallback(() => {
    cancelStreamingFlushesRef.current?.();
    cancelStreamingFlushesRef.current = null;
    pendingMessageRef.current = "";
    setPendingMessage("");
    setStreamingText("");
    setStreamingThinking("");
    setStreamingToolCalls([]);
    setIsStreaming(false);
  }, []);

  const attachIfGenerating = useCallback((sessionId: string, inFlightGeneration?: ChatInFlightGenerationState | null) => {
    if (streamRef.current || !sessionId) {
      return true;
    }

    cancelledByUserRef.current = false;
    if (inFlightGeneration) {
      setStreamingText(inFlightGeneration.streamingText);
      setStreamingThinking(inFlightGeneration.streamingThinking);
      setStreamingToolCalls(inFlightGeneration.toolCalls);
    }
    setIsStreaming(true);

    const { handlers } = createChatStreamHandlers({
      sessionId,
      tempUserMessageId: "",
      setStreamingText,
      setStreamingThinking,
      setStreamingToolCalls,
      cancelStreamingFlushesRef,
      addToast,
      onFallbackSession: (data, fallbackSessionId) => {
        const nextModel = parseModelDescriptor(data.fallbackModel);
        setSessions((prev) => prev.map((session) =>
          session.id === fallbackSessionId ? { ...session, ...nextModel } : session,
        ));
        setActiveSession((prev) => prev && prev.id === fallbackSessionId ? { ...prev, ...nextModel } : prev);
      },
      onDone: () => {
        setStreamingText("");
        setStreamingThinking("");
        setStreamingToolCalls([]);
        setIsStreaming(false);
        isStreamingRef.current = false;
        streamRef.current = null;
        void loadMessages(sessionId);
      },
      onError: (data) => {
        setStreamingText("");
        setStreamingThinking("");
        setStreamingToolCalls([]);
        setIsStreaming(false);
        isStreamingRef.current = false;
        streamRef.current = null;
        const failureInfo = normalizeFailureInfo(data);
        addToast?.(failureInfo.summary, "error");
        void loadMessages(sessionId);
      },
    });

    const stream = attachChatStream(sessionId, handlers, projectId, {
      ...(typeof inFlightGeneration?.replayFromEventId === "number"
        ? { lastEventId: inFlightGeneration.replayFromEventId }
        : {}),
    });
    streamRef.current = stream;
    return true;
  }, [addToast, loadMessages, projectId]);

  // Select a session
  const selectSession = useCallback(
    (id: string, sessionOverride?: ChatSessionInfo) => {
      const currentActiveSessionId = activeSessionRef.current?.id ?? null;
      if (id && currentActiveSessionId === id && !sessionOverride) {
        return;
      }

      // Close any existing stream
      if (streamRef.current) {
        streamRef.current.close();
        streamRef.current = null;
      }

      // Find and set active session
      const session = sessionOverride ?? sessions.find((s) => s.id === id);
      setActiveSession(session || null);

      // Reset transient state
      resetTransientComposerState();
      setHasMoreMessages(true);

      // Load messages for this session
      if (id) {
        loadMessages(id);
      } else {
        setMessages([]);
      }

      // Recover streaming state if the server reports an active generation.
      // After a reload/HMR, the server keeps generating but the UI loses
      // all streaming state. Showing "Connecting…" immediately tells the
      // user the AI is still working.
      if (session?.isGenerating) {
        attachIfGenerating(session.id, session.inFlightGeneration);
      }

      // Persist active session to localStorage
      if (id) {
        setScopedItem(ACTIVE_SESSION_STORAGE_KEY, id, projectId);
      } else {
        removeScopedItem(ACTIVE_SESSION_STORAGE_KEY, projectId);
      }
    },
    [attachIfGenerating, sessions, loadMessages, projectId, resetTransientComposerState],
  );

  // Update the ref to point to the actual selectSession function
  // This is needed to avoid circular dependencies in useEffect
  selectSessionRef.current = selectSession;

  // Create a new session
  const createSession = useCallback(
    async (input: { agentId: string; title?: string; modelProvider?: string; modelId?: string }) => {
      const data = await apiCreateChatSession(input, projectId);

      if (streamRef.current) {
        streamRef.current.close();
        streamRef.current = null;
      }
      const newSession: ChatSessionInfo = {
        id: data.session.id,
        title: data.session.title,
        agentId: data.session.agentId,
        status: data.session.status,
        modelProvider: data.session.modelProvider,
        modelId: data.session.modelId,
        createdAt: data.session.createdAt,
        updatedAt: data.session.updatedAt,
      };

      setSessions((prev) => {
        if (prev.some((s) => s.id === newSession.id)) return prev;
        return [newSession, ...prev];
      });

      resetTransientComposerState();
      selectSession(newSession.id, newSession);
      setMessages([]);

      return newSession;
    },
    [projectId, resetTransientComposerState, selectSession],
  );

  // Archive a session
  const archiveSession = useCallback(
    async (id: string) => {
      await updateChatSession(id, { status: "archived" }, projectId);
      // Remove from sessions list
      setSessions((prev) => prev.filter((s) => s.id !== id));
      // If it was the active session, clear it
      if (activeSession?.id === id) {
        setActiveSession(null);
        setMessages([]);
      }
    },
    [activeSession, projectId],
  );

  // Delete a session
  const deleteSession = useCallback(
    async (id: string) => {
      // Close stream if active
      if (activeSession?.id === id && streamRef.current) {
        streamRef.current.close();
        streamRef.current = null;
      }

      await deleteChatSession(id, projectId);
      // Remove from sessions list
      setSessions((prev) => prev.filter((s) => s.id !== id));
      // If it was the active session, clear it
      if (activeSession?.id === id) {
        setActiveSession(null);
        setMessages([]);
      }
    },
    [activeSession, projectId],
  );

  // Load more messages (pagination)
  const loadMoreMessages = useCallback(async () => {
    if (!activeSession || !hasMoreMessages) return;
    await loadMessages(activeSession.id, { offset: messages.length });
  }, [activeSession, hasMoreMessages, loadMessages, messages.length]);

  const stopStreaming = useCallback(() => {
    if (!activeSession) return;

    cancelledByUserRef.current = true;
    cancelStreamingFlushesRef.current?.();
    cancelStreamingFlushesRef.current = null;
    streamRef.current?.close();
    streamRef.current = null;

    void cancelChatResponse(activeSession.id, projectId).catch(() => {
      // Best-effort cancellation; ignore backend errors.
    });

    setIsStreaming(false);
    setStreamingText("");
    setStreamingThinking("");
    setStreamingToolCalls([]);
  }, [activeSession, projectId]);

  const clearPendingMessage = useCallback(() => {
    pendingMessageRef.current = "";
    setPendingMessage("");
  }, []);

  /**
   * Send a user message to the active chat session.
   * @param content Message text content to send.
   * @param attachments Optional files to upload with the message in the same request.
   */
  const sendMessageRef = useRef<(content: string, attachments?: File[]) => void>(() => {
    // no-op until sendMessage is defined
  });
  const visibilitySuspension = useTabVisibilitySuspension();

  const sendMessage = useCallback(
    (content: string, attachments?: File[]) => {
      if (!activeSession) return;

      if (isStreamingRef.current) {
        pendingMessageRef.current = content;
        setPendingMessage(content);
        return;
      }

      cancelledByUserRef.current = false;

      // Close any existing stream
      if (streamRef.current) {
        streamRef.current.close();
        streamRef.current = null;
      }

      // Optimistically add user message
      const tempId = `temp-${Date.now()}`;
      const userMessage: ChatMessageInfo = {
        id: tempId,
        sessionId: activeSession.id,
        role: "user",
        content,
        createdAt: new Date().toISOString(),
      };
      setMessages((prev) => [...prev, userMessage]);

      // Clear streaming state
      setStreamingText("");
      setStreamingThinking("");
      setStreamingToolCalls([]);
      setIsStreaming(true);

      const { handlers } = createChatStreamHandlers({
        sessionId: activeSession.id,
        tempUserMessageId: tempId,
        setStreamingText,
        setStreamingThinking,
        setStreamingToolCalls,
        cancelStreamingFlushesRef,
        addToast,
        onFallbackSession: (data, sessionId) => {
          const nextModel = parseModelDescriptor(data.fallbackModel);
          setSessions((prev) => prev.map((session) =>
            session.id === sessionId ? { ...session, ...nextModel } : session,
          ));
          setActiveSession((prev) => prev && prev.id === sessionId ? { ...prev, ...nextModel } : prev);
        },
        onDone: ({ messageId, message: finalMessage, accumulated }) => {
          const assistantMessage: ChatMessageInfo = finalMessage
            ? mapChatMessageToInfo(finalMessage)
            : {
                id: messageId || `msg-${Date.now()}`,
                sessionId: activeSession.id,
                role: "assistant",
                content: accumulated.text,
                thinkingOutput: accumulated.thinking,
                toolCalls: accumulated.toolCalls.length > 0 ? accumulated.toolCalls : undefined,
                fallbackInfo: accumulated.fallbackInfo,
                createdAt: new Date().toISOString(),
              };

          // Track this message ID so the SSE chatMessageAdded handler skips it
          // if the broadcast event arrives before our optimistic add settles.
          streamingMessageIdsRef.current.add(assistantMessage.id);

          // Preserve user message and add assistant message
          setMessages((prev) => [...prev, assistantMessage]);

          setStreamingText("");
          setStreamingThinking("");
          setStreamingToolCalls([]);
          setIsStreaming(false);
          isStreamingRef.current = false;
          streamRef.current = null;

          // Clean up tracked ID after a short delay (SSE event should arrive quickly)
          setTimeout(() => {
            streamingMessageIdsRef.current.delete(assistantMessage.id);
          }, 1000);

          refreshSessions();

          const queuedMessage = pendingMessageRef.current.trim();
          if (queuedMessage) {
            pendingMessageRef.current = "";
            setPendingMessage("");
            sendMessageRef.current(queuedMessage);
          }
        },
        onError: (data, tempUserMessageId) => {
          const failureInfo = normalizeFailureInfo(data);
          const shouldSuppressSuspensionError = typeof data === "string"
            && isLikelyTabSuspensionError(data)
            && (visibilitySuspension.isHiddenNow() || visibilitySuspension.wasRecentlyHidden(5000));

          setMessages((prev) => {
            const nextMessages = prev.filter((message) => message.id !== tempUserMessageId);
            if (shouldSuppressSuspensionError) {
              return nextMessages;
            }
            return [
              ...nextMessages,
              {
                id: `error-${Date.now()}`,
                sessionId: activeSession.id,
                role: "assistant",
                content: failureInfo.summary,
                failureInfo,
                createdAt: new Date().toISOString(),
              },
            ];
          });
          setStreamingText("");
          setStreamingThinking("");
          setStreamingToolCalls([]);
          setIsStreaming(false);
          isStreamingRef.current = false;
          streamRef.current = null;
          console.error("[useChat] Stream error:", data);

          if (shouldSuppressSuspensionError) {
            console.info("[useChat] Suppressed tab-suspension stream error:", data);
            if (activeSession?.id) {
              if (activeSession.isGenerating) {
                attachIfGenerating(activeSession.id, activeSession.inFlightGeneration);
              } else {
                void loadMessages(activeSession.id);
              }
            }
          } else {
            addToast?.(failureInfo.summary, "error");
            void refreshSessions();
          }

          if (!cancelledByUserRef.current) {
            const queuedMessage = pendingMessageRef.current.trim();
            if (queuedMessage) {
              pendingMessageRef.current = "";
              setPendingMessage("");
              sendMessageRef.current(queuedMessage);
            }
          }
        },
      });

      streamRef.current = streamChatResponse(activeSession.id, content, handlers, attachments, projectId);
    },
    [activeSession, projectId, refreshSessions, addToast, loadMessages, attachIfGenerating, visibilitySuspension],
  );

  sendMessageRef.current = sendMessage;

  // Filter sessions based on search query
  const filteredSessions = searchQuery
    ? sessions.filter(
        (s) =>
          s.title?.toLowerCase().includes(searchQuery.toLowerCase()) ||
          s.agentId.toLowerCase().includes(searchQuery.toLowerCase()),
      )
    : sessions;

  // Recovery mode polling: if reloaded mid-generation, keep waiting state alive
  // until generation finishes and messages can be reloaded.
  useEffect(() => {
    if (!activeSessionRef.current?.isGenerating) return;

    if (!streamRef.current) {
      attachIfGenerating(activeSessionRef.current.id, activeSessionRef.current.inFlightGeneration);
    }

    if (!isStreamingRef.current || streamRef.current || !activeSessionRef.current) return;

    const interval = setInterval(async () => {
      if (!isStreamingRef.current || streamRef.current || !activeSessionRef.current) {
        clearInterval(interval);
        return;
      }

      try {
        const data: ChatSessionListResponse = await fetchChatSessions(projectId);
        const session = data.sessions.find((candidate) => candidate.id === activeSessionRef.current?.id);
        if (!session?.isGenerating) {
          clearInterval(interval);
          await loadMessages(activeSessionRef.current.id);
          setStreamingText("");
          setStreamingThinking("");
          setStreamingToolCalls([]);
          setIsStreaming(false);
        }
      } catch {
        // Silently fail - will retry next interval
      }
    }, 3000);

    return () => clearInterval(interval);
  }, [attachIfGenerating, loadMessages, projectId, activeSession]);

  // SSE real-time updates
  useEffect(() => {
    const contextVersionAtStart = projectContextVersionRef.current;
    const query = projectId ? `?projectId=${encodeURIComponent(projectId)}` : "";

    const isStale = () => projectContextVersionRef.current !== contextVersionAtStart;

    const handleChatSessionCreated = (e: MessageEvent) => {
      if (isStale()) return;
      const session: ChatSessionInfo = JSON.parse(e.data);
      // Avoid duplicates
      setSessions((prev) => {
        if (prev.some((s) => s.id === session.id)) return prev;
        // Add at the top (sessions are sorted by updatedAt desc)
        return [session, ...prev];
      });
    };

    const handleChatSessionUpdated = (e: MessageEvent) => {
      if (isStale()) return;
      const updatedSession: ChatSessionInfo = JSON.parse(e.data);
      setSessions((prev) => {
        const updated = prev.map((s) => (s.id === updatedSession.id ? updatedSession : s));
        return [...updated];
      });
      // If this is the active session, update it too
      if (activeSessionRef.current?.id === updatedSession.id) {
        setActiveSession(updatedSession);
        if (updatedSession.isGenerating && !streamRef.current) {
          attachIfGenerating(updatedSession.id, updatedSession.inFlightGeneration);
        }
      }
    };

    const handleChatSessionDeleted = (e: MessageEvent) => {
      if (isStale()) return;
      const { id: sessionId }: { id: string } = JSON.parse(e.data);
      setSessions((prev) => prev.filter((s) => s.id !== sessionId));
      // If this was the active session, clear it
      if (activeSessionRef.current?.id === sessionId) {
        setActiveSession(null);
        setMessages([]);
      }
    };

    const handleChatMessageAdded = (e: MessageEvent) => {
      if (isStale()) return;
      const rawMessage = JSON.parse(e.data) as ChatMessage;
      const message = mapChatMessageToInfo(rawMessage);

      // Skip if this message was already added via streaming completion
      // (SSE event may arrive before streaming state clears)
      if (streamingMessageIdsRef.current.has(message.id)) {
        return;
      }

      // Recovery mode: isStreaming is true but there's no active stream (streamRef is null).
      // This happens after a page reload/HMR when the server is still generating.
      // When the assistant message arrives via SSE, add it and clear the recovery state.
      if (
        activeSessionRef.current?.id === message.sessionId &&
        isStreamingRef.current &&
        !streamRef.current &&
        message.role === "assistant"
      ) {
        setMessages((prev) => {
          if (prev.some((m) => m.id === message.id)) return prev;
          return [...prev, message];
        });
        setStreamingText("");
        setStreamingThinking("");
        setStreamingToolCalls([]);
        setIsStreaming(false);
        return;
      }

      // Only add if this is the active session AND we're not streaming
      // (during streaming, messages are managed locally to avoid duplicates)
      // Use ref to get the current value (state may not be updated yet when handler runs)
      if (activeSessionRef.current?.id === message.sessionId && !isStreamingRef.current) {
        setMessages((prev) => {
          // Avoid duplicates by persisted id first.
          if (prev.some((m) => m.id === message.id)) return prev;

          // Reconcile optimistic local user messages against persisted SSE echoes.
          // The optimistic message uses a temp id and should be replaced instead of appended.
          if (message.role === "user") {
            const optimisticIndex = prev.findIndex((candidate) =>
              candidate.role === "user"
              && candidate.id.startsWith("temp-")
              && candidate.content.trim() === message.content.trim(),
            );
            if (optimisticIndex >= 0) {
              const next = [...prev];
              next[optimisticIndex] = message;
              return next;
            }
          }

          return [...prev, message];
        });
      }
    };

    const handleChatMessageDeleted = (e: MessageEvent) => {
      if (isStale()) return;
      const { id: messageId }: { id: string } = JSON.parse(e.data);
      setMessages((prev) => prev.filter((m) => m.id !== messageId));
    };

    const unsubscribe = subscribeSse(`/api/events${query}`, {
      events: {
        "chat:session:created": handleChatSessionCreated,
        "chat:session:updated": handleChatSessionUpdated,
        "chat:session:deleted": handleChatSessionDeleted,
        "chat:message:added": handleChatMessageAdded,
        "chat:message:deleted": handleChatMessageDeleted,
      },
    });

    return unsubscribe;
  }, [attachIfGenerating, projectId]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (streamRef.current) {
        streamRef.current.close();
        streamRef.current = null;
      }
    };
  }, []);

  return {
    sessions,
    activeSession,
    sessionsLoading,
    messages,
    messagesLoading,
    isStreaming,
    streamingText,
    streamingThinking,
    streamingToolCalls,
    pendingMessage,
    selectSession,
    createSession,
    archiveSession,
    deleteSession,
    sendMessage,
    stopStreaming,
    clearPendingMessage,
    loadMoreMessages,
    hasMoreMessages,
    searchQuery,
    setSearchQuery,
    filteredSessions,
    refreshSessions,
    agentsMap,
  };
}