import { ChainId } from '@lifi/types';
import { config } from '@lifi/sdk';
import { PublicKey, Connection } from '@solana/web3.js';

const TokenProgramAddress = 'TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA';
const SolSystemProgram = '11111111111111111111111111111111';

const getRpcUrl = async (chainId: ChainId): Promise<string> => {
  const rpcUrls = await getRpcUrls(chainId);
  return rpcUrls[0];
};

const getRpcUrls = async (chainId: ChainId): Promise<string[]> => {
  const rpcUrls = (await config.getRPCUrls())[chainId];
  if (!rpcUrls?.length) {
    throw new Error('RPC URL not found');
  }
  return rpcUrls;
};

let connection: Connection | undefined = undefined;

/**
 * getSolanaConnection is just a thin wrapper around getting the connection (RPC provider) for Solana
 * @returns - Solana RPC connection
 */
const getSolanaConnection = async (): Promise<Connection> => {
  if (!connection) {
    const rpcUrl = await getRpcUrl(ChainId.SOL);
    connection = new Connection(rpcUrl);
    return connection;
  } else {
    return connection;
  }
};

export const getSolanaBalance = async (
  walletAddress: string,
  tokens: any[],
): Promise<any[]> => {
  if (tokens.length === 0) {
    return [];
  }
  const { chainId } = tokens[0];
  tokens.forEach((token) => {
    if (token.chainId !== chainId) {
      console.warn(`Requested tokens have to be on the same chain.`);
    }
  });

  return getSolanaBalanceDefault(chainId, tokens, walletAddress);
};

const getSolanaBalanceDefault = async (
  _chainId: any,
  tokens: any[],
  walletAddress: string,
): Promise<any[]> => {
  const connection = await getSolanaConnection();
  const accountPublicKey = new PublicKey(walletAddress);
  const tokenProgramPublicKey = new PublicKey(TokenProgramAddress);
  const [slot, balance, tokenAccountsByOwner] = await Promise.allSettled([
    connection.getSlot(),
    connection.getBalance(accountPublicKey),
    connection.getParsedTokenAccountsByOwner(accountPublicKey, {
      programId: tokenProgramPublicKey,
    }),
  ]);
  const blockNumber = slot.status === 'fulfilled' ? BigInt(slot.value) : 0n;
  const solBalance =
    balance.status === 'fulfilled' ? BigInt(balance.value) : 0n;
  const walletTokenAmounts =
    tokenAccountsByOwner.status === 'fulfilled'
      ? tokenAccountsByOwner.value.value.reduce(
          (tokenAmounts: any, value: any) => {
            const amount = BigInt(
              value.account.data.parsed.info.tokenAmount.amount,
            );
            if (amount > 0n) {
              tokenAmounts[value.account.data.parsed.info.mint] = amount;
            }
            return tokenAmounts;
          },
          {} as Record<string, bigint>,
        )
      : {};
  walletTokenAmounts[SolSystemProgram] = solBalance;
  const tokenAmounts: any[] = tokens.map((token) => {
    if (walletTokenAmounts[token.address]) {
      return {
        ...token,
        amount: walletTokenAmounts[token.address],
        blockNumber,
      };
    }
    return {
      ...token,
      blockNumber,
    };
  });
  return tokenAmounts;
};