Messenger/web/src/hooks/useRealtime.ts

import { useEffect, useMemo, useRef } from "react";
import { useAuthStore } from "../store/authStore";
import { useChatStore } from "../store/chatStore";
import type { Message } from "../chat/types";
import { buildWsUrl } from "../utils/ws";

interface RealtimeEnvelope {
  event: string;
  payload: Record<string, unknown>;
  timestamp: string;
}

export function useRealtime() {
  const accessToken = useAuthStore((s) => s.accessToken);
  const meId = useAuthStore((s) => s.me?.id ?? null);
  const typingByChat = useRef<Record<number, Set<number>>>({});
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | null>(null);
  const heartbeatIntervalRef = useRef<number | null>(null);
  const watchdogIntervalRef = useRef<number | null>(null);
  const lastPongAtRef = useRef<number>(Date.now());
  const reconnectAttemptsRef = useRef(0);
  const manualCloseRef = useRef(false);

  const wsUrl = useMemo(() => {
    return accessToken ? buildWsUrl(accessToken) : null;
  }, [accessToken]);

  useEffect(() => {
    if (!wsUrl) {
      manualCloseRef.current = true;
      if (reconnectTimeoutRef.current !== null) {
        window.clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      wsRef.current?.close();
      wsRef.current = null;
      return;
    }
    manualCloseRef.current = false;

    const connect = () => {
      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onopen = () => {
        reconnectAttemptsRef.current = 0;
        lastPongAtRef.current = Date.now();
        if (heartbeatIntervalRef.current !== null) {
          window.clearInterval(heartbeatIntervalRef.current);
        }
        if (watchdogIntervalRef.current !== null) {
          window.clearInterval(watchdogIntervalRef.current);
        }
        heartbeatIntervalRef.current = window.setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ event: "ping", payload: {} }));
          }
        }, 20000);
        watchdogIntervalRef.current = window.setInterval(() => {
          if (Date.now() - lastPongAtRef.current > 65000 && ws.readyState === WebSocket.OPEN) {
            ws.close();
          }
        }, 15000);
        void useChatStore.getState().loadChats();
      };

      ws.onmessage = (messageEvent) => {
        let event: RealtimeEnvelope;
        try {
          event = JSON.parse(messageEvent.data) as RealtimeEnvelope;
        } catch {
          return;
        }
        const chatStore = useChatStore.getState();
        const authStore = useAuthStore.getState();
        if (event.event === "receive_message") {
          const chatId = Number(event.payload.chat_id);
          const message = event.payload.message as Message;
          const clientMessageId = event.payload.client_message_id as string | undefined;
          if (!Number.isFinite(chatId) || !message?.id) {
            return;
          }
          let wasInserted = false;
          if (clientMessageId && message.sender_id === authStore.me?.id) {
            chatStore.confirmMessageByClientId(chatId, clientMessageId, message);
          } else {
            wasInserted = chatStore.prependMessage(chatId, message);
          }
          if (message.sender_id !== authStore.me?.id) {
            ws.send(JSON.stringify({ event: "message_delivered", payload: { chat_id: chatId, message_id: message.id } }));
            if (chatId === chatStore.activeChatId) {
              ws.send(JSON.stringify({ event: "message_read", payload: { chat_id: chatId, message_id: message.id } }));
              chatStore.clearUnread(chatId);
            } else if (wasInserted) {
              chatStore.incrementUnread(chatId);
            }
          }
          if (!chatStore.chats.some((chat) => chat.id === chatId)) {
            void chatStore.loadChats();
          }
        }
        if (event.event === "pong") {
          lastPongAtRef.current = Date.now();
        }
        if (event.event === "typing_start") {
          const chatId = Number(event.payload.chat_id);
          const userId = Number(event.payload.user_id);
          if (!Number.isFinite(chatId) || !Number.isFinite(userId) || userId === authStore.me?.id) {
            return;
          }
          if (!typingByChat.current[chatId]) {
            typingByChat.current[chatId] = new Set<number>();
          }
          typingByChat.current[chatId].add(userId);
          chatStore.setTypingUsers(chatId, [...typingByChat.current[chatId]]);
        }
        if (event.event === "typing_stop") {
          const chatId = Number(event.payload.chat_id);
          const userId = Number(event.payload.user_id);
          if (!Number.isFinite(chatId) || !Number.isFinite(userId)) {
            return;
          }
          typingByChat.current[chatId]?.delete(userId);
          chatStore.setTypingUsers(chatId, [...(typingByChat.current[chatId] ?? [])]);
        }
        if (event.event === "message_delivered") {
          const chatId = Number(event.payload.chat_id);
          const messageId = Number(event.payload.message_id);
          const lastDeliveredMessageId = Number(event.payload.last_delivered_message_id);
          const userId = Number(event.payload.user_id);
          if (!Number.isFinite(chatId) || !Number.isFinite(messageId) || !Number.isFinite(userId)) {
            return;
          }
          if (userId !== authStore.me?.id) {
            const maxId = Number.isFinite(lastDeliveredMessageId) ? lastDeliveredMessageId : messageId;
            chatStore.setMessageDeliveryStatusUpTo(chatId, maxId, "delivered", authStore.me?.id ?? -1);
          }
        }
        if (event.event === "message_read") {
          const chatId = Number(event.payload.chat_id);
          const messageId = Number(event.payload.message_id);
          const lastReadMessageId = Number(event.payload.last_read_message_id);
          const userId = Number(event.payload.user_id);
          if (!Number.isFinite(chatId) || !Number.isFinite(messageId) || !Number.isFinite(userId)) {
            return;
          }
          if (userId !== authStore.me?.id) {
            const maxId = Number.isFinite(lastReadMessageId) ? lastReadMessageId : messageId;
            chatStore.setMessageDeliveryStatusUpTo(chatId, maxId, "read", authStore.me?.id ?? -1);
          }
        }
        if (event.event === "user_online" || event.event === "user_offline") {
          const chatId = Number(event.payload.chat_id);
          const userId = Number(event.payload.user_id);
          const isOnline = Boolean(event.payload.is_online);
          const lastSeenAtRaw = event.payload.last_seen_at;
          const lastSeenAt = typeof lastSeenAtRaw === "string" ? lastSeenAtRaw : undefined;
          if (!Number.isFinite(chatId) || !Number.isFinite(userId)) {
            return;
          }
          if (userId !== authStore.me?.id) {
            chatStore.applyPresenceEvent(chatId, userId, isOnline, lastSeenAt);
          }
        }
      };

      ws.onclose = () => {
        if (heartbeatIntervalRef.current !== null) {
          window.clearInterval(heartbeatIntervalRef.current);
          heartbeatIntervalRef.current = null;
        }
        if (watchdogIntervalRef.current !== null) {
          window.clearInterval(watchdogIntervalRef.current);
          watchdogIntervalRef.current = null;
        }
        if (manualCloseRef.current) {
          return;
        }
        reconnectAttemptsRef.current += 1;
        const delay = Math.min(15000, 1000 * 2 ** Math.min(reconnectAttemptsRef.current, 4));
        reconnectTimeoutRef.current = window.setTimeout(connect, delay);
      };

      ws.onerror = () => {
        ws.close();
      };
    };

    connect();

    return () => {
      manualCloseRef.current = true;
      if (heartbeatIntervalRef.current !== null) {
        window.clearInterval(heartbeatIntervalRef.current);
        heartbeatIntervalRef.current = null;
      }
      if (watchdogIntervalRef.current !== null) {
        window.clearInterval(watchdogIntervalRef.current);
        watchdogIntervalRef.current = null;
      }
      if (reconnectTimeoutRef.current !== null) {
        window.clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      wsRef.current?.close();
      wsRef.current = null;
      typingByChat.current = {};
      useChatStore.setState({ typingByChat: {} });
    };
  }, [wsUrl, meId]);

  return null;
}