kandev/apps/web/lib/ws/use-websocket.tsx

'use client';

import { useEffect, useRef } from 'react';
import type { StoreApi } from 'zustand';
import { WebSocketClient } from '@/lib/ws/client';
import { registerWsHandlers } from '@/lib/ws/router';
import type { AppState } from '@/lib/state/store';
import { setWebSocketClient } from '@/lib/ws/connection';

export function useWebSocket(store: StoreApi<AppState>, url: string) {
  const clientRef = useRef<WebSocketClient | null>(null);

  useEffect(() => {
    const client = new WebSocketClient(
      url,
      (status) => {
        const setConnectionStatus = store.getState().setConnectionStatus;
        switch (status) {
          case 'connecting':
            setConnectionStatus('connecting', null);
            break;
          case 'open':
            setConnectionStatus('connected', null);
            client.subscribeUser();
            break;
          case 'reconnecting':
            setConnectionStatus('reconnecting', null);
            break;
          case 'error':
            setConnectionStatus('error', 'WebSocket connection failed');
            break;
          case 'closed':
          case 'idle':
          default:
            setConnectionStatus('disconnected', null);
            break;
        }
      },
      {
        enabled: true,
        maxAttempts: 10,
        initialDelay: 1000,
        maxDelay: 30000,
        backoffMultiplier: 1.5,
      }
    );
    clientRef.current = client;
    client.connect();
    setWebSocketClient(client);

    const handlers = registerWsHandlers(store);
    const unsubscribers = Object.entries(handlers).map(([type, handler]) =>
      client.on(type as keyof typeof handlers, handler as never)
    );

    return () => {
      unsubscribers.forEach((unsubscribe) => unsubscribe());
      client.disconnect();
      setWebSocketClient(null);
    };
  }, [store, url]);

  return clientRef;
}