import { useEffect, useMemo, useRef, useState } from 'react'
import type { Address, WalletCapabilities } from 'viem'
import { useConfig, } from 'wagmi'
import { getWalletClient } from 'wagmi/actions'

import { useAccounts } from './useAccounts'

const TIMEOUT_MS = 10000

// optional module-level cache (survives across hook instances)
const cache = new Map<string, WalletCapabilities | null>()

export function useWalletCapabilities(chainId?: number) {
  const config = useConfig();
  const [capabilities, setCapabilities] = useState<WalletCapabilities | null | undefined>(undefined)
  const [isLoading, setIsLoading] = useState<boolean>(false);
  const calledRef = useRef(false)
  const {accounts} = useAccounts();

  // evm account
  const account = accounts.find((account) => account.chainType === "EVM");
  const address = account?.address as Address;

  useEffect(() => {
    if (!address || calledRef.current || !chainId) return
    const cacheKey = `${chainId}:${address}`
    if (cache.has(cacheKey)) {
      setCapabilities(cache.get(cacheKey))
      calledRef.current = true
      return
    }

    let timeoutId: ReturnType<typeof setTimeout> | undefined

    const timeoutPromise = new Promise((_, reject) => {
      timeoutId = setTimeout(() => reject(new Error('getCapabilities timeout')), TIMEOUT_MS)
    })

    async function fetchCapabilities() {
      calledRef.current = true
      setIsLoading(true);

      try {
        const walletClient = await getWalletClient(config)
        const result = await Promise.race([
          walletClient?.getCapabilities({ account: address }),
          timeoutPromise,
        ]);
        if (result && typeof result === 'object') {
          cache.set(cacheKey, result as WalletCapabilities)
          setCapabilities(result as WalletCapabilities)
        } else {
          throw new Error(`Invalid capabilities format: ${JSON.stringify(result)}`)
        }
      } catch (err) {
        console.warn('[Capabilities]', err)
        cache.set(cacheKey, null)
        setCapabilities(null)
      } finally {
        if (timeoutId) clearTimeout(timeoutId)
        setIsLoading(false);
      }
    }

    fetchCapabilities()
  }, [config, address, chainId])

  // if the status is ready, the wallet will switch the eoa to smart account before proceeding
  const supportsAtomicBatching = useMemo(
		() => {
      if(!chainId) return false;
      const _capabilities = capabilities?.[chainId];
      return _capabilities?.atomicBatch?.supported || _capabilities?.atomic?.status === 'supported' || _capabilities?.atomic?.status === "ready"
    },
		[capabilities, chainId],
	);

  return {
    capabilities,
    supportsAtomicBatching,
    isLoading
  }
}
