bambuddy/frontend/src/hooks/useWebSocket.ts

import { useQueryClient } from '@tanstack/react-query';
import { useCallback, useEffect, useRef, useState } from 'react';
import { useToast } from '../contexts/ToastContext';
import { useTranslation } from 'react-i18next';
import { api, ApiError } from '../api/client';
import { inventoryLocationsQueryKey } from '../utils/inventoryQueries';

// The only auth-failure close code /api/v1/ws emits (websocket.py
// _WS_CLOSE_UNAUTHORIZED). A 4401 means the ws-token was missing / invalid /
// expired, or the caller lacks WEBSOCKET_CONNECT — none of which a reconnect can
// fix without a fresh login (which remounts this provider anyway). Treat it as
// terminal so we don't respawn the /auth/ws-token loop.
const WS_CLOSE_UNAUTHORIZED = 4401;

interface WebSocketMessage {
  type: string;
  printer_id?: number;
  data?: Record<string, unknown>;
  printer_name?: string;
  missing_slots?: Array<{ slot?: string }>;
  // Slicer Pipeline run events (#1425 PR C). ``run`` carries the full
  // PipelineRunResponse payload — typed loosely here so the WebSocket hook
  // doesn't pull the full client.ts types in.
  run?: { pipeline_id?: number | null };
}

