import { useEffect, useRef } from "react";
import { type Address, erc20Abi, isAddress } from "viem";
import { usePublicClient } from "wagmi";

import type { Token } from "@/types/index";
import {
	type UseQueryOptions,
	useQuery,
	useQueryClient,
} from "@tanstack/react-query";
import { NATIVE_TOKEN_ADDRESS } from "@utils/constants";
import { getValidAddress } from "@utils/getValidAddress";
import { isDummyAddress } from "@utils/isDummyAddress";
import { isPublicKey } from "@utils/isPublicKey";
import { getQueryKey as getTokenBalanceQueryKey } from "./apis/common/useTokens";
import { useAccounts } from "./useAccounts";
import { useConfig } from "./useConfig";
import { getSolanaBalance } from "./useSolanaBalance";

type TokenBalance = {
	balance: string | bigint;
	isLoading: boolean;
};

const REFETCH_INTERVAL = 60_000;
const MAX_SOLANA_RETRIES = 10;
const SOLANA_POLLING_RETRY_INTERVAL = 1_000;

type TokenBalanceParams = {
	token: Token | undefined;
	shouldFetch?: boolean;
	query?: Omit<
		UseQueryOptions<
			string | bigint,
			Error,
			string | bigint,
			readonly unknown[]
		>,
		"queryKey" | "queryFn" | "enabled"
	>;
};

export const getQueryKey = ({
	tokenAddress,
	chainId,
	userAddress,
	solanaRpc,
}: {
	tokenAddress: string;
	chainId: number;
	userAddress: string;
	solanaRpc?: string;
}) => {
	return [
		"tokenBalance",
		tokenAddress?.toLowerCase(),
		chainId,
		userAddress?.toLowerCase(),
		solanaRpc,
	];
};

export const useTokenBalance = ({
	token,
	shouldFetch = true,
	query,
}: TokenBalanceParams): TokenBalance => {
	const { accounts } = useAccounts();

	const chainId = token?.chainId;
	const tokenAddress = token?.address;
	const isNativeToken = tokenAddress?.toLowerCase() === NATIVE_TOKEN_ADDRESS;

	const evmAddress = accounts.find(
		(account) => account.chainType === "EVM",
	)?.address;
	const solanaAddress = accounts.find(
		(account) => account.chainType === "SVM",
	)?.address;

	const userAddress = getValidAddress({
		chainId: chainId ?? null,
		address: evmAddress || null,
		fallbackAddress: solanaAddress || null,
	});

	const isValidEvmAddress = Boolean(
		userAddress && isAddress(userAddress) && !isDummyAddress(userAddress),
	);

	const isValidSolanaAddress = Boolean(
		userAddress && isPublicKey(userAddress) && !isDummyAddress(userAddress),
	);

	const publicClient = usePublicClient({ chainId });

	const { apiKey, rpcs } = useConfig();

	const queryClient = useQueryClient();

	const userTokenBalances = queryClient.getQueryData<{
		result: Record<number, Token[]>;
	}>(getTokenBalanceQueryKey({ address: userAddress, apiKey }));

	const tokenBalance =
		chainId !== undefined
			? userTokenBalances?.result?.[chainId]?.find(
					(token: Token) =>
						token.address.toLowerCase() === tokenAddress?.toLowerCase(),
				)?.balance
			: undefined;

	const currentQueryKey = getQueryKey({
		userAddress: userAddress || "",
		chainId: chainId ?? 0,
		tokenAddress: tokenAddress || "",
		solanaRpc: rpcs?.solana,
	});

	const shouldStopPolling = useRef(false);

	// Get previous balance from React Query cache to persist across unmount/mount
	const getCachedBalance = () => {
		const cachedData = queryClient.getQueryData<string | bigint>(
			currentQueryKey,
		);
		return cachedData ? cachedData.toString() : null;
	};

	const { data: balance = "0", isLoading } = useQuery<string | bigint>({
		queryKey: currentQueryKey,
		enabled: !!token && shouldFetch && !!userAddress,
		refetchInterval: shouldStopPolling.current ? false : REFETCH_INTERVAL,
		queryFn: async () => {
			if (!userAddress || !token) return "0";

			// 🟣 Solana - with polling for balance changes
			if (isValidSolanaAddress) {
				const fetchSolanaBalance = async (retryCount = 0): Promise<string> => {
					const balances = await getSolanaBalance({
						userAddress,
						rpc: rpcs?.solana,
						tokenAddress,
					});

					const newBalance = isNativeToken
						? balances.nativeBalance
						: balances.tokenBalance || "0";

					// Get previous balance from cache (persists across component unmount/mount)
					const previousBalance = getCachedBalance();

					// If this is the first fetch ever or we've reached max retries, return the balance
					if (!previousBalance || retryCount >= MAX_SOLANA_RETRIES) {
						return newBalance;
					}

					// If balance hasn't changed and we haven't maxed out retries, poll again
					if (
						newBalance === previousBalance &&
						retryCount < MAX_SOLANA_RETRIES
					) {
						await new Promise((resolve) =>
							setTimeout(resolve, SOLANA_POLLING_RETRY_INTERVAL),
						); // Wait 1 second before retry
						return fetchSolanaBalance(retryCount + 1);
					}

					// Balance has changed or max retries reached
					return newBalance;
				};

				return fetchSolanaBalance(0);
			}

			// 🔵 EVM Native ETH
			if (isNativeToken && isValidEvmAddress) {
				const balance = await publicClient?.getBalance({
					address: userAddress as Address,
				});
				return balance?.toString() || "0";
			}

			// 🟡 EVM ERC20 Token
			if (!isNativeToken && isValidEvmAddress) {
				const balance = await publicClient?.readContract({
					abi: erc20Abi,
					address: tokenAddress as `0x${string}`,
					functionName: "balanceOf",
					args: [userAddress as Address],
				});

				return balance?.toString() || "0";
			}

			return "0";
		},
		retry: (value) => {
			if (value > 1) {
				shouldStopPolling.current = true;
				return false;
			}
			return true;
		},
		initialData: tokenBalance,
		...query,
	});

	return {
		balance,
		isLoading,
	};
};
