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;
  filename?: string;
  missing_slots?: Array<{ slot?: string }>;
  // Spool-assignment read-back verification (#2582).
  slot?: string;
  verified?: boolean;
  kprofile_applied?: boolean;
  saw_tray?: boolean;
  // 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) {
        handleMessageRef.current(message);
        // Small delay between messages to prevent overwhelming the browser.
        // This setTimeout is the yield; a requestAnimationFrame around the
        // handler used to sit here too, which stalled the whole queue in a
        // hidden tab (see the note on the rAF removal below).
        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]);

  // Write every pending printer status into the query cache.
  //
  // Extracted so the hidden-tab path below can run it inline: both paths share
  // this one body, so the merge semantics cannot drift apart. Cancels any
  // scheduled coalescing timer, since everything it was going to write has
  // just been written and re-running it would re-apply stale data over newer.
  const flushPrinterStatus = useCallback(() => {
    if (printerStatusTimeoutRef.current) {
      clearTimeout(printerStatusTimeoutRef.current);
      printerStatusTimeoutRef.current = null;
    }

    const updates = new Map(pendingPrinterStatus.current);
    pendingPrinterStatus.current.clear();

    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;
      });
    });
  }, [queryClient]);

  // Printer status update — coalesced while the tab is visible, written
  // straight through while it is not.
  //
  // #2754 (reporter @mic4rd), in two stages. First, these writes ran inside a
  // requestAnimationFrame: a hidden tab gets no rendering opportunities, so
  // the browser *holds* queued frame callbacks rather than throttling them,
  // and nothing reached the cache until the tab was shown again. Removing the
  // frame callback fixed that total stall but not the report, because a second
  // timer-shaped dependency was left behind — this 100ms coalescing window.
  //
  // Browsers clamp timers in a hidden page to at best once a second, and drop
  // pages hidden for more than five minutes to roughly one wake-up a minute.
  // The reporter saw a tab title stuck at 2% beside a page at 40%.
  //
  // The coalescing exists to stop rapid messages triggering a render cascade.
  // A hidden tab is not painting, so there is no cascade to prevent there —
  // the timer is pure cost, and it is exactly the thing being throttled. So
  // when hidden, skip it and write immediately.
  //
  // Note "hidden", not "unfocused": on Windows a fully-occluded window reports
  // visibilityState 'hidden' too, which is why the reporter saw this from
  // merely clicking away rather than only from switching tabs.
  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 });

    if (document.hidden) {
      flushPrinterStatus();
      return;
    }

    // Schedule update if not already scheduled
    if (!printerStatusTimeoutRef.current) {
      printerStatusTimeoutRef.current = window.setTimeout(flushPrinterStatus, 100);
    }
  }, [flushPrinterStatus]);

  // 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.
      // The 500ms stagger is the anti-cascade measure; a frame callback around
      // each invalidation used to sit inside it and stalled these refreshes in
      // a hidden tab for the same reason as the status writes above (#2754).
      let delay = 0;
      keys.forEach((key) => {
        setTimeout(() => {
          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 'kill_switch_triggered': {
        const printer = message.printer_name || `Printer ${message.printer_id ?? '?'}`;
        const filename = message.filename || t('common.unknown');
        showToast(t('printers.toast.killSwitchTriggered', { printer, filename }), 'error');
        break;
      }

      case 'billing_charge_failed': {
        const printer = message.printer_name || `Printer ${message.printer_id ?? '?'}`;
        const filename = message.filename || t('common.unknown');
        showToast(t('printers.toast.billingChargeFailed', { printer, filename }), 'error');
        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_assignment_verified': {
        // #2582: the backend read the AMS telemetry back after an assignment
        // and either confirmed the tray accepted it or timed out. Toast the
        // outcome so the AMS→Studio hand-off is no longer silent.
        // Backend always supplies printer_name (falls back to "Printer <id>"),
        // so the '||' here only guards a malformed payload.
        const printer = message.printer_name || 'Printer';
        const slot = message.slot || '?';
        if (message.verified) {
          if (message.kprofile_applied === false) {
            // Filament id landed but the K-profile (cali_idx) did not — the
            // exact "loaded but no flow profile" case the reporter chased.
            showToast(
              t('printers.toast.assignmentVerifiedNoKprofile', { slot, printer }),
              'warning'
            );
          } else {
            showToast(t('printers.toast.assignmentVerified', { slot, printer }), 'success');
          }
        } else {
          showToast(t('printers.toast.assignmentNotConfirmed', { slot, printer }), 'warning');
        }
        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 };
}