import { useEffect, useRef, useCallback } from "react";

/**
 * Derives the WebSocket base URL from VITE_APP_API_URL.
 * e.g. "https://api.brandxtech.ca" -> "wss://api.brandxtech.ca"
 * Falls back to current page host for local dev.
 */
function getWsBaseUrl(): string {
  const apiUrl = import.meta.env.VITE_APP_API_URL;
  if (apiUrl) {
    return apiUrl.replace(/^http/, "ws");
  }
  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
  return `${protocol}//${window.location.host}`;
}

interface UseWebSocketOptions<T> {
  /** The API path, e.g. "/v1/live/devices/123/components" */
  path: string;
  /** Transform the raw MessageEvent into your domain type */
  parse: (event: MessageEvent) => T;
  /** Called whenever a new parsed message arrives */
  onMessage: (data: T) => void;
  /** Auth token passed as ?token= query param */
  token?: string;
  /** Minimum seconds between updates, passed as ?rate= to the backend */
  rate?: number;
  /** Whether the connection should be active. Default true. */
  enabled?: boolean;
}

/**
 * useWebSocket manages a WebSocket connection with automatic reconnection.
 *
 * Usage:
 *   useWebSocket({
 *     path: `/v1/live/devices/${deviceID}/components`,
 *     parse: (e) => Component.any(JSON.parse(e.data)),
 *     onMessage: (component) => handleUpdate(component),
 *     token: authToken,
 *   });
 */
export function useWebSocket<T>(options: UseWebSocketOptions<T>) {
  const { path, parse, onMessage, token, rate = 0, enabled = true } = options;

  // Keep latest callbacks in refs so reconnects don't use stale closures
  const onMessageRef = useRef(onMessage);
  onMessageRef.current = onMessage;
  const parseRef = useRef(parse);
  parseRef.current = parse;

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout>>();
  const retriesRef = useRef(0);

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

    const params = new URLSearchParams();
    params.set("token", token);
    if (rate > 0) params.set("rate", rate.toString());

    const url = `${getWsBaseUrl()}${path}?${params.toString()}`;
    const ws = new WebSocket(url);
    wsRef.current = ws;

    ws.onopen = () => {
      console.debug(`[ws] connected: ${path}`);
      retriesRef.current = 0;
    };

    ws.onmessage = (event) => {
      try {
        const parsed = parseRef.current(event);
        onMessageRef.current(parsed);
      } catch (err) {
        console.warn(`[ws] parse error on ${path}:`, err);
      }
    };

    ws.onerror = (err) => {
      console.warn(`[ws] error on ${path}:`, err);
    };

    ws.onclose = (event) => {
      console.debug(`[ws] closed: ${path} (code: ${event.code})`);
      wsRef.current = null;

      // Don't reconnect on clean close or auth rejection
      if (event.code === 1000 || event.code === 4001 || event.code === 4003) {
        return;
      }

      // Exponential backoff: 1s, 2s, 4s, 8s, ... capped at 30s
      const delay = Math.min(1000 * Math.pow(2, retriesRef.current), 30000);
      retriesRef.current += 1;
      console.debug(`[ws] reconnecting in ${delay}ms...`);
      reconnectTimeoutRef.current = setTimeout(connect, delay);
    };
  }, [path, token, rate]);

  useEffect(() => {
    if (!enabled || !token) {
      if (wsRef.current) {
        wsRef.current.close(1000, "disabled");
        wsRef.current = null;
      }
      return;
    }

    connect();

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (wsRef.current) {
        wsRef.current.close(1000, "cleanup");
        wsRef.current = null;
      }
    };
  }, [connect, enabled, token]);
}