export function useWebSocket() {
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<number | null>(null);
  // Set true by the effect cleanup so a close event fired *during* unmount
  // can't schedule a reconnect after the provider is gone (the old code cleared
  // reconnectTimeoutRef, then .close() ran ws.onclose which set a *fresh*
  // timeout — a leaked reconnect that kept minting ws-tokens post-logout).
  const disposedRef = useRef(false);
  const queryClient = useQueryClient();
  const [isConnected, setIsConnected] = useState(false);
  const lastMissingSpoolWarningRef = useRef<Map<number, string>>(new Map());
  const { showToast } = useToast();
  const { t } = useTranslation();

  // Debounce invalidations to prevent rapid re-render cascades
  const pendingInvalidations = useRef<Set<string>>(new Set());
  const invalidationTimeoutRef = useRef<number | null>(null);

  // Throttle printer status updates to prevent freeze during rapid messages
  const pendingPrinterStatus = useRef<Map<number, Record<string, unknown>>>(new Map());
  const printerStatusTimeoutRef = useRef<number | null>(null);

  // Throttle message processing to prevent browser freeze
  const messageQueueRef = useRef<WebSocketMessage[]>([]);
  const processingRef = useRef(false);

  // Use ref for handleMessage to avoid stale closure in connect
  const handleMessageRef = useRef<(message: WebSocketMessage) => void>(() => {});

  // Process message queue with throttling to prevent UI freeze
  const processMessageQueue = useCallback(() => {
    if (processingRef.current || messageQueueRef.current.length === 0) {
      return;
    }

    processingRef.current = true;

    const processNext = () => {
      const message = messageQueueRef.current.shift();
      if (message) {
        // Use requestAnimationFrame to yield to the browser
        requestAnimationFrame(() => {
          handleMessageRef.current(message);
          // Small delay between messages to prevent overwhelming the browser
          if (messageQueueRef.current.length > 0) {
            setTimeout(processNext, 16); // ~60fps
          } else {
            processingRef.current = false;
          }
        });
      } else {
        processingRef.current = false;
      }
    };

    processNext();
  }, []);

  const connect = useCallback(async () => {
    if (disposedRef.current || wsRef.current?.readyState === WebSocket.OPEN) {
      return;
    }

    // GHSA-r2qv follow-up: when auth is enabled, /ws now requires a token
    // minted by POST /api/v1/auth/ws-token. We use the shared ``api.request``
    // helper (via ``api.getWebSocketToken``) so the JWT Authorization header
    // is attached — a raw ``fetch()`` with ``credentials: 'include'`` would
    // miss it (Bambuddy uses Bearer tokens, not cookies, for JWT auth).
    // Auth-disabled deployments accept connections without a token.
    let token: string | undefined;
    try {
      const resp = await api.getWebSocketToken();
      token = resp.token;
    } catch (err) {
      // A 401/403 from the token mint is an AUTH decision, not a transient
      // blip, so retrying is pointless and hammers /auth/ws-token every 3s:
      //   401 — the JWT expired. ``request()`` already cleared it and
      //         dispatched ``auth:expired``, so the route guard is redirecting
      //         to /login and this provider is about to unmount.
      //   403 — the user is validly logged in but their group lacks
      //         WEBSOCKET_CONNECT. They stay logged in; live updates simply
      //         degrade to the REST polling the query cache already does.
      // Either way: do NOT open a tokenless socket (the server just closes it
      // 4401) and do NOT reconnect. The old catch-all fell through to a
      // tokenless socket whose 4401 close rescheduled connect() forever. A
      // network/5xx error is not auth — fall through and let the socket + its
      // reconnect loop handle it (auth-disabled deployments also land here with
      // no token and connect fine).
      const status = err instanceof ApiError ? err.status : 0;
      if (status === 401 || status === 403) {
        return;
      }
    }

    if (disposedRef.current) {
      return;
    }

    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const tokenParam = token ? `?token=${encodeURIComponent(token)}` : '';
    const wsUrl = `${protocol}//${window.location.host}/api/v1/ws${tokenParam}`;

    const ws = new WebSocket(wsUrl);

    let pingInterval: number | null = null;

    ws.onopen = () => {
      if (import.meta.env.MODE !== 'test') console.log('[WebSocket] Connected');
      setIsConnected(true);
      // Start ping interval
      pingInterval = window.setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: 'ping' }));
        }
      }, 30000);
    };

    ws.onmessage = (event) => {
      try {
        const message: WebSocketMessage = JSON.parse(event.data);
        // Handle printer_status directly (already throttled) to avoid queue delays
        // This prevents the "timelapse" effect where status updates are applied slowly
        if (message.type === 'printer_status' && message.printer_id !== undefined && message.data) {
          handleMessageRef.current(message);
        } else {
          // Queue other messages for throttled processing
          messageQueueRef.current.push(message);
          processMessageQueue();
        }
      } catch {
        // Ignore parse errors
      }
    };

    ws.onclose = (event) => {
      if (import.meta.env.MODE !== 'test') console.log('[WebSocket] Closed', event.code, event.reason);
      if (pingInterval) {
        clearInterval(pingInterval);
        pingInterval = null;
      }
      setIsConnected(false);
      wsRef.current = null;

      // Don't reconnect after an auth rejection (4401) or once the provider has
      // unmounted — both would just respawn the /auth/ws-token loop. A 4401 is
      // terminal (needs a fresh login, which remounts us); every other close
      // code is treated as a network drop and gets the 3s reconnect.
      if (disposedRef.current || event.code === WS_CLOSE_UNAUTHORIZED) {
        return;
      }

      // Reconnect after 3 seconds
      reconnectTimeoutRef.current = window.setTimeout(() => {
        connect();
      }, 3000);
    };

    ws.onerror = (error) => {
      if (import.meta.env.MODE !== 'test') console.error('[WebSocket] Error', error);
      ws.close();
    };

    wsRef.current = ws;
  }, [processMessageQueue]);

  // Throttled printer status update - coalesces rapid updates per printer
  const throttledPrinterStatusUpdate = useCallback((printerId: number, data: Record<string, unknown>) => {
    // Merge with any pending data for this printer
    const existing = pendingPrinterStatus.current.get(printerId) || {};
    pendingPrinterStatus.current.set(printerId, { ...existing, ...data });

    // Schedule update if not already scheduled
    if (!printerStatusTimeoutRef.current) {
      printerStatusTimeoutRef.current = window.setTimeout(() => {
        const updates = new Map(pendingPrinterStatus.current);
        pendingPrinterStatus.current.clear();
        printerStatusTimeoutRef.current = null;

        // Apply all pending updates
        requestAnimationFrame(() => {
          updates.forEach((statusData, id) => {
            queryClient.setQueryData(
              ['printerStatus', id],
              (old: Record<string, unknown> | undefined) => {
                const merged = { ...old, ...statusData };
                if (merged.wifi_signal == null && old?.wifi_signal != null) {
                  merged.wifi_signal = old.wifi_signal;
                }
                return merged;
              }
            );
          });
        });
      }, 100); // Update at most every 100ms
    }
  }, [queryClient]);

  // Debounced invalidation helper - coalesces multiple rapid invalidations
  const debouncedInvalidate = useCallback((queryKey: string) => {
    pendingInvalidations.current.add(queryKey);

    // Clear existing timeout
    if (invalidationTimeoutRef.current) {
      clearTimeout(invalidationTimeoutRef.current);
    }

    // Schedule invalidation after a delay (3s to prevent browser freeze on print completion)
    invalidationTimeoutRef.current = window.setTimeout(() => {
      const keys = Array.from(pendingInvalidations.current);
      pendingInvalidations.current.clear();
      invalidationTimeoutRef.current = null;

      // Invalidate queries one at a time with delays to prevent freeze
      let delay = 0;
      keys.forEach((key) => {
        setTimeout(() => {
          requestAnimationFrame(() => {
            queryClient.invalidateQueries({ queryKey: [key] });
          });
        }, delay);
        delay += 500; // 500ms between each invalidation
      });
    }, 3000);
  }, [queryClient]);

  const handleMessage = useCallback((message: WebSocketMessage) => {
    switch (message.type) {
      case 'printer_status':
        if (message.printer_id !== undefined && message.data) {
          throttledPrinterStatusUpdate(message.printer_id, message.data);
        }
        break;

      case 'print_start':
        // Refetch printer status immediately when print starts to get printable_objects_count
        if (message.printer_id !== undefined) {
          queryClient.invalidateQueries({ queryKey: ['printerStatus', message.printer_id] });
        }
        break;

      case 'missing_spool_assignment': {
        if (message.printer_id === undefined || !Array.isArray(message.missing_slots)) {
          break;
        }

        const missingSlotLabels = message.missing_slots
          .map((slot) => (slot && typeof slot.slot === 'string' ? slot.slot : 'Unknown'))
          .filter((slot) => slot.length > 0);

        if (missingSlotLabels.length === 0) {
          lastMissingSpoolWarningRef.current.delete(message.printer_id);
          break;
        }

        const signature = missingSlotLabels.join('|');
        if (lastMissingSpoolWarningRef.current.get(message.printer_id) === signature) {
          break;
        }
        lastMissingSpoolWarningRef.current.set(message.printer_id, signature);

        const printerName = message.printer_name || `Printer ${message.printer_id}`;
        const toastMsg = t('printers.toast.missingSpoolAssignment', {
          printer: printerName,
          slots: missingSlotLabels.join(', '),
        });
        showToast(toastMsg, 'warning');
        break;
      }

      case 'print_complete':
        // Don't invalidate printerStatus here - it causes re-render cascade and browser freeze
        // The printer_status websocket messages will naturally update the status
        debouncedInvalidate('archives');
        debouncedInvalidate('archiveStats');
        break;

      case 'archive_created':
        debouncedInvalidate('archives');
        debouncedInvalidate('archiveStats');
        break;

      case 'archive_updated':
        debouncedInvalidate('archives');
        break;

      case 'pong':
        // Keepalive response, ignore
        break;

      case 'plate_not_empty':
        // Plate detection found objects - print was paused
        // Dispatch event for toast notification
        window.dispatchEvent(new CustomEvent('plate-not-empty', {
          detail: {
            printer_id: message.printer_id,
            printer_name: (message as unknown as { printer_name?: string }).printer_name,
            message: (message as unknown as { message?: string }).message,
          }
        }));
        break;

      case 'inventory_changed':
        // Spool created/updated/deleted/archived/restored - refresh inventory across all tabs
        debouncedInvalidate('inventory-spools');
        debouncedInvalidate('spoolman-inventory-spools');
        debouncedInvalidate(inventoryLocationsQueryKey[0]);
        break;

      case 'spool_assignment_changed':
        // Spool assigned/unassigned - refresh assignment data across all tabs
        debouncedInvalidate('spool-assignments');
        debouncedInvalidate('slotPresets');
        break;

      case 'spool_auto_assigned':
        // RFID tag matched - refresh inventory and assignment data
        debouncedInvalidate('inventory-spools');
        debouncedInvalidate('spool-assignments');
        break;

      case 'spool_usage_logged':
        // Filament consumption recorded - refresh spool data
        debouncedInvalidate('inventory-spools');
        break;

      case 'unknown_tag': {
        // Unknown RFID tag detected — dispatch event for UI. The backend
        // ships the slot's current tray data alongside the event so
        // consumers don't have to look it up from the (frequently stale)
        // cached printerStatus query.
        const m = message as unknown as {
          printer_id?: number;
          ams_id?: number;
          tray_id?: number;
          tag_uid?: string;
          tray_uuid?: string;
          tray_type?: string | null;
          tray_color?: string | null;
          tray_sub_brands?: string | null;
          tray_count?: number | null;
        };
        window.dispatchEvent(new CustomEvent('unknown-tag', {
          detail: {
            printer_id: m.printer_id,
            ams_id: m.ams_id,
            tray_id: m.tray_id,
            tag_uid: m.tag_uid,
            tray_uuid: m.tray_uuid,
            tray_type: m.tray_type,
            tray_color: m.tray_color,
            tray_sub_brands: m.tray_sub_brands,
            tray_count: m.tray_count,
          }
        }));
        break;
      }

      case 'spoolbuddy_weight':
        window.dispatchEvent(new CustomEvent('spoolbuddy-weight', { detail: message }));
        break;

      case 'spoolbuddy_tag_matched':
        window.dispatchEvent(new CustomEvent('spoolbuddy-tag-matched', { detail: message }));
        debouncedInvalidate('inventory-spools');
        break;

      case 'spoolbuddy_unknown_tag':
        window.dispatchEvent(new CustomEvent('spoolbuddy-unknown-tag', { detail: message }));
        break;

      case 'spoolbuddy_tag_removed':
        window.dispatchEvent(new CustomEvent('spoolbuddy-tag-removed', { detail: message }));
        break;

      case 'spoolbuddy_tag_written':
        window.dispatchEvent(new CustomEvent('spoolbuddy-tag-written', { detail: message }));
        debouncedInvalidate('inventory-spools');
        break;

      case 'spoolbuddy_tag_write_failed':
        window.dispatchEvent(new CustomEvent('spoolbuddy-tag-write-failed', { detail: message }));
        break;

      case 'spoolbuddy_online':
        window.dispatchEvent(new CustomEvent('spoolbuddy-online', { detail: message }));
        debouncedInvalidate('spoolbuddy-devices');
        debouncedInvalidate('spoolbuddy-update-check');
        break;

      case 'spoolbuddy_offline':
        window.dispatchEvent(new CustomEvent('spoolbuddy-offline', { detail: message }));
        debouncedInvalidate('spoolbuddy-devices');
        break;

      case 'spoolbuddy_update':
        debouncedInvalidate('spoolbuddy-devices');
        debouncedInvalidate('spoolbuddy-update-check');
        break;

      // Dispatch toast lifecycle (#1625 follow-up — restored the upload
      // progress UI that the scheduler unification removed). Four backend
      // event types collapse to one frontend channel. No
      // `queue_item_queued` (the toast must wait for the upload to
      // actually start) and no `queue_item_dispatched` (the legacy
      // background-dispatch flow kept status='processing' from upload
      // start until printer ack — the "Awaiting printer…" subtitle is
      // derived from upload_progress_pct >= 99.9, not from a separate
      // event).
      case 'queue_item_uploading':
      case 'queue_item_upload_progress':
      case 'queue_item_acked':
      case 'queue_item_failed':
        window.dispatchEvent(new CustomEvent('bambuddy:dispatch-toast', { detail: message }));
        break;
      // Slicer Pipeline runs (#1425 PR C). State transitions on the run
      // refresh both the dashboard list AND the per-pipeline "Last run"
      // chip in Settings → Pipelines.
      case 'pipeline_run_updated':
        queryClient.invalidateQueries({ queryKey: ['pipeline-runs-all'] });
        if (message.run?.pipeline_id) {
          queryClient.invalidateQueries({ queryKey: ['pipeline-runs', message.run.pipeline_id] });
        }
        break;
    }
  }, [queryClient, debouncedInvalidate, throttledPrinterStatusUpdate, showToast, t]);

  // Keep the ref updated with latest handleMessage
  useEffect(() => {
    handleMessageRef.current = handleMessage;
  }, [handleMessage]);

  useEffect(() => {
    // connect() is async after the GHSA-r2qv fix (mints a ws-token first).
    // Fire-and-forget at mount; the inner reconnect loop also calls
    // connect() in the ws.onclose handler.
    disposedRef.current = false;
    void connect();

    return () => {
      // Mark disposed BEFORE closing so the ws.onclose triggered by close()
      // sees it and won't schedule a post-unmount reconnect.
      disposedRef.current = true;
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (invalidationTimeoutRef.current) {
        clearTimeout(invalidationTimeoutRef.current);
      }
      if (printerStatusTimeoutRef.current) {
        clearTimeout(printerStatusTimeoutRef.current);
      }
      if (wsRef.current) {
        wsRef.current.close();
      }
    };
  }, [connect]);

  const sendMessage = useCallback((message: Record<string, unknown>) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    }
  }, []);

  return { isConnected, sendMessage };
}