@swapper-finance/sdk/src/contexts/swapProvider/useTokensState.ts

JavaScript SDK form Swapper

import { useState, useCallback, useMemo, useEffect } from "react";
import {
  TokenWithChain,
  Token,
  TokenPrices,
  SwapperAppPayload,
} from "@src/models";
import {
  isServer,
  deepMergeObjects,
  convertTokensBalances,
  logger,
} from "@src/utils";
import { Chain } from "@src/models";
import { getPrices, getBalances } from "@src/services";
import { Balances } from "@src/interfaces";
import {
  CASH_TOKEN_ADDRESS_BY_CHAIN,
  DEFAULT_CHAIN_ID,
  DEFAULT_TOKEN_ADDR,
} from "@src/config";
import { ethers } from "ethers";
import { ERC20Abi } from "@src/contractAbis";

const CUSTOM_TOKENS_KEY = "swapper-custom-tokens";

export const useTokensState = (
  chain: Chain | undefined,
  setChain: React.Dispatch<React.SetStateAction<Chain | undefined>>,
  supportedChains: Chain[],
  integrationConfig: Partial<SwapperAppPayload>,
  smartWalletAddress: string | undefined,
) => {
  const [token, setToken] = useState<TokenWithChain | undefined>(undefined);
  const [displayToken, setDisplayToken] = useState<TokenWithChain | undefined>(
    undefined,
  );
  const [transferToken, setTransferToken] = useState<
    TokenWithChain | undefined
  >(undefined);
  const [tokensPrices, setTokensPrices] = useState<TokenPrices>({});
  const [tokensBalances, setTokensBalances] = useState<
    Record<string, Record<string, Balances>>
  >({});
  const [isFetchingBalances, setIsFetchingBalances] = useState(false);
  const [customTokensMap, setCustomTokensMap] = useState<
    Record<string, Record<string, Token>>
  >(() => {
    if (isServer) return {};
    try {
      return JSON.parse(localStorage.getItem(CUSTOM_TOKENS_KEY) || "{}");
    } catch {
      return {};
    }
  });

  const saveCustomTokens = useCallback((newMap) => {
    setCustomTokensMap(newMap);
    try {
      localStorage.setItem(CUSTOM_TOKENS_KEY, JSON.stringify(newMap));
    } catch (e) {
      console.error("Failed to save custom tokens:", e);
    }
  }, []);

  const { dstChainId, dstTokenAddr, dstTokenImage } = integrationConfig;

  const supportedTokens: Record<string, Record<string, Token>> = useMemo(() => {
    const chainIdToTokenAddressToData: {
      [chainId: string]: {
        [tokenAddress: string]: Token;
      };
    } = {};
    supportedChains.forEach((chain) => {
      chainIdToTokenAddressToData[chain.chainId] = {};
      chain.tokens.forEach((token) => {
        // @ts-expect-error TODO
        chainIdToTokenAddressToData[chain.chainId][token.address] = token;
      });
    });
    return deepMergeObjects(chainIdToTokenAddressToData, customTokensMap);
  }, [supportedChains, customTokensMap]);

  const chainToTokenOptionsMap = useMemo(() => {
    const tempChainToTokenOptionsMap: Record<string, TokenWithChain[]> = {};
    Object.entries(supportedTokens).forEach(([chainId, chainData]) => {
      tempChainToTokenOptionsMap[chainId] = Object.entries(chainData)
        .filter(([, tokenData]) => tokenData.supported)
        .sort(([, tokenDataA], [, tokenDataB]) => {
          return tokenDataB.priority - tokenDataA.priority;
        })
        .map(([, tokenData]) => {
          return {
            ...tokenData,
            chainId: chainId,
          };
        });
    });
    return tempChainToTokenOptionsMap;
  }, [supportedTokens]);

  const allTokens = useMemo(() => {
    const tokens: TokenWithChain[] = [];
    const chainsIds = Object.keys(chainToTokenOptionsMap);
    chainsIds.forEach((chainId) => {
      const chainTokens = chainToTokenOptionsMap[chainId]!;
      chainTokens.forEach((token) => {
        tokens.push(token);
      });
    });
    return tokens;
  }, [chainToTokenOptionsMap]);

  const cashToken = useMemo(
    () =>
      allTokens.find(
        (token) =>
          token.address ===
          CASH_TOKEN_ADDRESS_BY_CHAIN[chain?.chainId || token.chainId],
      ),
    [allTokens, chain?.chainId],
  );

  const fetchAllTokenPrices = useCallback(async () => {
    try {
      if (supportedChains.length > 0) {
        const prices = await getPrices({
          chainId: supportedChains.map((chain) => chain.chainId),
        });
        setTokensPrices(prices);
      }
    } catch (err) {
      console.error(err);
    }
  }, [supportedChains]);

  // Fetch balances function (moved from useBalancesState)
  const fetchAllBalances = useCallback(async () => {
    logger.log("fetchingBalances", { smartWalletAddress, allTokens });
    if (!smartWalletAddress || allTokens.length === 0) {
      setIsFetchingBalances(false);
      return;
    }
    try {
      logger.log("fetchingBalances2");

      setIsFetchingBalances(true);
      const balances = await getBalances({ walletAddress: smartWalletAddress });
      setTokensBalances(convertTokensBalances(balances, allTokens));
    } catch (err) {
      console.error(err);
    } finally {
      setIsFetchingBalances(false);
    }
  }, [smartWalletAddress, allTokens]);

  // Initialize prices
  useEffect(() => {
    fetchAllTokenPrices();
  }, [fetchAllTokenPrices]);

  // Initialize balances
  useEffect(() => {
    fetchAllBalances();
  }, [fetchAllBalances]);

  useEffect(() => {
    if (supportedChains.length === 0) {
      return;
    }
    if (!dstChainId || !dstTokenAddr) {
      // Prefer default chain; otherwise fallback to first available; if none, throw
      let selectedChain: Chain | undefined;
      const defaultChainCandidate = supportedChains.find(
        (c) => c.chainId === DEFAULT_CHAIN_ID,
      );
      if (defaultChainCandidate) {
        selectedChain = defaultChainCandidate;
      } else {
        const firstSupportedChain = supportedChains[0];
        selectedChain = firstSupportedChain;
      }

      if (!selectedChain) {
        throw new Error("No supported chains available");
      }

      // Select a token for the selected chain
      let selectedToken = undefined as TokenWithChain | undefined;
      const selectedChainToken = chainToTokenOptionsMap[
        selectedChain.chainId
      ]?.find(
        (t) => t.address.toLowerCase() === DEFAULT_TOKEN_ADDR.toLowerCase(),
      );

      if (selectedChainToken) {
        selectedToken = selectedChainToken;
      } else {
        const tokensForSelectedChain =
          chainToTokenOptionsMap[selectedChain?.chainId || DEFAULT_CHAIN_ID];
        selectedToken = tokensForSelectedChain?.[0];
      }

      setChain(selectedChain);
      setToken(selectedToken);
      return;
    }

    const addr = dstTokenAddr.toLowerCase();

    const matchedToken = chainToTokenOptionsMap[dstChainId]?.find(
      (t) => t.address.toLowerCase() === addr,
    );

    if (matchedToken) {
      if (matchedToken.isCustom && matchedToken.image !== dstTokenImage) {
        const updatedToken = { ...matchedToken, image: dstTokenImage };
        saveCustomTokens({
          ...customTokensMap,
          [dstChainId]: {
            ...customTokensMap[dstChainId],
            [addr]: updatedToken,
          },
        });
        matchedToken.image = dstTokenImage;
      }

      setChain(supportedChains.find((c) => c.chainId === dstChainId));
      setToken(matchedToken);
      return;
    }

    (async () => {
      try {
        const chainObj = supportedChains.find((c) => c.chainId === dstChainId)!;
        if (!chainObj) return;

        const provider = new ethers.providers.JsonRpcProvider(
          chainObj.publicRpcUrls[0],
        );
        const contract = new ethers.Contract(addr, ERC20Abi, provider);
        const [name, symbol, decimals] = await Promise.all([
          contract.name(),
          contract.symbol(),
          contract.decimals(),
        ]);

        const newToken: TokenWithChain = {
          address: addr,
          name,
          symbol,
          tokenId: addr,
          decimals: Number(decimals),
          image: dstTokenImage,
          supported: true,
          quickPick: false,
          priority: 0,
          chainId: dstChainId,
          isCustom: true,
        };

        const updatedMap = {
          ...customTokensMap,
          [dstChainId]: {
            ...(customTokensMap[dstChainId] || {}),
            [addr]: newToken,
          },
        };
        saveCustomTokens(updatedMap);

        setChain(chainObj);
        setToken(newToken);
      } catch (e) {
        console.error("Preload custom token failed:", e);
      }
    })();
  }, [
    dstChainId,
    dstTokenAddr,
    dstTokenImage,
    supportedChains,
    chainToTokenOptionsMap,
    customTokensMap,
    saveCustomTokens,
    setChain,
    setToken,
  ]);

  return {
    // Token state
    token,
    setToken,
    displayToken,
    setDisplayToken,
    transferToken,
    setTransferToken,
    tokensPrices,
    allTokens,
    cashToken,
    chainToTokenOptionsMap,

    // Balance state (merged from useBalancesState)
    tokensBalances,
    setTokensBalances,
    isFetchingBalances,
    fetchAllBalances,
  };
};