import os

from web3 import Web3

native_aliases = [
    "0x0000000000000000000000000000000000000000",
    "0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee",
]

erc20_abi = [
    {
        "constant": True,
        "inputs": [{"name": "_owner", "type": "address"}],
        "name": "balanceOf",
        "outputs": [{"name": "balance", "type": "uint256"}],
        "type": "function",
    }
]


def get_token_balance(token_address, wallet_address, block_number):
    rpc_url = os.getenv("RPC_URL")

    if rpc_url is None:
        raise EnvironmentError("RPC_URL environment variable not set")

    web3 = Web3(Web3.HTTPProvider(rpc_url))

    if not web3.isConnected():
        raise ConnectionError("Failed to connect to the Ethereum node")

    # Check if the token_address is a native token alias
    if token_address.lower() in native_aliases:
        balance = web3.eth.get_balance(
            Web3.toChecksumAddress(wallet_address), block_identifier=block_number
        )
    else:
        contract = web3.eth.contract(
            address=Web3.toChecksumAddress(token_address), abi=erc20_abi
        )
        balance = contract.functions.balanceOf(
            Web3.toChecksumAddress(wallet_address)
        ).call(block_identifier=block_number)

    return balance


def get_block_header(block_number):
    rpc_url = os.getenv("RPC_URL")

    if rpc_url is None:
        raise EnvironmentError("RPC_URL environment variable not set")

    web3 = Web3(Web3.HTTPProvider(rpc_url))

    if not web3.isConnected():
        raise ConnectionError("Failed to connect to the Ethereum node")

    block = web3.eth.get_block(block_number)

    return block
