dumpsterChat/web/src/stores/ws.ts

import { create } from 'zustand';
import { useMessageStore } from './message.ts';
import { useChannelStore } from './channel.ts';
import { useServerStore } from './server.ts';
import { usePresenceStore } from './presence.ts';
import { useTypingStore } from './typing.ts';
import { useVoiceStore } from './voice.ts';
import { useAuthStore } from './auth.ts';
import { useConversationStore } from './conversation.ts';
import { useReadStatesStore } from './readStates.ts';
import type { Message } from './message.ts';
import type { ConversationMessage } from './conversation.ts';
import type { Channel } from './channel.ts';
import type { Server } from './server.ts';
import type { UserStatus } from './auth.ts';

type UnknownPayload = Record<string, unknown>;

export interface WsEvent {
  type: string;
  data?: UnknownPayload;
}

interface WebSocketState {
  socket: WebSocket | null;
  connected: boolean;
  connect: () => void;
  disconnect: () => void;
  send: (event: WsEvent) => void;
}

let unreadCount = 0;

if (typeof window !== 'undefined') {
  window.addEventListener('focus', () => {
    unreadCount = 0;
    if (document.title.match(/^\(\d+\)\s/)) {
      document.title = document.title.replace(/^\(\d+\)\s/, '');
    }
  });
}

function getWsHost(): string {
  const { protocol, host } = window.location;
  const wsProtocol = protocol === 'https:' ? 'wss:' : 'ws:';
  return `${wsProtocol}//${host}`;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null;
}

function extractIds(payload: UnknownPayload | undefined): { channel_id?: string; conversation_id?: string; message_id: string } | null {
  if (!payload || typeof payload.message_id !== 'string') return null;
  if (typeof payload.channel_id === 'string') return { channel_id: payload.channel_id, message_id: payload.message_id };
  if (typeof payload.conversation_id === 'string') return { conversation_id: payload.conversation_id, message_id: payload.message_id };
  return null;
}

