import contextlib
import logging
from typing import Any, Callable, Iterable, List, Optional, Tuple, Union

import a_sync
import brownie
import dank_mids
from a_sync import igather
from faster_eth_abi.exceptions import InsufficientDataBytes
from multicall import Call
from web3.exceptions import CannotHandleRequest

from y import convert
from y._decorators import stuck_coro_debugger
from y.contracts import Contract, contract_creation_block
from y.datatypes import AddressOrContract, AnyAddressType, Block
from y.exceptions import continue_if_call_reverted
from y.interfaces.multicall2 import MULTICALL2_ABI
from y.networks import Network
from y.utils.raw_calls import _decimals, _totalSupply

logger = logging.getLogger(__name__)

MULTICALL2 = {
    Network.Mainnet: "0x5BA1e12693Dc8F9c48aAD8770482f4739bEeD696",
    Network.Arbitrum: "0x842eC2c7D803033Edf55E478F461FC547Bc54EB2",
    Network.Avalanche: "0xdf2122931FEb939FB8Cf4e67Ea752D1125e18858",
    Network.BinanceSmartChain: "0xfF6FD90A470Aaa0c1B8A54681746b07AcdFedc9B",
    Network.Fantom: "0xBAD2B082e2212DE4B065F636CA4e5e0717623d18",
    Network.Harmony: "0x34b415f4d3b332515e66f70595ace1dcf36254c5",
    Network.Heco: "0xd1F3BE686D64e1EA33fcF64980b65847aA43D79C",
    Network.Moonriver: "0xaeF00A0Cf402D9DEdd54092D9cA179Be6F9E5cE3",
    Network.Polygon: "0xc8E51042792d7405184DfCa245F2d27B94D013b6",
    Network.xDai: "0x9903f30c1469d8A2f415D4E8184C93BD26992573",
    Network.Aurora: "0xe0e3887b158F7F9c80c835a61ED809389BC08d1b",
    Network.Cronos: "0x5e954f5972EC6BFc7dECd75779F10d848230345F",
    Network.Optimism: "0xcA11bde05977b3631167028862bE2a173976CA11",  # Multicall 3
    Network.Base: "0xcA11bde05977b3631167028862bE2a173976CA11",  # mc3
}.get(brownie.chain.id)

multicall = None
multicall2 = (
    brownie.Contract.from_abi("Multicall2", MULTICALL2, MULTICALL2_ABI)
    if brownie.chain.id in [Network.Harmony, Network.Cronos]
    else Contract(MULTICALL2)
)

# the address doesn't matter, it just needs to have code
code = " "

multicall_deploy_block = contract_creation_block(multicall2.address)


@a_sync.a_sync(default="sync")
@stuck_coro_debugger
async def multicall_same_func_no_input(
    addresses: Iterable[AnyAddressType],
    method: str,
    block: Optional[Block] = None,
    apply_func: Optional[Callable] = None,
    return_None_on_failure: bool = False,
) -> List[Any]:

    addresses = [await convert.to_address_async(address) for address in addresses]
    results: List[dict] = await igather(
        Call(address, [method], ((address, apply_func)), block_id=block) for address in addresses
    )
    return [v for call in results for v in call.values()]


@a_sync.a_sync(default="sync")
@stuck_coro_debugger
async def multicall_same_func_same_contract_different_inputs(
    address: AnyAddressType,
    method: str,
    inputs: Union[List, Tuple],
    block: Optional[Block] = None,
    apply_func: Optional[Callable] = None,
    return_None_on_failure: bool = False,
) -> List[Any]:
    assert inputs
    address = await convert.to_address_async(address)
    results: List[dict] = await igather(
        (Call(address, [method, input], [(input, apply_func)], block_id=block) for input in inputs),
        return_exceptions=return_None_on_failure,
    )
    if return_None_on_failure:
        for i, result in enumerate(results[:]):
            if isinstance(result, Exception):
                continue_if_call_reverted(result)
                results[i] = {i: None}
    return [result for call in results for result in call.values()]


@a_sync.a_sync(default="sync")
@stuck_coro_debugger
async def multicall_decimals(
    addresses: Iterable[AddressOrContract],
    block: Optional[Block] = None,
    return_None_on_failure: bool = True,
) -> List[int]:

    addresses = tuple(map(str, addresses))
    try:
        calls = a_sync.map(Call, addresses, function="decimals()(uint256)", block_id=block)
        return [decimals async for address, decimals in calls]
    except (CannotHandleRequest, InsufficientDataBytes):
        pass  # TODO investigate these
    except Exception as e:
        continue_if_call_reverted(e)

    return await a_sync.map(
        _decimals, addresses, block=block, return_None_on_failure=return_None_on_failure
    ).values(pop=True)


@a_sync.a_sync(default="sync")
@stuck_coro_debugger
async def multicall_totalSupply(
    addresses: Iterable[AddressOrContract],
    block: Optional[Block] = None,
    return_None_on_failure: bool = True,
) -> List[int]:

    with contextlib.suppress(CannotHandleRequest, InsufficientDataBytes):
        return await multicall_same_func_no_input(
            addresses, "totalSupply()(uint256)", block=block, sync=False
        )
    return [
        await _totalSupply(
            address,
            block=block,
            return_None_on_failure=return_None_on_failure,
            sync=False,
        )
        for address in addresses
    ]


# yLazyLogger(logger)
@a_sync.a_sync(default="sync")
@stuck_coro_debugger
async def fetch_multicall(*calls: Any, block: Optional[Block] = None) -> List[Optional[Any]]:
    # https://github.com/makerdao/multicall
    multicall_input = []
    fn_list = []
    decoded = []

    for contract, fn_name, *fn_inputs in calls:
        fn = getattr(contract, fn_name)

        # check that there aren't multiple functions with the same name
        if hasattr(fn, "_get_fn_from_args"):
            fn = fn._get_fn_from_args(fn_inputs)

        fn_list.append(fn)
        multicall_input.append((contract, fn.encode_input(*fn_inputs)))

    if isinstance(block, int) and block < multicall_deploy_block:
        # use state override to resurrect the contract prior to deployment
        data = multicall2.tryAggregate.encode_input(False, multicall_input)
        call = await dank_mids.eth.call(
            {"to": str(multicall2), "data": data},
            block or "latest",
            {str(multicall2): {"code": f"0x{multicall2.bytecode}"}},  # type: ignore [dict-item, typeddict-item]
        )
        result = multicall2.tryAggregate.decode_output(call)
    else:
        result = await multicall2.tryAggregate.coroutine(
            False, multicall_input, block_identifier=block or "latest"
        )

    for fn, (ok, data) in zip(fn_list, result):
        try:
            assert ok, "call failed"
            decoded.append(fn.decode_output(data))
        except (AssertionError, InsufficientDataBytes):
            decoded.append(None)

    return decoded
