frontend/src/lib/hooks/useWebSocketChat.ts

/**
 * WebSocket chat hook — manages connection lifecycle, reconnection,
 * heartbeat, and streaming message assembly.
 */
"use client";

import { useCallback, useEffect, useRef, useState } from "react";
import { createChatWebSocket } from "@/lib/api";

export type WsStatus = "connecting" | "connected" | "disconnected" | "error";

interface UseWebSocketChatOptions {
  userId?: string;
  onChunk: (chunk: string) => void;
  onStart: () => void;
  onEnd: () => void;
  onError: (msg: string) => void;
}

export function useWebSocketChat({
  userId,
  onChunk,
  onStart,
  onEnd,
  onError,
}: UseWebSocketChatOptions) {
  const wsRef = useRef<WebSocket | null>(null);
  const [status, setStatus] = useState<WsStatus>("disconnected");
  const reconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const heartbeatTimer = useRef<ReturnType<typeof setInterval> | null>(null);
  const reconnectAttempts = useRef(0);
  const MAX_RECONNECT = 5;
  const isMounted = useRef(true);

  const clearTimers = () => {
    if (reconnectTimer.current) clearTimeout(reconnectTimer.current);
    if (heartbeatTimer.current) clearInterval(heartbeatTimer.current);
  };

  const startHeartbeat = (ws: WebSocket) => {
    heartbeatTimer.current = setInterval(() => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ type: "ping" }));
      }
    }, 25000); // ping every 25s
  };

  const connect = useCallback(() => {
    if (!isMounted.current) return;
    if (wsRef.current?.readyState === WebSocket.OPEN) return;

    setStatus("connecting");
    const ws = createChatWebSocket(userId);
    wsRef.current = ws;

    ws.onopen = () => {
      if (!isMounted.current) return;
      setStatus("connected");
      reconnectAttempts.current = 0;
      startHeartbeat(ws);
    };

    ws.onmessage = (event) => {
      if (!isMounted.current) return;
      try {
        const msg = JSON.parse(event.data);
        switch (msg.type) {
          case "chat_start":
            onStart();
            break;
          case "chat_chunk":
            if (msg.content) onChunk(msg.content);
            break;
          case "chat_end":
            onEnd();
            break;
          case "pong":
            break; // heartbeat ack
          case "error":
            onError(msg.message || "Unknown error");
            break;
        }
      } catch {
        // ignore malformed messages
      }
    };

    ws.onerror = () => {
      if (!isMounted.current) return;
      setStatus("error");
    };

    ws.onclose = (event) => {
      if (!isMounted.current) return;
      clearTimers();
      setStatus("disconnected");

      // Reconnect with exponential backoff (unless intentional close)
      if (event.code !== 1000 && reconnectAttempts.current < MAX_RECONNECT) {
        const delay = Math.min(1000 * 2 ** reconnectAttempts.current, 30000);
        reconnectAttempts.current++;
        reconnectTimer.current = setTimeout(connect, delay);
      }
    };
  }, [userId, onChunk, onStart, onEnd, onError]);

  const disconnect = useCallback(() => {
    isMounted.current = false;
    clearTimers();
    if (wsRef.current) {
      wsRef.current.close(1000, "Component unmounted");
      wsRef.current = null;
    }
  }, []);

  const sendMessage = useCallback((message: string) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: "chat", message }));
      return true;
    }
    return false;
  }, []);

  useEffect(() => {
    isMounted.current = true;
    connect();
    return () => {
      disconnect();
    };
  }, [connect, disconnect]);

  return { status, sendMessage, reconnect: connect };
}