export const useWebSocketStore = create<WebSocketState>((set, get) => ({
  socket: null,
  connected: false,

  connect: () => {
    const existing = get().socket;
    if (existing && (existing.readyState === WebSocket.OPEN || existing.readyState === WebSocket.CONNECTING)) {
      return;
    }

    const socket = new WebSocket(`${getWsHost()}/ws`);

    let reconnectDelay = 1000;
    const maxReconnectDelay = 30000;
    let reconnectTimeout: number | null = null;

    const scheduleReconnect = () => {
      if (reconnectTimeout) {
        window.clearTimeout(reconnectTimeout);
      }
      reconnectTimeout = window.setTimeout(() => {
        if (!get().connected) {
          get().connect();
        }
      }, reconnectDelay);
      reconnectDelay = Math.min(reconnectDelay * 2, maxReconnectDelay);
    };

    socket.onopen = () => {
      // Cookie-based auth: browser sends session cookie automatically.
      // No need to send a token frame.
    };

    socket.onmessage = (event) => {
      let data: WsEvent;
      try {
        data = JSON.parse(event.data) as WsEvent;
      } catch {
        return;
      }

      if (data.type === 'ready') {
        set({ connected: true });
        reconnectDelay = 1000;
        return;
      }

      if (data.type === 'error') {
        console.error('ws error:', data);
        socket.close();
        return;
      }

      if (data.type === 'ping') {
        get().send({ type: 'pong', data: {} });
        return;
      }

      const addMessage = useMessageStore.getState().addMessage;
      const updateMessage = useMessageStore.getState().updateMessage;
      const removeMessage = useMessageStore.getState().removeMessage;
      const addReaction = useMessageStore.getState().addReaction;
      const removeReaction = useMessageStore.getState().removeReaction;
      const addChannel = useChannelStore.getState().addChannel;
      const updateChannel = useChannelStore.getState().updateChannel;
      const removeChannel = useChannelStore.getState().removeChannel;
      const updateServer = useServerStore.getState().updateServer;

      const payload = data.data || {};

      switch (data.type) {
        case 'MESSAGE_CREATE': {
          if (isRecord(payload)) {
            if (payload.conversation_id) {
              const msg = payload as unknown as ConversationMessage;
              useConversationStore.getState().addMessage(msg);
              // auto-mark DM as read if this conversation is active
              const activeConvId = useConversationStore.getState().activeConversationId;
              if (msg.conversation_id === activeConvId && document.hasFocus()) {
                useReadStatesStore.getState().markConvRead(msg.conversation_id, msg.id);
              }
            } else {
              const msg = payload as unknown as Message;
              addMessage(msg);
              
              // Desktop notification
              const currentUserId = useAuthStore.getState().user?.id;
              if (msg.author_id !== currentUserId && !document.hasFocus()) {
                // Update browser tab title
                unreadCount++;
                if (document.title.match(/^\(\d+\)\s/)) {
                  document.title = document.title.replace(/^\(\d+\)\s/, `(${unreadCount}) `);
                } else {
                  document.title = `(${unreadCount}) ` + document.title;
                }

                if (Notification.permission === 'granted') {
                  const title = msg.author_username ? `New message from ${msg.author_display_name || msg.author_username}` : 'New message';
                  const notification = new Notification(title, {
                    body: msg.content,
                    icon: '/icons/icon-192.png',
                  });
                  notification.onclick = () => {
                    window.focus();
                    notification.close();
                  };
                }
              }
            }
          }
          break;
        }
        case 'MESSAGE_UPDATE': {
          if (isRecord(payload)) {
            if (payload.conversation_id) {
               useConversationStore.getState().updateMessage(payload as unknown as ConversationMessage);
            } else {
               updateMessage(payload as unknown as Message);
            }
          }
          break;
        }
        case 'MESSAGE_DELETE': {
          const ids = extractIds(payload);
          if (ids) {
            if (ids.conversation_id) {
              useConversationStore.getState().deleteMessage(ids.conversation_id, ids.message_id);
            } else if (ids.channel_id) {
              removeMessage(ids.channel_id, ids.message_id);
            }
          }
          break;
        }
        case 'REACTION_ADD': {
          const channelId = typeof payload.channel_id === 'string' ? payload.channel_id : '';
          const convId = typeof payload.conversation_id === 'string' ? payload.conversation_id : '';
          const messageId = typeof payload.message_id === 'string' ? payload.message_id : '';
          const reaction = isRecord(payload.reaction) ? payload.reaction : null;
          if (messageId && reaction) {
            const emoji = typeof reaction.emoji === 'string' ? reaction.emoji : '';
            const userId = typeof reaction.user_id === 'string' ? reaction.user_id : '';
            if (emoji && userId) {
              if (convId) {
                useConversationStore.getState().addReaction(convId, messageId, emoji, userId);
              } else if (channelId) {
                addReaction(channelId, messageId, emoji, userId);
              }
            }
          }
          break;
        }
        case 'REACTION_REMOVE': {
          const channelId = typeof payload.channel_id === 'string' ? payload.channel_id : '';
          const convId = typeof payload.conversation_id === 'string' ? payload.conversation_id : '';
          const messageId = typeof payload.message_id === 'string' ? payload.message_id : '';
          const emoji = typeof payload.emoji === 'string' ? payload.emoji : '';
          const userId = typeof payload.user_id === 'string' ? payload.user_id : '';
          if (messageId && emoji && userId) {
            if (convId) {
              useConversationStore.getState().removeReaction(convId, messageId, emoji, userId);
            } else if (channelId) {
              removeReaction(channelId, messageId, emoji, userId);
            }
          }
          break;
        }
        case 'POLL_UPDATE': {
          const channelId = typeof payload.channel_id === 'string' ? payload.channel_id : '';
          const poll = isRecord(payload.poll) ? payload.poll : null;
          if (channelId && poll) {
            const updatePoll = useMessageStore.getState().updatePoll;
            updatePoll(channelId, poll as unknown as import('./message.ts').Poll);
          }
          break;
        }
        case 'CHANNEL_CREATE': {
          if (isRecord(payload)) addChannel(payload as unknown as Channel);
          break;
        }
        case 'CHANNEL_UPDATE': {
          if (isRecord(payload)) updateChannel(payload as unknown as Channel);
          break;
        }
        case 'CHANNEL_DELETE': {
          const channelId = typeof payload.channel_id === 'string' ? payload.channel_id : null;
          if (channelId) removeChannel(channelId);
          break;
        }
        case 'SERVER_UPDATE': {
          if (isRecord(payload)) updateServer(payload as unknown as Server);
          break;
        }
        case 'PRESENCE_UPDATE': {
          const { user_id, username, status } = payload as Record<string, string>;
          if (user_id && username && status) {
            usePresenceStore.getState().updatePresence(user_id, username, status as UserStatus);
          }
          break;
        }
        case 'TYPING_START': {
          const { channel_id, user_id, username } = payload as Record<string, string>;
          if (channel_id && user_id && username) {
            useTypingStore.getState()._handleTypingEvent({ channel_id, user_id, username });
          }
          break;
        }
        case 'VOICE_WHISPER': {
          // Forwarded from backend — only intended recipient gets this
          const { from_user_id, from_username, message } = payload as Record<string, string>;
          useVoiceStore.getState()._addWhisper({
            fromUserId: from_user_id || '',
            fromUsername: from_username || 'unknown',
            message: message || '',
            timestamp: Date.now(),
          });
          break;
        }
        default:
          break;
      }
    };

    socket.onclose = () => {
      set({ connected: false, socket: null });
      scheduleReconnect();
    };

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

    set({ socket });
  },

  disconnect: () => {
    const socket = get().socket;
    if (socket) {
      socket.close();
    }
    set({ socket: null, connected: false });
  },

  send: (event) => {
    const socket = get().socket;
    if (socket && socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify(event));
    }
  },
}));