wormhole/bridge_ui/src/hooks/useGetBalanceEffect.ts

import { CHAIN_ID_ETH, CHAIN_ID_SOLANA } from "@certusone/wormhole-sdk";
import { Connection, PublicKey } from "@solana/web3.js";
import { formatUnits } from "ethers/lib/utils";
import { useEffect } from "react";
import { useDispatch, useSelector } from "react-redux";
import { useEthereumProvider } from "../contexts/EthereumProviderContext";
import { useSolanaWallet } from "../contexts/SolanaWalletContext";
import { TokenImplementation__factory } from "../ethers-contracts";
import {
  selectTransferSourceAsset,
  selectTransferSourceChain,
  selectTransferTargetAsset,
  selectTransferTargetChain,
} from "../store/selectors";
import {
  setSourceParsedTokenAccount,
  setTargetParsedTokenAccount,
} from "../store/transferSlice";
import { hexToUint8Array } from "../utils/array";
import { SOLANA_HOST } from "../utils/consts";

function createParsedTokenAccount(
  publicKey: PublicKey | undefined,
  amount: string,
  decimals: number,
  uiAmount: number,
  uiAmountString: string
) {
  return {
    publicKey: publicKey?.toString(),
    amount,
    decimals,
    uiAmount,
    uiAmountString,
  };
}

/**
 * Fetches the balance of an asset for the connected wallet
 * @param sourceOrTarget determines whether this will fetch balance for the source or target account. Not intended to be switched on the same hook!
 */
function useGetBalanceEffect(sourceOrTarget: "source" | "target") {
  const dispatch = useDispatch();
  const setAction =
    sourceOrTarget === "source"
      ? setSourceParsedTokenAccount
      : setTargetParsedTokenAccount;
  const lookupChain = useSelector(
    sourceOrTarget === "source"
      ? selectTransferSourceChain
      : selectTransferTargetChain
  );
  const lookupAsset = useSelector(
    sourceOrTarget === "source"
      ? selectTransferSourceAsset
      : selectTransferTargetAsset
  );
  const { wallet } = useSolanaWallet();
  const solPK = wallet?.publicKey;
  const { provider, signerAddress } = useEthereumProvider();
  useEffect(() => {
    // TODO: loading state
    dispatch(setAction(undefined));
    if (!lookupAsset) {
      return;
    }
    let cancelled = false;
    if (lookupChain === CHAIN_ID_SOLANA && solPK) {
      let mint;
      try {
        mint = new PublicKey(
          sourceOrTarget === "source"
            ? lookupAsset
            : hexToUint8Array(lookupAsset)
        );
      } catch (e) {
        return;
      }
      const connection = new Connection(SOLANA_HOST, "finalized");
      connection
        .getParsedTokenAccountsByOwner(solPK, { mint })
        .then(({ value }) => {
          if (!cancelled) {
            console.log("parsed token accounts", value);
            if (value.length) {
              // TODO: allow selection between these target accounts
              dispatch(
                setAction(
                  createParsedTokenAccount(
                    value[0].pubkey,
                    value[0].account.data.parsed?.info?.tokenAmount?.amount,
                    value[0].account.data.parsed?.info?.tokenAmount?.decimals,
                    value[0].account.data.parsed?.info?.tokenAmount?.uiAmount,
                    value[0].account.data.parsed?.info?.tokenAmount
                      ?.uiAmountString
                  )
                )
              );
            } else {
              // TODO: error state
            }
          }
        })
        .catch(() => {
          if (!cancelled) {
            // TODO: error state
          }
        });
    }
    if (lookupChain === CHAIN_ID_ETH && provider && signerAddress) {
      const token = TokenImplementation__factory.connect(lookupAsset, provider);
      token
        .decimals()
        .then((decimals) => {
          token.balanceOf(signerAddress).then((n) => {
            if (!cancelled) {
              dispatch(
                setAction(
                  // TODO: verify accuracy
                  createParsedTokenAccount(
                    undefined,
                    n.toString(),
                    decimals,
                    Number(formatUnits(n, decimals)),
                    formatUnits(n, decimals)
                  )
                )
              );
            }
          });
        })
        .catch(() => {
          if (!cancelled) {
            // TODO: error state
          }
        });
    }
    return () => {
      cancelled = true;
    };
  }, [
    dispatch,
    sourceOrTarget,
    setAction,
    lookupChain,
    lookupAsset,
    solPK,
    provider,
    signerAddress,
  ]);
}

export default useGetBalanceEffect;