import { useEffect, useRef, useState, useCallback } from 'react';

const WS_BASE = (() => {
  if (import.meta.env.VITE_WS_URL) return import.meta.env.VITE_WS_URL;
  const loc = window.location;
  const proto = loc.protocol === 'https:' ? 'wss:' : 'ws:';
  // Use the same origin as the page; in development the Vite dev server
  // proxies /ws to the backend, and in production nginx does the same.
  return `${proto}//${loc.host}/ws`;
})();

const RECONNECT_DELAYS = [1000, 2000, 4000, 8000, 15000, 30000]; // ms

/**
 * Maintains a WebSocket connection with automatic reconnection.
 *
 * @param {(message: object) => void} onMessage - Called for each parsed JSON message.
 * @returns {{ status: 'connecting'|'connected'|'disconnected'|'error' }}
 */
export default function useWebSocket(onMessage) {
  const [status, setStatus] = useState('connecting');
  const wsRef = useRef(null);
  const retryCountRef = useRef(0);
  const unmountedRef = useRef(false);
  const retryTimerRef = useRef(null);
  const onMessageRef = useRef(onMessage);

  // Keep onMessage ref up-to-date without re-running the effect
  useEffect(() => {
    onMessageRef.current = onMessage;
  }, [onMessage]);

  const connect = useCallback(() => {
    if (unmountedRef.current) return;

    setStatus('connecting');
    const ws = new WebSocket(WS_BASE);
    wsRef.current = ws;

    ws.onopen = () => {
      if (unmountedRef.current) { ws.close(); return; }
      retryCountRef.current = 0;
      setStatus('connected');
    };

    ws.onmessage = (event) => {
      try {
        const msg = JSON.parse(event.data);
        onMessageRef.current(msg);
      } catch {
        // ignore malformed messages
      }
    };

    ws.onclose = () => {
      if (unmountedRef.current) return;
      const delay = RECONNECT_DELAYS[Math.min(retryCountRef.current, RECONNECT_DELAYS.length - 1)];
      retryCountRef.current += 1;
      // After exhausting all backoff levels, keep status as 'error'
      if (retryCountRef.current > RECONNECT_DELAYS.length) {
        setStatus('error');
      } else {
        setStatus('disconnected');
      }
      retryTimerRef.current = setTimeout(connect, delay);
    };

    ws.onerror = () => {
      // onclose will fire right after, which handles reconnect
    };
  }, []);

  useEffect(() => {
    connect();
    return () => {
      unmountedRef.current = true;
      clearTimeout(retryTimerRef.current);
      wsRef.current?.close();
    };
  }, [connect]);

  return { status };
}