frontend/src/hooks/useDeviceStatusStreams.ts

import { pond } from "protobuf-ts/pond";
import { useEffect, useRef, useCallback, useMemo } from "react";

/**
 * Derives the WebSocket base URL from VITE_APP_API_URL.
 * Matches the logic in useWebSocket.ts.
 */
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}`;
}

export interface UseDeviceStatusStreamsOptions {
  /** Device IDs to stream status for (e.g. from the visible devices list) */
  deviceIds: string[];
  /** Called when a device status update is received */
  onStatusUpdate: (deviceId: string, status: pond.DeviceStatus) => void;
  /** Auth token passed as ?token= query param */
  token?: string;
  /** Whether the connections should be active. Default true. */
  enabled?: boolean;
}

/**
 * Subscribes to live device status streams for multiple devices.
 * Uses the /v1/live/devices/:device/status endpoint per device.
 * Readings (plenum, sen5x, co, co2, no2, o2, etc.) stream in real time.
 */
export function useDeviceStatusStreams(options: UseDeviceStatusStreamsOptions) {
  const { deviceIds, onStatusUpdate, token, enabled = true } = options;

  const onStatusUpdateRef = useRef(onStatusUpdate);
  onStatusUpdateRef.current = onStatusUpdate;

  const wsMapRef = useRef<Map<string, WebSocket>>(new Map());
  const retriesRef = useRef<Map<string, number>>(new Map());
  const reconnectTimeoutsRef = useRef<Map<string, ReturnType<typeof setTimeout>>>(new Map());

  const connect = useCallback(
    (deviceId: string) => {
      if (!token || !deviceId) return;

      const params = new URLSearchParams();
      params.set("token", token);

      const path = `/live/devices/${deviceId}/status`;
      const url = `${getWsBaseUrl()}${path}?${params.toString()}`;
      const ws = new WebSocket(url);
      wsMapRef.current.set(deviceId, ws);

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

      ws.onmessage = (event) => {
        try {
          const raw = JSON.parse(event.data);
          const status = pond.DeviceStatus.fromObject(raw ?? {});
          onStatusUpdateRef.current(deviceId, status);
        } catch (err) {
          console.warn(`[ws] failed to parse device status for ${deviceId}:`, err);
        }
      };

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

      ws.onclose = (event) => {
        wsMapRef.current.delete(deviceId);

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

        const retries = retriesRef.current.get(deviceId) ?? 0;
        const delay = Math.min(1000 * Math.pow(2, retries), 30000);
        retriesRef.current.set(deviceId, retries + 1);
        console.debug(`[ws] reconnecting ${path} in ${delay}ms...`);

        const timeout = setTimeout(() => connect(deviceId), delay);
        reconnectTimeoutsRef.current.set(deviceId, timeout);
      };
    },
    [token]
  );

  // Normalize for stable comparison: same set of IDs = same string regardless of order
  const deviceIdsKey = useMemo(
    () => [...new Set(deviceIds)].sort().join(","),
    [deviceIds.join(",")]
  );

  useEffect(() => {
    if (!enabled || !token || deviceIds.length === 0) {
      wsMapRef.current.forEach((ws) => ws.close(1000, "disabled"));
      wsMapRef.current.clear();
      reconnectTimeoutsRef.current.forEach((t) => clearTimeout(t));
      reconnectTimeoutsRef.current.clear();
      return;
    }

    const currentIds = new Set(deviceIds);

    // Connect to new devices
    deviceIds.forEach((id) => {
      if (!wsMapRef.current.has(id)) {
        connect(id);
      }
    });

    // Disconnect from devices no longer in the list
    wsMapRef.current.forEach((ws, id) => {
      if (!currentIds.has(id)) {
        ws.close(1000, "device removed");
        wsMapRef.current.delete(id);
        const t = reconnectTimeoutsRef.current.get(id);
        if (t) {
          clearTimeout(t);
          reconnectTimeoutsRef.current.delete(id);
        }
      }
    });

    return () => {
      wsMapRef.current.forEach((ws) => ws.close(1000, "cleanup"));
      wsMapRef.current.clear();
      reconnectTimeoutsRef.current.forEach((t) => clearTimeout(t));
      reconnectTimeoutsRef.current.clear();
    };
  }, [enabled, token, deviceIdsKey, deviceIds.length, connect]);
}