import { OrangeKitConnector } from "@mezo-org/orangekit"
import { useQuery } from "@tanstack/react-query"
import { useMemo } from "react"
import { Address } from "viem"
import { useAccount, UseAccountReturnType } from "wagmi"

type UseWalletAccountReturn = Omit<UseAccountReturnType, "address"> & {
  accountAddress?: Address
  walletAddress?: string
  networkFamily: "bitcoin" | "evm"
}

export function useWalletAccount(): UseWalletAccountReturn {
  const {
    address: evmAddress,
    connector,
    chainId,
    isConnected,
    ...rest
  } = useAccount()

  const networkFamily = useMemo(
    () => (connector?.type !== "orangekit" ? "evm" : "bitcoin"),
    [connector],
  )

  const { data: btcData } = useQuery({
    queryKey: [connector?.id],
    queryFn: async () => {
      const bitcoinProvider = (
        connector as unknown as OrangeKitConnector
      ).getBitcoinProvider()

      const address = await bitcoinProvider.getAddress()
      const balance = await bitcoinProvider.getBalance()

      return {
        address,
        balance,
      }
    },
    enabled: networkFamily === "bitcoin" && !!connector,
  })

  const { address: btcAddress } = btcData || {}

  return {
    accountAddress: evmAddress,
    walletAddress: btcAddress ?? evmAddress,
    isConnected:
      isConnected &&
      (networkFamily === "bitcoin" ? !!btcAddress : !!evmAddress),
    networkFamily,
    connector,
    chainId,
    ...rest,
  }
}