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 me = useAuthStore((s) => s.me);
  const prependMessage = useChatStore((s) => s.prependMessage);
  const typingByChat = useRef<Record<number, Set<number>>>({});

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

  useEffect(() => {
    if (!wsUrl) {
      return;
    }
    const ws = new WebSocket(wsUrl);

    ws.onmessage = (messageEvent) => {
      const event: RealtimeEnvelope = JSON.parse(messageEvent.data);
      if (event.event === "receive_message") {
        const chatId = Number(event.payload.chat_id);
        const message = event.payload.message as Message;
        prependMessage(chatId, message);
      }
      if (event.event === "typing_start") {
        const chatId = Number(event.payload.chat_id);
        const userId = Number(event.payload.user_id);
        if (userId === me?.id) {
          return;
        }
        if (!typingByChat.current[chatId]) {
          typingByChat.current[chatId] = new Set<number>();
        }
        typingByChat.current[chatId].add(userId);
        useChatStore.getState().setTypingUsers(chatId, [...typingByChat.current[chatId]]);
      }
      if (event.event === "typing_stop") {
        const chatId = Number(event.payload.chat_id);
        const userId = Number(event.payload.user_id);
        typingByChat.current[chatId]?.delete(userId);
        useChatStore.getState().setTypingUsers(chatId, [...(typingByChat.current[chatId] ?? [])]);
      }
    };

    return () => ws.close();
  }, [wsUrl, prependMessage, me?.id]);

  return null;
}