// Attempt connection. If a network mismatch occurs, fall back to the next candidate.
walletAPI = await connector.connect(networkId);
// Lace v4: getConfiguration() is on walletAPI (not connector)
const walletRaw = walletAPI as unknown as Record<string, unknown>;

let address = "";
let coinPublicKey = "";
let encryptionPublicKey = "";

if (typeof walletRaw.getShieldedAddresses === "function") {
  // getShieldedAddresses() may return an array (old versions) or a single object (new versions), so handle both cases
  const result = await (
    walletRaw.getShieldedAddresses as () => Promise<Record<string, unknown>>
  )();
  // Lace v4 returns a single object (not array):
  // { shieldedAddress, shieldedCoinPublicKey, shieldedEncryptionPublicKey }
  const entry = (Array.isArray(result) ? result[0] : result) as
    | Record<string, unknown>
    | undefined;
  if (entry) {
    address = String(entry.shieldedAddress ?? entry.address ?? "");
    coinPublicKey = String(
      entry.shieldedCoinPublicKey ?? entry.coinPublicKey ?? "",
    );
    encryptionPublicKey = String(
      entry.shieldedEncryptionPublicKey ?? entry.encryptionPublicKey ?? "",
    );
  }
}

import { formatBalance } from "@/lib/utils";
import type { BalanceState } from "@/utils/types";
import type { DAppConnectorWalletAPI } from "@midnight-ntwrk/dapp-connector-api";
import { useCallback, useState } from "react";

export type { BalanceState };

/**
 * Fetch shielded / unshielded / dust balances in parallel from the wallet API.
 *
 * The Lace SDK type definitions do not include these balance methods,
 * so this function casts to unknown and checks method availability dynamically.
 * Promise.allSettled allows us to receive remaining results even if one call fails.
 */
async function fetchBalances(walletAPI: DAppConnectorWalletAPI): Promise<{
  shielded: string;
  unshielded: string;
  dust: string;
}> {
  const raw = walletAPI as unknown as Record<string, unknown>;

  // Try fetching balances in parallel. If a method is unavailable, return null.
  const [shieldedResult, unshieldedResult, dustResult] =
    await Promise.allSettled([
      typeof raw.getShieldedBalances === "function"
        ? (raw.getShieldedBalances as () => Promise<unknown>)()
        : Promise.resolve(null),
      typeof raw.getUnshieldedBalances === "function"
        ? (raw.getUnshieldedBalances as () => Promise<unknown>)()
        : Promise.resolve(null),
      typeof raw.getDustBalance === "function"
        ? (raw.getDustBalance as () => Promise<unknown>)()
        : Promise.resolve(null),
    ]);

  return {
    shielded:
      shieldedResult.status === "fulfilled"
        ? formatBalance(shieldedResult.value)
        : "--",
    unshielded:
      unshieldedResult.status === "fulfilled"
        ? formatBalance(unshieldedResult.value)
        : "--",
    dust:
      dustResult.status === "fulfilled"
        ? formatBalance(dustResult.value)
        : "--",
  };
}

/**
 * Custom hook to manage loading and refreshing balances.
 *
 * @param walletAPI - Connected wallet API. If null, no action is taken.
 * @returns balanceState: current balance state / refresh: manual refresh trigger
 */
export function useBalance(walletAPI: DAppConnectorWalletAPI | null) {
  const [balanceState, setBalanceState] = useState<BalanceState>({
    status: "idle",
  });

  /**
   * Fetch balances from the wallet API and update state.
   * If walletAPI is null, do nothing.
   * Set status to "loading" while fetching, "loaded" on success with balances,
   * and "error" if anything fails.
   */
  const refresh = useCallback(async () => {
    if (!walletAPI) return;
    setBalanceState({ status: "loading" });
    try {
      const balances = await fetchBalances(walletAPI);
      setBalanceState({ status: "loaded", ...balances });
    } catch (e: unknown) {
      console.error("[useBalance] Failed to fetch balances:", e);
      setBalanceState({ status: "error" });
    }
  }, [walletAPI]);

  return { balanceState, refresh };
}

git clone https://github.com/mashharuki/midnight-lace-react-sample-app.git

# Move to the frontend app directory
cd app

bun install

bun run build

bun run dev