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

import a_sync
import brownie
from brownie import chain
from eth_abi.exceptions import InsufficientDataBytes
from multicall import Call
from web3.exceptions import CannotHandleRequest

from y import convert
from y.contracts import Contract, contract_creation_block
from y.datatypes import Address, AddressOrContract, AnyAddressType, Block
from y.decorators import stuck_coro_debugger
from y.exceptions import continue_if_call_reverted
from y.interfaces.multicall2 import MULTICALL2_ABI
from y.networks import Network
from y.utils.dank_mids import dank_w3
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(chain.id)

multicall = None
multicall2 = brownie.Contract.from_abi("Multicall2",MULTICALL2, MULTICALL2_ABI) if 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 = _clean_addresses(addresses)
    results = await asyncio.gather(*[Call(address, [method], [[address,apply_func]], block_id=block).coroutine() for address in addresses])
    return [v for call in results for k, v in call.items()]


@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 = convert.to_address(address)
    results = await asyncio.gather(
        *[Call(address, [method, input], [[input, apply_func]], block_id=block).coroutine() 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 key, result in call.items()]


@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]:

    try: 
        return await asyncio.gather(*[Call(str(address), ['decimals()(uint256)'], block_id=block).coroutine() for address in addresses])
    except (CannotHandleRequest, InsufficientDataBytes):
        pass # TODO investigate these
    except Exception as e:
        continue_if_call_reverted(e)

    return await asyncio.gather(*[_decimals(address,block=block,return_None_on_failure=return_None_on_failure) for address in addresses])

@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_w3.eth.call(
            {'to': str(multicall2), 'data': data},
            block or 'latest',
            {str(multicall2): {'code': f'0x{multicall2.bytecode}'}},
        )
        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


#yLazyLogger(logger)
def _clean_addresses(
    addresses: Iterable[AnyAddressType]
    ) -> List[Address]:

    return [convert.to_address(address) for address in addresses]
