visionA/local-tool/frontend/src/hooks/use-flash-progress.ts

'use client';

import { useEffect, useRef, useCallback } from 'react';
import { createWebSocket } from '@/lib/ws';
import { useFlashStore } from '@/stores/flash-store';
import type { FlashProgress } from '@/types/device';

/**
 * Manages flash progress WebSocket.
 * Returns a `connectAndWait` callback that creates the WebSocket and
 * returns a promise that resolves once the WS is open.
 */
export function useFlashProgress(deviceId: string) {
  const updateProgress = useFlashStore((s) => s.updateProgress);
  const wsRef = useRef<ReturnType<typeof createWebSocket> | null>(null);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      wsRef.current?.close();
      wsRef.current = null;
    };
  }, [deviceId]);

  /**
   * Creates the WebSocket connection and returns a promise that resolves
   * once the connection is open. This is called imperatively (not via
   * useEffect) to avoid React render-cycle timing issues.
   */
  const connectAndWait = useCallback(
    () =>
      new Promise<void>((resolve) => {
        // Close any existing connection
        wsRef.current?.close();

        let resolved = false;
        const doResolve = () => {
          if (!resolved) {
            resolved = true;
            resolve();
          }
        };

        const ws = createWebSocket(
          `/ws/devices/${deviceId}/flash-progress`,
          (data) => {
            updateProgress(data as FlashProgress);
          },
          () => {
            doResolve();
          },
        );
        wsRef.current = ws;

        // Safety timeout — don't block forever
        setTimeout(doResolve, 3000);
      }),
    [deviceId, updateProgress],
  );

  /** Close the WebSocket connection. */
  const disconnect = useCallback(() => {
    wsRef.current?.close();
    wsRef.current = null;
  }, []);

  return { connectAndWait, disconnect };
}