wrongmove/frontend/src/hooks/useTaskWebSocket.ts

import { useCallback, useEffect, useRef, useState } from 'react';
import type { AuthUser } from '@/auth/types';
import type { TaskState, WSMessage } from '@/types';
import { WS_TASKS_PATH } from '@/constants';

const KEEPALIVE_MS = 30_000;
const MAX_RECONNECT_DELAY_MS = 30_000;

function wsUrl(token: string): string {
  const proto = window.location.protocol === 'https:' ? 'wss' : 'ws';
  return `${proto}://${window.location.host}${WS_TASKS_PATH}?token=${encodeURIComponent(token)}`;
}

export interface UseTaskWebSocketReturn {
  tasks: Record<string, TaskState>;
  isConnected: boolean;
  subscribe: (taskId: string) => void;
}

export function useTaskWebSocket(user: AuthUser | null): UseTaskWebSocketReturn {
  const [tasks, setTasks] = useState<Record<string, TaskState>>({});
  const [isConnected, setIsConnected] = useState(false);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttempt = useRef(0);
  const reconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const keepaliveTimer = useRef<ReturnType<typeof setInterval> | null>(null);
  const mountedRef = useRef(true);

  const clearTimers = useCallback(() => {
    if (reconnectTimer.current) {
      clearTimeout(reconnectTimer.current);
      reconnectTimer.current = null;
    }
    if (keepaliveTimer.current) {
      clearInterval(keepaliveTimer.current);
      keepaliveTimer.current = null;
    }
  }, []);

  const connect = useCallback(() => {
    if (!user) return;

    const ws = new WebSocket(wsUrl(user.accessToken));
    wsRef.current = ws;

    ws.onopen = () => {
      if (!mountedRef.current) return;
      setIsConnected(true);
      reconnectAttempt.current = 0;

      // Start keepalive pings
      keepaliveTimer.current = setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: 'ping' }));
        }
      }, KEEPALIVE_MS);
    };

    ws.onmessage = (event) => {
      if (!mountedRef.current) return;
      try {
        const msg: WSMessage = JSON.parse(event.data);

        if (msg.type === 'init') {
          const initial: Record<string, TaskState> = {};
          for (const t of msg.tasks) {
            initial[t.task_id] = t;
          }
          setTasks(initial);
        } else if (msg.type === 'task_update') {
          const { type: _, ...taskData } = msg;
          setTasks((prev) => ({
            ...prev,
            [msg.task_id]: { ...prev[msg.task_id], ...taskData } as TaskState,
          }));
        }
        // pong messages are ignored
      } catch {
        // Ignore malformed messages
      }
    };

    ws.onclose = () => {
      if (!mountedRef.current) return;
      setIsConnected(false);
      if (keepaliveTimer.current) {
        clearInterval(keepaliveTimer.current);
        keepaliveTimer.current = null;
      }

      // Exponential backoff reconnect
      const delay = Math.min(
        1000 * 2 ** reconnectAttempt.current,
        MAX_RECONNECT_DELAY_MS,
      );
      reconnectAttempt.current += 1;
      reconnectTimer.current = setTimeout(() => {
        if (mountedRef.current) connect();
      }, delay);
    };

    ws.onerror = () => {
      // onclose will fire after this, triggering reconnect
    };
  }, [user, clearTimers]);

  useEffect(() => {
    mountedRef.current = true;
    connect();

    return () => {
      mountedRef.current = false;
      clearTimers();
      if (wsRef.current) {
        wsRef.current.close();
        wsRef.current = null;
      }
    };
  }, [connect, clearTimers]);

  const subscribe = useCallback((taskId: string) => {
    const ws = wsRef.current;
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify({ type: 'subscribe', task_id: taskId }));
    }
  }, []);

  return { tasks, isConnected, subscribe };
}