frontend/src/hooks/useWebSocket.ts

import { useCallback, useEffect, useRef, useState } from "react";
import { inferenceWebSocketUrl } from "@/lib/api";
import type { WsEvent } from "@/types/inference";

interface UseInferenceWebSocket {
  status: "idle" | "connecting" | "open" | "closed" | "error";
  events: WsEvent[];
  open: (sessionId?: string) => void;
  close: () => void;
  send: (data: object) => void;
  sessionId: string | null;
  reset: () => void;
}

export function useInferenceWebSocket(): UseInferenceWebSocket {
  const [status, setStatus] = useState<UseInferenceWebSocket["status"]>("idle");
  const [events, setEvents] = useState<WsEvent[]>([]);
  const [sessionId, setSessionId] = useState<string | null>(null);
  const wsRef = useRef<WebSocket | null>(null);

  const close = useCallback(() => {
    if (wsRef.current) {
      try {
        wsRef.current.close();
      } catch {
        // noop
      }
      wsRef.current = null;
    }
    setStatus("closed");
  }, []);

  const reset = useCallback(() => {
    close();
    setEvents([]);
    setSessionId(null);
    setStatus("idle");
  }, [close]);

  const open = useCallback((override?: string) => {
    const id =
      override ??
      (typeof crypto !== "undefined" && "randomUUID" in crypto
        ? crypto.randomUUID()
        : Math.random().toString(36).slice(2, 14));
    setSessionId(id);
    setEvents([]);
    setStatus("connecting");
    const ws = new WebSocket(inferenceWebSocketUrl(id));
    wsRef.current = ws;
    ws.onopen = () => setStatus("open");
    ws.onmessage = (msg) => {
      try {
        const evt = JSON.parse(msg.data) as WsEvent;
        setEvents((prev) => [...prev, evt]);
      } catch {
        // ignore malformed
      }
    };
    ws.onerror = () => setStatus("error");
    ws.onclose = () => setStatus("closed");
  }, []);

  const send = useCallback((data: object) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(data));
    }
  }, []);

  useEffect(() => () => close(), [close]);

  return { status, events, open, close, send, sessionId, reset };
}