import { Connection, PublicKey } from "@solana/web3.js";
import { useQuery } from "@tanstack/react-query";
import { isPublicKey } from "@utils/isPublicKey";
import { useConfig } from "./useConfig";

const TOKEN_PROGRAM_ID = "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA";

type QueryResult = {
	nativeBalance: string;
	tokenBalance: string | undefined;
};

const useSolanaBalance = ({
	userAddress,
	tokenAddress,
	enabled = true,
}: {
	userAddress: string;
	tokenAddress?: string;
	enabled?: boolean;
}) => {
	const { rpcs } = useConfig();
	return useQuery({
		queryKey: ["solanaBalance", { userAddress, tokenAddress, enabled, rpcs }],
		queryFn: () => {
			return getSolanaBalance({
				userAddress,
				rpc: rpcs?.solana,
				tokenAddress,
			});
		},
		enabled: Boolean(
			userAddress && tokenAddress && isPublicKey(userAddress) && enabled,
		),
		retry: false,
	});
};

export default useSolanaBalance;

export const getSolanaBalance = async ({
	userAddress,
	rpc,
	tokenAddress,
}: {
	userAddress: string;
	rpc: string | undefined;
	tokenAddress?: string;
}): Promise<QueryResult> => {
	const pubAddress = userAddress ? new PublicKey(userAddress) : "";

	if (!pubAddress || !rpc)
		return {
			nativeBalance: "0",
			tokenBalance: "0",
		};

	const connection = new Connection(rpc);
	const nativeBalance = (await connection.getBalance(pubAddress)).toString();

	const tokenInfo = (
		await connection.getParsedTokenAccountsByOwner(pubAddress, {
			programId: new PublicKey(TOKEN_PROGRAM_ID),
		})
	).value.map((token) => {
		return {
			mint: token.account.data.parsed?.info.mint as string,
			uiAmount: token.account.data.parsed?.info.tokenAmount.uiAmount as string,

			amount: token.account.data.parsed?.info.tokenAmount.amount as string,
		};
	});

	return {
		nativeBalance,
		tokenBalance: tokenInfo.find(
			(token) => token.mint.toLowerCase() === tokenAddress?.toLowerCase(),
		)?.amount,
	};
};
