gitdataai/src/contexts/room-context.tsx

import {
  createContext,
  type ReactNode,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import { useNavigate } from 'react-router-dom';
import { toast } from 'sonner';
import {
  type ProjectRepositoryItem,
  type RoomCategoryResponse,
  type RoomMemberResponse,
  type RoomMessageResponse,
  type RoomPinResponse,
  type RoomResponse,
  type RoomThreadResponse,
  categoryList as restCategoryList,
  roomList as restRoomList,
} from '@/client';
import {
  createRoomWsClient,
  type RoomWsClient,
  type RoomWsStatus,
  type RoomMessagePayload,
  type RoomCategoryResponse as WsRoomCategoryResponse,
  type RoomReactionUpdatedPayload,
  type ReactionListData,
} from '@/lib/room-ws-client';
import { requestWsToken } from '@/lib/ws-token';
import { useUser } from '@/contexts';

export type { RoomWsStatus, RoomWsClient } from '@/lib/room-ws-client';

export type PresenceStatus = 'online' | 'away' | 'dnd' | 'offline';
export type PresenceMap = Record<string, PresenceStatus>; // keyed by user_id

export interface RoomAiConfig {
  model: string;
  modelName?: string;
}

export interface ReactionGroup {
  emoji: string;
  count: number;
  reacted_by_me: boolean;
  users: string[];
}

const MAX_MESSAGES_IN_MEMORY = 1000;

function handleRoomError(context: string, error: unknown): void {
  const message = error instanceof Error ? error.message : 'Operation failed';
  console.error(`[RoomContext] ${context}:`, error);
  toast.error(`${context} failed`, {
    description: message,
  });
}

export type MessageWithMeta = RoomMessageResponse & {
  thread_id?: string;
  display_content?: string;
  is_streaming?: boolean;
  isOptimisticError?: boolean;
  /** True for messages sent by the current user that haven't been confirmed by the server */
  isOptimistic?: boolean;
  reactions?: ReactionGroup[];
  /** Attachment IDs for files uploaded with this message */
  attachment_ids?: string[];
  /** AI stream chunk type: "thinking", "tool_call", "tool_result", or undefined for normal text */
  chunk_type?: string;
  /** Accumulated thinking/reasoning content from AI stream (collapsible) */
  thinking_content?: string;
};

export type RoomWithCategory = RoomResponse & {
  category_info?: RoomCategoryResponse | null;
};

export type UiMessage = MessageWithMeta;

function wsMessageToUiMessage(wsMsg: RoomMessagePayload): MessageWithMeta {
  return {
    id: wsMsg.id,
    seq: wsMsg.seq,
    room: wsMsg.room_id,
    sender_type: wsMsg.sender_type,
    sender_id: wsMsg.sender_id,
    display_name: wsMsg.display_name,
    thread: wsMsg.thread_id,
    thread_id: wsMsg.thread_id,
    content: wsMsg.content,
    content_type: wsMsg.content_type,
    send_at: wsMsg.send_at,
    display_content: wsMsg.content,
    is_streaming: false,
    reactions: wsMsg.reactions,
  };
}

interface RoomContextValue {
  wsStatus: RoomWsStatus;
  wsError: string | null;
  wsClient: RoomWsClient | null;
  connectWs: () => Promise<void>;
  disconnectWs: () => void;
  presence: PresenceMap;

  rooms: RoomWithCategory[];
  roomsLoading: boolean;
  roomsError: Error | null;
  refreshRooms: () => Promise<void>;

  categories: RoomCategoryResponse[];
  categoriesLoading: boolean;
  refreshCategories: () => Promise<void>;
  createCategory: (name: string, position?: number) => Promise<WsRoomCategoryResponse>;
  updateCategory: (id: string, name?: string, position?: number) => Promise<void>;
  deleteCategory: (id: string) => Promise<void>;

  activeRoom: RoomResponse | null;
  activeRoomId: string | null;
  setActiveRoom: (roomId: string | null) => void;

  messages: MessageWithMeta[];
  messagesLoading: boolean;
  isHistoryLoaded: boolean;
  isLoadingMore: boolean;
  isTransitioningRoom: boolean;
  nextCursor: number | null;
  loadMore: (cursor?: number | null) => void;
  sendMessage: (content: string, contentType?: string, inReplyTo?: string, attachmentIds?: string[]) => Promise<void>;
  editMessage: (messageId: string, content: string) => Promise<void>;
  revokeMessage: (messageId: string) => Promise<void>;
  updateReadSeq: (seq: number) => Promise<void>;

  members: RoomMemberResponse[];
  membersLoading: boolean;
  refreshMembers: () => Promise<void>;
  addMember: (userId: string, role?: string) => Promise<void>;
  removeMember: (userId: string) => Promise<void>;
  updateMemberRole: (userId: string, role: string) => Promise<void>;

  pins: RoomPinResponse[];
  pinsLoading: boolean;
  refreshPins: () => Promise<void>;
  pinMessage: (messageId: string) => Promise<void>;
  unpinMessage: (messageId: string) => Promise<void>;

  threads: RoomThreadResponse[];
  refreshThreads: () => Promise<void>;
  createThread: (parentSeq: number) => Promise<RoomThreadResponse>;

  createRoom: (name: string, isPublic: boolean, categoryId?: string) => Promise<RoomResponse>;
  updateRoom: (roomId: string, name?: string, isPublic?: boolean, categoryId?: string) => Promise<void>;
  deleteRoom: (roomId: string) => Promise<void>;
  streamingMessages: Map<string, string>;
  /** Streaming thinking/reasoning content keyed by message_id */
  streamingThinkingContent: Map<string, string>;
  /** Active AI stream info for typing indicator */
  activeAiStream: { message_id: string; display_name: string } | null;

  /** Project repositories for @repository: mention suggestions */
  projectRepos: ProjectRepositoryItem[];
  reposLoading?: boolean;
  /** Room AI configs for @ai: mention suggestions */
  roomAiConfigs: RoomAiConfig[];
  aiConfigsLoading?: boolean;

  /** Typing users in the active room: roomId -> userId -> { username, avatar_url } */
  typingUsers: Record<string, Record<string, { username: string; avatar_url?: string; timeoutId?: ReturnType<typeof setTimeout> }>>;
}

const RoomContext = createContext<RoomContextValue | null>(null);

interface RoomProviderProps {
  projectName: string | null;
  /** Room to open immediately on mount */
  initialRoomId?: string | null;
  children: ReactNode;
}

export function RoomProvider({
  projectName,
  initialRoomId = null,
  children,
}: RoomProviderProps) {
  const { user } = useUser();
  const navigate = useNavigate();

  const [activeRoomId, setActiveRoomId] = useState<string | null>(initialRoomId);
  const [wsClient, setWsClient] = useState<RoomWsClient | null>(null);
  const wsClientRef = useRef<RoomWsClient | null>(null);
  const activeRoomIdRef = useRef<string | null>(activeRoomId);
  const [wsStatus, setWsStatus] = useState<RoomWsStatus>('idle');
  const [wsError, setWsError] = useState<string | null>(null);
  const [wsToken, setWsToken] = useState<string | null>(null);
  // Buffer for messages received while user is in a different room (Bug 3 fix).
  // Merged into state when the user switches to that room.
  const pendingRoomMessagesRef = useRef<Map<string, RoomMessagePayload[]>>(new Map());

  // Keep ref updated with latest activeRoomId
  activeRoomIdRef.current = activeRoomId;

  useEffect(() => {
    let cancelled = false;

    const fetchToken = async () => {
      try {
        const token = await requestWsToken();
        if (!cancelled) {
          setWsToken(token);
        }
      } catch (error) {
        console.error('[RoomContext] Failed to fetch WS token:', error);
        // Token fetch is not critical - client can fall back to cookie auth
        if (!cancelled) {
          setWsToken(null);
        }
      }
    };

    fetchToken();

    return () => {
      cancelled = true;
    };
  }, []);

  const [rooms, setRooms] = useState<RoomWithCategory[]>([]);
  const [roomsLoading, setRoomsLoading] = useState(false);
  const [roomsError, setRoomsError] = useState<Error | null>(null);

  const [categories, setCategories] = useState<RoomCategoryResponse[]>([]);
  const [categoriesLoading, setCategoriesLoading] = useState(false);

  // Merge category_info into rooms whenever either changes
  const roomsWithCategory = useMemo<RoomWithCategory[]>(() => {
    const catMap = new Map(categories.map((c) => [c.id, c]));
    return rooms.map((r) => ({
      ...r,
      category_info: r.category ? (catMap.get(r.category) ?? null) : null,
    }));
  }, [rooms, categories]);

  const [activeRoom, setActiveRoomState] = useState<RoomResponse | null>(null);

  const [messages, setMessages] = useState<MessageWithMeta[]>([]);
  const [isHistoryLoaded, setIsHistoryLoaded] = useState(false);
  const [isLoadingMore, setIsLoadingMore] = useState(false);
  const [nextCursor, setNextCursor] = useState<number | null>(null);

  const loadMessagesAbortRef = useRef<AbortController | null>(null);

  const prevRoomIdRef = useRef<string | null>(null);
  const [isTransitioningRoom, setIsTransitioningRoom] = useState(false);

  useEffect(() => {
    if (prevRoomIdRef.current !== activeRoomId) {
      prevRoomIdRef.current = activeRoomId;
      loadMessagesAbortRef.current?.abort();
      loadMessagesAbortRef.current = null;
      // Mark as transitioning room, show loading state instead of clearing immediately
      setIsTransitioningRoom(true);
      // Immediately clear old room's messages from state to prevent flicker
      setMessages([]);
      setIsHistoryLoaded(false);
      setNextCursor(null);

      // Merge any buffered messages for the new room (Bug 3 fix)
      if (activeRoomId) {
        const pending = pendingRoomMessagesRef.current.get(activeRoomId);
        if (pending && pending.length > 0) {
          pendingRoomMessagesRef.current.delete(activeRoomId);
          setMessages((prev) => {
            const merged = [...prev, ...pending.map(wsMessageToUiMessage)];
            merged.sort((a, b) => a.seq - b.seq);
            return merged;
          });
        }
      }
      // NOTE: intentionally NOT clearing IndexedDB — keeping it enables instant
      // load when the user returns to this room without waiting for API.
    }
  }, [activeRoomId]);

  // ── Subscribe to room (WS must already be connected) ───────────────────────
  useEffect(() => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) return;

    // Load messages via WS (with HTTP fallback)
    loadMore(null);

    // Subscribe to room events. connect() is already called at the provider
    // level — subscribe/unsubscribe only manage per-room event routing.
    client.subscribeRoom(activeRoomId).catch(() => {});

    return () => {
      client.unsubscribeRoom(activeRoomId).catch(() => {});
    };
  }, [activeRoomId, wsClient]);

  /**
   * Fetch reactions for a batch of messages via WS (with HTTP fallback),
   * then merge them into the messages state.  Fires-and-forgets so it
   * does not block the caller.
   */
  const thisLoadReactions = (
    roomId: string,
    client: RoomWsClient,
    msgs: MessageWithMeta[],
  ) => {
    const msgIds = msgs.map((m) => m.id);
    if (msgIds.length === 0) return;
    const doLoad = async () => {
      let reactionResults: ReactionListData[];
      if (client.getStatus() === 'open') {
        try {
          reactionResults = await client.reactionListBatchWs(roomId, msgIds);
        } catch {
          reactionResults = await client.reactionListBatch(roomId, msgIds);
        }
      } else {
        reactionResults = await client.reactionListBatch(roomId, msgIds);
      }
      const reactionMap = new Map<string, ReactionListData['reactions']>();
      for (const result of reactionResults) {
        if (result.reactions.length > 0) {
          reactionMap.set(result.message_id, result.reactions);
        }
      }
      if (reactionMap.size > 0) {
        setMessages((prev) =>
          prev.map((m) =>
            reactionMap.has(m.id) ? { ...m, reactions: reactionMap.get(m.id) } : m,
          ),
        );
      }
    };
    doLoad().catch(() => {});
  };

  const loadMore = useCallback(
    async (cursor?: number | null) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;

      // Cancel any in-flight request
      loadMessagesAbortRef.current?.abort();
      loadMessagesAbortRef.current = new AbortController();
      const abortController = loadMessagesAbortRef.current;

      setIsLoadingMore(true);
      try {
        const isInitial = cursor === null || cursor === undefined;
        const limit = isInitial ? 200 : 50;

        // Try WebSocket first; fall back to HTTP on failure
        let resp: import('@/lib/room-ws-client').RoomMessageListResponse;
        if (client.getStatus() === 'open') {
          try {
            resp = await client.messageListWs(activeRoomId, {
              beforeSeq: cursor ?? undefined,
              limit,
            });
          } catch {
            // WS failed — fall back to HTTP
            resp = await client.messageList(activeRoomId, {
              beforeSeq: cursor ?? undefined,
              limit,
            });
          }
        } else {
          resp = await client.messageList(activeRoomId, {
            beforeSeq: cursor ?? undefined,
            limit,
          });
        }

        if (abortController.signal.aborted) return;

        const newMessages = resp.messages.map((m) => ({
          ...m,
          thread_id: m.thread,
          display_content: m.content,
          is_streaming: false,
        }));

        setMessages((prev) => {
          if (abortController.signal.aborted) return prev;
          if (isInitial) {
            setIsTransitioningRoom(false);
            return newMessages;
          }
          const existingIds = new Set(prev.map((m) => m.id));
          const filtered = newMessages.filter((m) => !existingIds.has(m.id));
          let merged = [...filtered, ...prev];
          merged.sort((a, b) => a.seq - b.seq);
          if (merged.length > MAX_MESSAGES_IN_MEMORY) {
            merged = merged.slice(-MAX_MESSAGES_IN_MEMORY);
          }
          return merged;
        });

        if (resp.messages.length < limit) {
          setIsHistoryLoaded(true);
        }
        setNextCursor(resp.messages.length > 0 ? resp.messages[resp.messages.length - 1].seq : null);

        // Fetch reactions for all loaded messages
        thisLoadReactions(activeRoomId, client, newMessages);
      } catch (error) {
        if (abortController.signal.aborted) return;
        handleRoomError('Load messages', error);
      } finally {
        setIsLoadingMore(false);
        loadMessagesAbortRef.current = null;
      }
    },
    [activeRoomId],
  );

  const [members, setMembers] = useState<RoomMemberResponse[]>([]);
  const [membersLoading, setMembersLoading] = useState(false);

  const [pins, setPins] = useState<RoomPinResponse[]>([]);
  const [pinsLoading, setPinsLoading] = useState(false);

  const [threads, setThreads] = useState<RoomThreadResponse[]>([]);

  // User presence map: user_id -> status
  const [presence, setPresence] = useState<PresenceMap>({});

  // Typing users map: roomId -> Map<userId, { username, avatar_url, timeoutId }>
  const [typingUsers, setTypingUsers] = useState<Record<string, Record<string, { username: string; avatar_url?: string; timeoutId?: ReturnType<typeof setTimeout> }>>>({});



  const [streamingContent, setStreamingContent] = useState<Map<string, string>>(new Map());
  const [streamingThinkingContent, setStreamingThinkingContent] = useState<Map<string, string>>(new Map());
  const [activeAiStream, setActiveAiStream] = useState<{ message_id: string; display_name: string } | null>(null);

  // Streaming timeout: if no chunk received for 60s, force-end the stream
  // to prevent UI hanging forever when done=true is never delivered.
  const streamingTimersRef = useRef<Map<string, ReturnType<typeof setTimeout>>>(new Map());

  // Ref to latest streamingThinkingContent so done handler can read it (setState is async)
  const streamingThinkingContentRef = useRef<Map<string, string>>(new Map());

  const clearStreamingTimer = useCallback((msgId: string) => {
    const timer = streamingTimersRef.current.get(msgId);
    if (timer) {
      clearTimeout(timer);
      streamingTimersRef.current.delete(msgId);
    }
  }, []);

  const startStreamingTimer = useCallback((msgId: string) => {
    clearStreamingTimer(msgId);
    const timer = setTimeout(() => {
      // Force-end: mark message as not-streaming and keep whatever content we have
      setActiveAiStream((prev) => prev?.message_id === msgId ? null : prev);
      setStreamingContent((prev) => {
        prev.delete(msgId);
        return new Map(prev);
      });
      setStreamingThinkingContent((prev) => {
        prev.delete(msgId);
        return new Map(prev);
      });
      streamingThinkingContentRef.current.delete(msgId);
      setMessages((prev) =>
        prev.map((m) =>
          m.id === msgId && m.is_streaming
            ? { ...m, is_streaming: false, content: m.content || '[Stream timed out — no completion signal received]' }
            : m,
        ),
      );
      streamingTimersRef.current.delete(msgId);
    }, 60000);
    streamingTimersRef.current.set(msgId, timer);
  }, []);

  // Project repos for @repository: mention suggestions
  const [projectRepos, setProjectRepos] = useState<ProjectRepositoryItem[]>([]);
  const [reposLoading, setReposLoading] = useState(false);
  // Room AI configs for @ai: mention suggestions
  const [roomAiConfigs, setRoomAiConfigs] = useState<RoomAiConfig[]>([]);
  const [aiConfigsLoading, setAiConfigsLoading] = useState(false);

  // ── Update WS token on existing client (instead of recreating client) ────────
  useEffect(() => {
    if (wsToken && wsClientRef.current) {
      wsClientRef.current.setWsToken(wsToken);
    }
  }, [wsToken]);

  // ── Create WS client ONCE on mount ──────────────────────────────────────────
  // Recreating the client on wsToken change caused multiple invalid connections
  // (Bug 1). Instead, create once and update the token in-place.
  useEffect(() => {
    const baseUrl = import.meta.env.VITE_API_BASE_URL ?? window.location.origin;
    const client = createRoomWsClient(
      baseUrl,
      {
        onRoomMessage: (payload) => {
          // Use ref to get current activeRoomId to avoid stale closure
          if (payload.room_id === activeRoomIdRef.current) {
            setMessages((prev) => {
              const existingIdx = prev.findIndex((m) => m.id === payload.id);
              if (existingIdx !== -1) {
                // Message already exists (e.g. created by streaming chunk) —
                // merge server-side fields (display_name, reactions) that the
                // chunk didn't have.
                const existing = prev[existingIdx];
                const needsUpdate =
                  (!existing.display_name && payload.display_name) ||
                  (payload.reactions !== undefined && existing.reactions === undefined);
                if (needsUpdate) {
                  const updated = [...prev];
                  updated[existingIdx] = {
                    ...existing,
                    display_name: payload.display_name ?? existing.display_name,
                    reactions: payload.reactions ?? existing.reactions,
                  };
                  return updated;
                }
                return prev;
              }
              // Replace optimistic message with server-confirmed one.
              // Streaming messages have seq=0 in optimistic state, so match by id instead.
              const optimisticIdx = prev.findIndex(
                (m) => m.isOptimistic && m.id === payload.id,
              );
              if (optimisticIdx !== -1) {
                const confirmed: MessageWithMeta = {
                  ...wsMessageToUiMessage(payload),
                  reactions: prev[optimisticIdx].reactions,
                };
                const next = [...prev];
                next[optimisticIdx] = confirmed;
                return next;
              }
              const newMsg = wsMessageToUiMessage(payload);
              let updated = [...prev, newMsg];
              updated.sort((a, b) => a.seq - b.seq);
              if (updated.length > MAX_MESSAGES_IN_MEMORY) {
                updated = updated.slice(-MAX_MESSAGES_IN_MEMORY);
              }
              return updated;
            });
          } else {
            // Buffer messages for non-active rooms (Bug 3 fix).
            // When user switches to that room, pending messages are merged.
            pendingRoomMessagesRef.current.set(payload.room_id, [
              ...pendingRoomMessagesRef.current.get(payload.room_id) ?? [],
              payload,
            ]);
          }
        },
    onAiStreamChunk: (chunk: { done: boolean; message_id: string; room_id: string; content: string; display_name?: string; chunk_type?: string }) => {
      const isToolCall = chunk.chunk_type === 'tool_call' || chunk.chunk_type === 'tool_result';

      if (chunk.done) {
        // Clear the timeout timer since stream completed normally
        clearStreamingTimer(chunk.message_id);
        // Set activeAiStream to null since streaming is done
        setActiveAiStream(null);
        // Clear streaming content maps
        setStreamingContent((prev) => {
          prev.delete(chunk.message_id);
          return new Map(prev);
        });
        setStreamingThinkingContent((prev) => {
          prev.delete(chunk.message_id);
          return new Map(prev);
        });
        // Finalize message: keep thinking_content from accumulator, set content from done chunk
        setMessages((prev) =>
          prev.map((m) => {
            if (m.id !== chunk.message_id) return m;
            // Get thinking_content from the accumulator before it was cleared
            const tc = streamingThinkingContentRef.current.get(chunk.message_id);
            return {
              ...m,
              content: chunk.content,
              display_content: chunk.content,
              is_streaming: false,
              thinking_content: tc ?? m.thinking_content,
              chunk_type: chunk.chunk_type,
            };
          }),
        );
      } else {
        // Reset the timeout timer on each chunk — stream is still alive
        startStreamingTimer(chunk.message_id);
        // Update activeAiStream for typing indicator (skip tool call / result)
        if (!isToolCall && chunk.display_name) {
          setActiveAiStream({ message_id: chunk.message_id, display_name: chunk.display_name });
        }

        if (chunk.chunk_type === 'thinking') {
          // Accumulate thinking content separately
          setStreamingThinkingContent((prev) => {
            const next = new Map(prev);
            const prevContent = next.get(chunk.message_id) ?? '';
            const newContent =
              prevContent === '' || !chunk.content.startsWith(prevContent)
                ? chunk.content
                : prevContent + chunk.content.slice(prevContent.length);
            next.set(chunk.message_id, newContent);
            // Sync ref for done handler access
            streamingThinkingContentRef.current = new Map(next);
            return next;
          });
          // Ensure message entry exists (with minimal content to show streaming state)
          setMessages((msgs) => {
            const idx = msgs.findIndex((m) => m.id === chunk.message_id);
            if (idx !== -1) return msgs;
            const newMsg: MessageWithMeta = {
              id: chunk.message_id,
              room: chunk.room_id,
              seq: 0,
              sender_type: 'ai',
              display_name: chunk.display_name,
              content: '',
              display_content: '',
              content_type: 'text',
              send_at: new Date().toISOString(),
              is_streaming: true,
              chunk_type: 'thinking',
            };
            return [...msgs, newMsg];
          });
        } else if (chunk.chunk_type === 'answer') {
          // Accumulate answer content (existing behavior)
          setStreamingContent((prev) => {
            const next = new Map(prev);
            const prevContent = next.get(chunk.message_id) ?? '';
            const newContent =
              prevContent === '' || !chunk.content.startsWith(prevContent)
                ? chunk.content
                : prevContent + chunk.content.slice(prevContent.length);
            next.set(chunk.message_id, newContent);
            setMessages((msgs) => {
              const idx = msgs.findIndex((m) => m.id === chunk.message_id);
              if (idx !== -1) {
                const m = msgs[idx];
                if (m.content === newContent && m.is_streaming === true) return msgs;
                const updated = [...msgs];
                updated[idx] = { ...m, content: newContent, display_content: newContent };
                return updated;
              }
              if (!newContent) return msgs;
              const newMsg: MessageWithMeta = {
                id: chunk.message_id,
                room: chunk.room_id,
                seq: 0,
                sender_type: 'ai',
                display_name: chunk.display_name,
                content: newContent,
                display_content: newContent,
                content_type: 'text',
                send_at: new Date().toISOString(),
                is_streaming: true,
                chunk_type: chunk.chunk_type,
              };
              return [...msgs, newMsg];
            });
            return next;
          });
        }
        // tool_call / tool_result: skip content update entirely — don't pollute display
      }
    },
        onRoomReactionUpdated: (payload: RoomReactionUpdatedPayload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          setMessages((prev) => {
            const existingIdx = prev.findIndex((m) => m.id === payload.message_id);
            if (existingIdx === -1) return prev;
            const updated = [...prev];
            updated[existingIdx] = { ...updated[existingIdx], reactions: payload.reactions };
            return updated;
          });
        },
        onMessageEdited: async (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          const client = wsClientRef.current;
          if (!client) return;

          let rollbackEditedAt: string | null = null;
          setMessages((prev) => {
            const msg = prev.find((m) => m.id === payload.message_id);
            rollbackEditedAt = msg?.edited_at ?? null;
            return prev.map((m) =>
              m.id === payload.message_id ? { ...m, edited_at: payload.edited_at } : m,
            );
          });

          try {
            const updatedMsg = await client.messageGet(payload.message_id);
            if (!updatedMsg) return;
            setMessages((prev) =>
              prev.map((m) =>
                m.id === payload.message_id
                  ? {
                      ...m,
                      content: updatedMsg.content,
                      display_content: updatedMsg.content,
                      edited_at: payload.edited_at,
                    }
                  : m,
              ),
            );
          } catch {
            if (rollbackEditedAt !== null) {
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === payload.message_id ? { ...m, edited_at: rollbackEditedAt! } : m,
                ),
              );
            }
          }
        },
        onMessageRevoked: async (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          setMessages((prev) =>
            prev.map((m) =>
              m.id === payload.message_id
                ? {
                    ...m,
                    revoked: payload.revoked_at,
                    revoked_by: payload.revoked_by,
                    content: '',
                    display_content: '',
                  }
                : m,
            ),
          );
        },
        onMessagePinned: async (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          const client = wsClientRef.current;
          if (!client) return;
          try {
            const pins = await client.pinList(payload.room_id);
            setPins(pins.map((p) => ({
              room: p.room,
              message: p.message,
              pinned_by: p.pinned_by,
              pinned_at: p.pinned_at,
            })));
          } catch {
            // Silently ignore
          }
        },
        onMessageUnpinned: async (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          setPins((prev) => prev.filter((p) => p.message !== payload.message_id));
        },
        onUserPresence: (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          setPresence((prev) => ({ ...prev, [payload.user_id]: payload.status }));
        },
        onTypingStart: (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          if (payload.user_id === user?.uid) return;
          setTypingUsers((prev) => {
            const roomMap = prev[payload.room_id] ?? {};
            // Clear existing timeout for this user
            const existing = roomMap[payload.user_id];
            if (existing?.timeoutId) clearTimeout(existing.timeoutId);
            const timeoutId = setTimeout(() => {
              setTypingUsers((p) => {
                const rm = { ...p[payload.room_id] };
                delete rm[payload.user_id];
                return { ...p, [payload.room_id]: rm };
              });
            }, 4000);
            const next = {
              ...prev,
              [payload.room_id]: {
                ...roomMap,
                [payload.user_id]: { username: payload.username, avatar_url: payload.avatar_url, timeoutId },
              },
            };
            return next;
          });
        },
        onTypingStop: (payload) => {
          if (payload.room_id !== activeRoomIdRef.current) return;
          setTypingUsers((prev) => {
            const roomMap = prev[payload.room_id] ?? {};
            const existing = roomMap[payload.user_id];
            if (existing?.timeoutId) clearTimeout(existing.timeoutId);
            const newRoomMap = { ...roomMap };
            delete newRoomMap[payload.user_id];
            return { ...prev, [payload.room_id]: newRoomMap };
          });
        },
        onStatusChange: (status) => {
          setWsStatus(status);
          if (status === 'closed' || status === 'error') {
            setWsError('Connection lost');
          } else {
            setWsError(null);
          }
        },
        onError: (error) => {
          setWsError(error.message);
        },
      },
    );

    setWsClient(client);
    wsClientRef.current = client;

    // Connect immediately — connect() fetches its own token if needed
    client.connect().catch((e) => {
      console.error('[RoomContext] WS connect error:', e);
    });

    return () => {
      client.disconnect();  // Intentional disconnect on unmount — no reconnect
      wsClientRef.current = null;
      // Clear all streaming timeout timers
      for (const timer of streamingTimersRef.current.values()) {
        clearTimeout(timer);
      }
      streamingTimersRef.current.clear();
    };
  }, []);  // ← empty deps: create once on mount

  const connectWs = useCallback(async () => {
    const client = wsClientRef.current;
    if (!client) return;
    try {
      await client.connect();
      setWsError(null);
    } catch (e) {
      setWsError(e instanceof Error ? e.message : 'Failed to connect');
    }
  }, []);

  const disconnectWs = useCallback(() => {
    wsClientRef.current?.disconnect();
  }, []);

  const fetchRooms = useCallback(async () => {
    if (!projectName) {
      setRooms([]);
      return;
    }
    setRoomsLoading(true);
    setRoomsError(null);
    try {
      const resp = await restRoomList({ path: { project_name: projectName } });
      const data = resp.data?.data;
      if (Array.isArray(data)) {
        setRooms(data.map((r) => ({ ...r, category_info: null })));
      } else {
        setRooms([]);
      }
    } catch (err) {
      setRoomsError(err instanceof Error ? err : new Error('Failed to load rooms'));
    } finally {
      setRoomsLoading(false);
    }
  }, [projectName]);

  useEffect(() => {
    fetchRooms();
  }, [fetchRooms]);

  const fetchCategories = useCallback(async () => {
    if (!projectName) return;
    setCategoriesLoading(true);
    try {
      const resp = await restCategoryList({ path: { project_name: projectName } });
      const data = resp.data?.data;
      setCategories(Array.isArray(data) ? data : []);
    } catch (error) {
      handleRoomError('Load categories', error);
    } finally {
      setCategoriesLoading(false);
    }
  }, [projectName]);

  useEffect(() => {
    fetchCategories();
  }, [fetchCategories]);

  const createCategory = useCallback(
    async (name: string, position?: number) => {
      const client = wsClientRef.current;
      if (!projectName || !client) throw new Error('No project');
      const cat = await client.categoryCreate(projectName, name, position);
      setCategories((prev) => [...prev, cat]);
      return cat;
    },
    [projectName],
  );

  const updateCategory = useCallback(
    async (id: string, name?: string, position?: number) => {
      const client = wsClientRef.current;
      if (!client) return;
      await client.categoryUpdate(id, { name, position });
      setCategories((prev) =>
        prev.map((c) =>
          c.id === id
            ? { ...c, ...(name !== undefined && { name }), ...(position !== undefined && { position }) }
            : c,
        ),
      );
    },
    [],
  );

  const deleteCategory = useCallback(
    async (id: string) => {
      const client = wsClientRef.current;
      if (!client) return;
      await client.categoryDelete(id);
      setCategories((prev) => prev.filter((c) => c.id !== id));
    },
    [],
  );

  const setActiveRoom = useCallback(
    (roomId: string | null) => {
      setActiveRoomId(roomId);
      if (roomId) {
        navigate(`../room/${roomId}`, { replace: true });
      } else {
        navigate('../room', { replace: true });
      }
    },
    [navigate],
  );

  useEffect(() => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) {
      setActiveRoomState(null);
      return;
    }
    client
      .roomGet(activeRoomId)
      .then((resp) => resp && setActiveRoomState(resp))
      .catch(() => setActiveRoomState(null));
  }, [activeRoomId]);

  const fetchMembers = useCallback(async () => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) {
      setMembers([]);
      return;
    }
    setMembersLoading(true);
    try {
      const resp = await client.memberList(activeRoomId);
      setMembers(resp.map((m) => ({ ...m, do_not_disturb: false })));
    } catch (error) {
      handleRoomError('Load members', error);
      setMembers([]);
    } finally {
      setMembersLoading(false);
    }
  }, [activeRoomId]);

  useEffect(() => {
    fetchMembers();
  }, [fetchMembers]);

  const addMember = useCallback(
    async (userId: string, role = 'member') => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;

      const tempMember: RoomMemberResponse = {
        room: activeRoomId,
        user: userId,
        role,
        user_info: undefined,
        do_not_disturb: false,
      };
      setMembers((prev) => [...prev, tempMember]);

      try {
        await client.memberAdd(activeRoomId, userId, role);
        await fetchMembers();
      } catch (error) {
        handleRoomError('Add member', error);
        setMembers((prev) => prev.filter((m) => m.user !== userId));
      }
    },
    [activeRoomId, fetchMembers],
  );

  const removeMember = useCallback(
    async (userId: string) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;
      await client.memberRemove(activeRoomId, userId);
      setMembers((prev) => prev.filter((m) => m.user !== userId));
    },
    [activeRoomId],
  );

  const updateMemberRole = useCallback(
    async (userId: string, role: string) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;
      await client.memberUpdateRole(activeRoomId, userId, role);
      setMembers((prev) => prev.map((m) => (m.user === userId ? { ...m, role } : m)));
    },
    [activeRoomId],
  );

  // Guard against double-sending while a previous send is in-flight.
  // Without this, rapid clicking can queue multiple optimistic messages and
  // create duplicate sends on the server.
  const sendingRef = useRef(false);

  const sendMessage = useCallback(
    async (content: string, contentType = 'text', inReplyTo?: string, attachmentIds?: string[]) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;
      if (sendingRef.current) return;
      sendingRef.current = true;

      // Optimistic update: add message immediately so user sees it instantly
      const optimisticId = `optimistic-${crypto.randomUUID()}`;
      const optimisticMsg: MessageWithMeta = {
        id: optimisticId,
        room: activeRoomId,
        seq: 0,
        sender_type: 'member',
        sender_id: user?.uid ?? null,
        content,
        content_type: contentType,
        send_at: new Date().toISOString(),
        display_content: content,
        is_streaming: false,
        isOptimistic: true,
        thread: inReplyTo,
        thread_id: inReplyTo,
        in_reply_to: inReplyTo,
        reactions: [],
        attachment_ids: attachmentIds,
      };

      setMessages((prev) => [...prev, optimisticMsg]);

      try {
        const confirmedMsg = await client.messageCreate(activeRoomId, content, {
          contentType,
          inReplyTo,
          attachmentIds,
        });
        // Replace optimistic message with server-confirmed one
        setMessages((prev) => {
          const without = prev.filter((m) => m.id !== optimisticId);
          const confirmed: MessageWithMeta = {
            ...confirmedMsg,
            thread_id: confirmedMsg.thread,
            display_content: confirmedMsg.content,
            is_streaming: false,
            isOptimistic: false,
            reactions: [],
          };
          return [...without, confirmed];
        });
      } catch (err) {
        // Mark optimistic message as failed
        setMessages((prev) =>
          prev.map((m) =>
            m.id === optimisticId ? { ...m, isOptimisticError: true } : m,
          ),
        );
        handleRoomError('Send message', err);
      } finally {
        sendingRef.current = false;
      }
    },
    [activeRoomId, user],
  );

  const editMessage = useCallback(
    async (messageId: string, content: string) => {
      const client = wsClientRef.current;
      if (!client) return;

      let rollbackContent: string | null = null;
      setMessages((prev) => {
        const msg = prev.find((m) => m.id === messageId);
        rollbackContent = msg?.content ?? null;
        return prev.map((m) =>
          m.id === messageId ? { ...m, content, display_content: content } : m,
        );
      });

      try {
        if (client.getStatus() === 'open') {
          try {
            await client.messageUpdateWs(messageId, content);
          } catch {
            await client.messageUpdate(messageId, content);
          }
        } else {
          await client.messageUpdate(messageId, content);
        }
      } catch (err) {
        if (rollbackContent !== null) {
          setMessages((prev) =>
            prev.map((m) =>
              m.id === messageId ? { ...m, content: rollbackContent!, display_content: rollbackContent! } : m,
            ),
          );
        }
        handleRoomError('Edit message', err);
      }
    },
    [],
  );

  const revokeMessage = useCallback(
    async (messageId: string) => {
      const client = wsClientRef.current;
      if (!client) return;

      let rollbackMsg: MessageWithMeta | null = null;
      setMessages((prev) => {
        rollbackMsg = prev.find((m) => m.id === messageId) ?? null;
        return prev.filter((m) => m.id !== messageId);
      });

      try {
        if (client.getStatus() === 'open') {
          try {
            await client.messageRevokeWs(messageId);
          } catch {
            await client.messageRevoke(messageId);
          }
        } else {
          await client.messageRevoke(messageId);
        }
      } catch (err) {
        if (rollbackMsg) {
          setMessages((prev) => [...prev, rollbackMsg!]);
        }
        handleRoomError('Delete message', err);
      }
    },
    [],
  );

  const updateReadSeq = useCallback(
    async (seq: number) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !user || !client) return;
      try {
        await client.memberSetReadSeq(activeRoomId, seq);
      } catch (error) {
        console.warn('[RoomContext] Failed to update read position:', error);
      }
    },
    [activeRoomId, user],
  );

  const fetchPins = useCallback(async () => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) {
      setPins([]);
      return;
    }
    setPinsLoading(true);
    try {
      const resp = await client.pinList(activeRoomId);
      setPins(
        resp.map((p) => ({
          room: p.room,
          message: p.message,
          pinned_by: p.pinned_by,
          pinned_at: p.pinned_at,
        })),
      );
    } catch (error) {
      handleRoomError('Load pinned messages', error);
      setPins([]);
    } finally {
      setPinsLoading(false);
    }
  }, [activeRoomId]);

  useEffect(() => {
    fetchPins();
  }, [fetchPins]);

  const pinMessage = useCallback(
    async (messageId: string) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;

      const tempPin: RoomPinResponse = {
        room: activeRoomId,
        message: messageId,
        pinned_by: '',
        pinned_at: new Date().toISOString(),
      };
      setPins((prev) => [...prev, tempPin]);

      try {
        await client.pinAdd(activeRoomId, messageId);
        await fetchPins();
      } catch (error) {
        handleRoomError('Pin message', error);
        setPins((prev) => prev.filter((p) => p.message !== messageId));
      }
    },
    [activeRoomId, fetchPins],
  );

  const unpinMessage = useCallback(
    async (messageId: string) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) return;

      setPins((prev) => prev.filter((p) => p.message !== messageId));

      try {
        await client.pinRemove(activeRoomId, messageId);
      } catch (error) {
        handleRoomError('Unpin message', error);
        fetchPins();
      }
    },
    [activeRoomId, fetchPins],
  );

  const fetchThreads = useCallback(async () => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) {
      setThreads([]);
      return;
    }
    try {
      const resp = await client.threadList(activeRoomId);
      setThreads(resp);
    } catch (error) {
      handleRoomError('Load threads', error);
      setThreads([]);
    }
  }, [activeRoomId]);

  useEffect(() => {
    fetchThreads();
  }, [fetchThreads]);

  const createThread = useCallback(
    async (parentSeq: number) => {
      const client = wsClientRef.current;
      if (!activeRoomId || !client) throw new Error('No active room');
      const thread = await client.threadCreate(activeRoomId, parentSeq);
      setThreads((prev) => [...prev, thread]);
      return thread;
    },
    [activeRoomId],
  );

  // Fetch project repos for @repository: mention suggestions
  const fetchProjectRepos = useCallback(async () => {
    if (!projectName) {
      setProjectRepos([]);
      return;
    }
    setReposLoading(true);
    try {
      const baseUrl = import.meta.env.VITE_API_BASE_URL ?? window.location.origin;
      const resp = await fetch(`${baseUrl}/api/projects/${encodeURIComponent(projectName)}/repos`);
      if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
      const json: { data?: { items?: ProjectRepositoryItem[] } } = await resp.json();
      setProjectRepos(json.data?.items ?? []);
    } catch {
      setProjectRepos([]);
    } finally {
      setReposLoading(false);
    }
  }, [projectName]);

  // Fetch room AI configs for @ai: mention suggestions
  const fetchRoomAiConfigs = useCallback(async () => {
    const client = wsClientRef.current;
    if (!activeRoomId || !client) {
      setRoomAiConfigs([]);
      return;
    }
    setAiConfigsLoading(true);
    try {
      const configs = await client.aiList(activeRoomId);
      setRoomAiConfigs(
        configs.map((cfg) => ({
          model: cfg.model,
          modelName: cfg.model_name,
        })),
      );
    } catch {
      setRoomAiConfigs([]);
    } finally {
      setAiConfigsLoading(false);
    }
  }, [activeRoomId]);

  useEffect(() => {
    fetchProjectRepos();
  }, [fetchProjectRepos]);

  useEffect(() => {
    fetchRoomAiConfigs();
  }, [fetchRoomAiConfigs]);

  const createRoom = useCallback(
    async (name: string, isPublic: boolean, categoryId?: string) => {
      const client = wsClientRef.current;
      if (!projectName || !client) throw new Error('No project');
      const room = await client.roomCreate(projectName, name, isPublic, categoryId);
      setRooms((prev) => [...prev, { ...room, category_info: null }]);
      return room;
    },
    [projectName],
  );

  const updateRoom = useCallback(
    async (roomId: string, name?: string, isPublic?: boolean, categoryId?: string) => {
      const client = wsClientRef.current;
      if (!client) return;

      // Save previous values for rollback (use functional update to avoid closure issues)
      let rollbackRooms: RoomWithCategory[] | null = null;
      let rollbackActiveRoom: RoomResponse | null = null;

      // Capture current state for rollback
      setRooms((prev) => {
        rollbackRooms = prev;
        return prev.map((r) =>
          r.id === roomId
            ? {
                ...r,
                ...(name !== undefined && { room_name: name }),
                ...(isPublic !== undefined && { public: isPublic }),
                ...(categoryId !== undefined && { category: categoryId }),
              }
            : r,
        );
      });
      if (activeRoomId === roomId) {
        setActiveRoomState((prev) => {
          rollbackActiveRoom = prev;
          return prev
            ? {
                ...prev,
                ...(name !== undefined && { room_name: name }),
                ...(isPublic !== undefined && { public: isPublic }),
                ...(categoryId !== undefined && { category: categoryId }),
              }
            : prev;
        });
      }

      try {
        await client.roomUpdate(roomId, { roomName: name, isPublic, categoryId });
      } catch (error) {
        handleRoomError('Update room', error);
        // Use captured snapshot for rollback, not the old value in closure
        if (rollbackRooms) {
          setRooms(rollbackRooms);
        }
        if (activeRoomId === roomId && rollbackActiveRoom) {
          setActiveRoomState(rollbackActiveRoom);
        }
      }
    },
    [activeRoomId],
  );

  const deleteRoom = useCallback(
    async (roomId: string) => {
      const client = wsClientRef.current;
      if (!client) return;
      await client.roomDelete(roomId);
      setRooms((prev) => prev.filter((r) => r.id !== roomId));
      if (activeRoomId === roomId) {
        setActiveRoomId(null);
        setActiveRoom(null);
      }
    },
    [activeRoomId],
  );

  const value = useMemo<RoomContextValue>(
    () => ({
      wsStatus,
      wsError,
      wsClient: wsClientRef.current,
      connectWs,
      disconnectWs,
      rooms: roomsWithCategory,
      roomsLoading,
      roomsError,
      refreshRooms: fetchRooms,
      categories,
      categoriesLoading,
      refreshCategories: fetchCategories,
      createCategory,
      updateCategory,
      deleteCategory,
      activeRoom,
      activeRoomId,
      setActiveRoom,
      messages,
      messagesLoading: false,
      isHistoryLoaded,
      isLoadingMore,
      isTransitioningRoom,
      nextCursor,
      loadMore,
      sendMessage,
      editMessage,
      revokeMessage,
      updateReadSeq,
      members,
      membersLoading,
      refreshMembers: fetchMembers,
      addMember,
      removeMember,
      updateMemberRole,
      pins,
      pinsLoading,
      refreshPins: fetchPins,
      pinMessage,
      unpinMessage,
      presence,
      threads,
      refreshThreads: fetchThreads,
      createThread,
      createRoom,
      updateRoom,
      deleteRoom,
      streamingMessages: streamingContent,
      streamingThinkingContent,
      activeAiStream,
      projectRepos,
      reposLoading,
      roomAiConfigs,
      aiConfigsLoading,
      typingUsers,
    }),
    [
      wsStatus,
      wsError,
      connectWs,
      disconnectWs,
      wsClientRef.current,
      roomsWithCategory,
      roomsLoading,
      roomsError,
      fetchRooms,
      categories,
      categoriesLoading,
      fetchCategories,
      createCategory,
      updateCategory,
      deleteCategory,
      activeRoom,
      activeRoomId,
      setActiveRoom,
      messages,
      isHistoryLoaded,
      isLoadingMore,
      isTransitioningRoom,
      nextCursor,
      loadMore,
      sendMessage,
      editMessage,
      revokeMessage,
      updateReadSeq,
      members,
      membersLoading,
      fetchMembers,
      addMember,
      removeMember,
      updateMemberRole,
      pins,
      pinsLoading,
      fetchPins,
      pinMessage,
      unpinMessage,
      presence,
      threads,
      fetchThreads,
      createThread,
      createRoom,
      updateRoom,
      deleteRoom,
      streamingContent,
      streamingThinkingContent,
      activeAiStream,
      projectRepos,
      reposLoading,
      roomAiConfigs,
      aiConfigsLoading,
      typingUsers,
    ],
  );

  return <RoomContext.Provider value={value}>{children}</RoomContext.Provider>;
}

export function useRoom() {
  const ctx = useContext(RoomContext);
  if (!ctx) throw new Error('useRoom must be used within RoomProvider');
  return ctx;
}