dchain/client-app/hooks/useBalance.ts

/**
 * Balance hook — uses the WebSocket gateway to receive instant updates when
 * a tx involving the current address is committed, with HTTP polling as a
 * graceful fallback for old nodes that don't expose /api/ws.
 *
 * Flow:
 *   1. On mount: immediate HTTP fetch so the UI has a non-zero balance ASAP
 *   2. Subscribe to `addr:<my_pubkey>` on the WS hub
 *   3. On every `tx` event, re-fetch balance (cheap — one Badger read server-side)
 *   4. If WS disconnects for >15s, fall back to 10-second polling until it reconnects
 */

import { useEffect, useCallback, useRef } from 'react';
import { getBalance } from '@/lib/api';
import { getWSClient } from '@/lib/ws';
import { useStore } from '@/lib/store';

const FALLBACK_POLL_INTERVAL  = 10_000;   // HTTP poll when WS is down
const WS_GRACE_BEFORE_POLLING = 15_000;   // don't start polling immediately on disconnect

export function useBalance() {
  const keyFile    = useStore(s => s.keyFile);
  const setBalance = useStore(s => s.setBalance);

  const refresh = useCallback(async () => {
    if (!keyFile) return;
    try {
      const bal = await getBalance(keyFile.pub_key);
      setBalance(bal);
    } catch {
      // transient — next call will retry
    }
  }, [keyFile, setBalance]);

  // --- fallback polling management ---
  const pollTimerRef    = useRef<ReturnType<typeof setInterval> | null>(null);
  const disconnectSinceRef = useRef<number | null>(null);
  const disconnectTORef = useRef<ReturnType<typeof setTimeout> | null>(null);

  const startPolling = useCallback(() => {
    if (pollTimerRef.current) return;
    console.log('[useBalance] WS down for grace period — starting HTTP poll');
    refresh();
    pollTimerRef.current = setInterval(refresh, FALLBACK_POLL_INTERVAL);
  }, [refresh]);

  const stopPolling = useCallback(() => {
    if (pollTimerRef.current) {
      clearInterval(pollTimerRef.current);
      pollTimerRef.current = null;
    }
    if (disconnectTORef.current) {
      clearTimeout(disconnectTORef.current);
      disconnectTORef.current = null;
    }
    disconnectSinceRef.current = null;
  }, []);

  useEffect(() => {
    if (!keyFile) return;
    const ws  = getWSClient();

    // Immediate HTTP fetch so the UI is not empty while the WS hello arrives.
    refresh();

    // Refresh balance whenever a tx for our address is committed.
    const offTx = ws.subscribe('addr:' + keyFile.pub_key, (frame) => {
      if (frame.event === 'tx') {
        refresh();
      }
    });

    // Manage fallback polling based on WS connection state.
    const offConn = ws.onConnectionChange((ok) => {
      if (ok) {
        stopPolling();
        refresh();  // catch up anything we missed while disconnected
      } else if (disconnectTORef.current === null) {
        disconnectSinceRef.current = Date.now();
        disconnectTORef.current    = setTimeout(startPolling, WS_GRACE_BEFORE_POLLING);
      }
    });

    ws.connect();

    return () => {
      offTx();
      offConn();
      stopPolling();
    };
  }, [keyFile, refresh, startPolling, stopPolling]);

  return { refresh };
}