import { useEffect, useRef } from 'react';

function buildWsUrl() {
  if (import.meta.env.VITE_WS_URL) {
    return import.meta.env.VITE_WS_URL;
  }

  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  return `${protocol}//${window.location.host}/ws`;
}

export function useWebSocket({ onMessage }) {
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;

  useEffect(() => {
    let socket;
    let reconnectTimer;
    let progressFlushTimer = null;
    let progressFlushRaf = null;
    let latestProgressMessage = null;
    let isUnmounted = false;

    const clearProgressFlush = () => {
      if (progressFlushRaf !== null && typeof window !== 'undefined' && typeof window.cancelAnimationFrame === 'function') {
        window.cancelAnimationFrame(progressFlushRaf);
      }
      if (progressFlushTimer) {
        clearTimeout(progressFlushTimer);
      }
      progressFlushRaf = null;
      progressFlushTimer = null;
    };

    const flushProgressMessage = () => {
      progressFlushRaf = null;
      progressFlushTimer = null;
      const message = latestProgressMessage;
      latestProgressMessage = null;
      if (!message) {
        return;
      }
      onMessageRef.current?.(message);
    };

    const scheduleProgressFlush = () => {
      if (progressFlushRaf !== null || progressFlushTimer) {
        return;
      }
      if (typeof window !== 'undefined' && typeof window.requestAnimationFrame === 'function') {
        progressFlushRaf = window.requestAnimationFrame(flushProgressMessage);
        return;
      }
      progressFlushTimer = setTimeout(flushProgressMessage, 16);
    };

    const connect = () => {
      if (isUnmounted) {
        return;
      }

      socket = new WebSocket(buildWsUrl());

      socket.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data);
          if (message?.type === 'PIPELINE_PROGRESS') {
            latestProgressMessage = message;
            scheduleProgressFlush();
            return;
          }
          onMessageRef.current?.(message);
        } catch (error) {
          // ignore invalid json
        }
      };

      socket.onclose = () => {
        if (!isUnmounted) {
          reconnectTimer = setTimeout(connect, 1500);
        }
      };

      socket.onerror = () => {
        if (socket && socket.readyState !== WebSocket.CLOSED) {
          socket.close();
        }
      };
    };

    connect();

    return () => {
      isUnmounted = true;
      clearProgressFlush();
      if (reconnectTimer) {
        clearTimeout(reconnectTimer);
      }
      if (socket && socket.readyState !== WebSocket.CLOSED) {
        socket.close();
      }
    };
  }, []);
}