@shogun-sdk/one-shot/src/hooks/useBalance.ts

Shogun SDK - One Shot: React Components and hooks for cross-chain swaps

import {
  compareAddresses,
  getRpcUrls,
  getSolanaProvider,
  isEVMChain,
  isNativeToken,
  isSolanaChain,
  SOL_NATIVE,
} from '@shogun-sdk/money-legos';
import { TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID } from '@solana/spl-token';
import { PublicKey } from '@solana/web3.js';
import { useQuery } from '@tanstack/react-query';
import { ethers } from 'ethers';
import { useCallback } from 'react';

// ERC20 ABI for token balance queries
const ERC20_ABI = ['function balanceOf(address owner) view returns (uint256)'];

// Constants
const REFETCH_INTERVAL = 5000; // 5 seconds

// Types
type BalanceResult = {
  balance: bigint | undefined;
  loading: boolean;
  error: Error | null;
};

type TokenAccount = {
  account: {
    data: {
      parsed: {
        info: {
          mint: string;
          tokenAmount: { amount: string };
        };
      };
    };
  };
};

// Function to get EVM provider based on chainId
const getEVMProvider = (chainId: number): ethers.JsonRpcProvider => {
  const urls = getRpcUrls();
  const rpcUrl = urls[chainId as keyof typeof urls]?.rpc[0];
  if (!rpcUrl) {
    throw new Error(`No RPC URL configured for chain ID ${chainId}`);
  }

  return new ethers.JsonRpcProvider(rpcUrl);
};

// Function to fetch EVM balance
const fetchEVMBalance = async (address: string, tokenAddress: string, chainId: number): Promise<bigint> => {
  if (!address.startsWith('0x')) {
    return BigInt(0);
  }

  const provider = getEVMProvider(chainId);

  if (isNativeToken(tokenAddress)) {
    const balance = await provider.getBalance(address);
    return BigInt(balance.toString());
  } else {
    const tokenContract = new ethers.Contract(tokenAddress, ERC20_ABI, provider);
    if (!tokenContract) {
      return BigInt(0);
    }
    const balance = await tokenContract.balanceOf?.(address);
    return BigInt(balance.toString());
  }
};

// Function to fetch Solana balance
const fetchSolanaBalance = async (address: string, tokenAddress: string): Promise<bigint> => {
  try {
    const ownerPublicKey = new PublicKey(address);
    const solanaProvider = await getSolanaProvider();

    if (compareAddresses(tokenAddress, SOL_NATIVE)) {
      const balance = await solanaProvider.getBalance(ownerPublicKey);
      return BigInt(balance);
    } else {
      const ownerPubKey = new PublicKey(address);
      const mintPubKey = new PublicKey(tokenAddress);

      const tokenAccounts = await Promise.all([
        solanaProvider.getParsedTokenAccountsByOwner(ownerPubKey, { programId: TOKEN_PROGRAM_ID }),
        solanaProvider.getParsedTokenAccountsByOwner(ownerPubKey, { programId: TOKEN_2022_PROGRAM_ID }),
      ]);

      const mint58 = mintPubKey.toBase58();

      const flatAccounts = tokenAccounts.flatMap((accounts) => accounts.value);
      const tokenAccount = flatAccounts.find(
        (account: TokenAccount) => account?.account?.data?.parsed?.info?.mint === mint58,
      );

      if (!tokenAccount) {
        return BigInt(0);
      }
      return BigInt(tokenAccount.account.data.parsed.info.tokenAmount.amount);
    }
  } catch (error) {
    console.error('Solana balance fetch error:', error);
    return BigInt(0);
  }
};

export const useBalance = (userChainAddress: string, tokenAddress: string, chainId: number): BalanceResult => {
  const address = userChainAddress;

  const fetchBalance = useCallback(async (): Promise<bigint> => {
    try {
      // Return 0n if no address is available
      if (!address) return BigInt(0);

      if (isEVMChain(chainId)) {
        return await fetchEVMBalance(address, tokenAddress, chainId);
      } else if (isSolanaChain(chainId)) {
        return await fetchSolanaBalance(address, tokenAddress);
      } else {
        throw new Error('Unsupported chain');
      }
    } catch (err) {
      console.error('Balance fetch error:', err);
      return BigInt(0);
    }
  }, [address, tokenAddress, chainId]);

  const {
    data: balance,
    isLoading: loading,
    error,
  } = useQuery({
    queryKey: ['userAccountBalance', address, tokenAddress, chainId],
    queryFn: fetchBalance,
    refetchInterval: REFETCH_INTERVAL,
    enabled: !!address && !!tokenAddress && !!chainId,
  });

  return { balance, loading, error: error as Error | null };
};