import logging
from collections import defaultdict
from itertools import count, product
from operator import itemgetter
from typing import Any, Callable, Iterable, List, Optional, Tuple, Union

import brownie
import requests
from brownie import chain, web3
from eth_abi.exceptions import InsufficientDataBytes
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.exceptions import continue_if_call_reverted
from y.interfaces.multicall2 import MULTICALL2_ABI
from y.networks import Network
from y.utils.logging import yLazyLogger
from y.utils.raw_calls import _decimals, _totalSupply

from multicall import Call
from multicall.utils import await_awaitable, gather

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",
}.get(chain.id, None)

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)


#yLazyLogger(logger)
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]:

    return await_awaitable(
        multicall_same_func_no_input_async(addresses, method, block=block, apply_func=apply_func, return_None_on_failure=return_None_on_failure)
    )
    

#yLazyLogger(logger)
async def multicall_same_func_no_input_async(
    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 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()]


#yLazyLogger(logger)
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]:

    return await_awaitable(
        multicall_same_func_same_contract_different_inputs_async(address, method, inputs, block=block, apply_func=apply_func, return_None_on_failure=return_None_on_failure)
    )


#yLazyLogger(logger)
async def multicall_same_func_same_contract_different_inputs_async(
    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 gather(Call(address, [method, input], [[input,apply_func]], block_id=block).coroutine() for input in inputs)
    return [result for call in results for key, result in call.items()]


#yLazyLogger(logger)
def multicall_decimals(
    addresses: Iterable[AddressOrContract], 
    block: Optional[Block] = None,
    return_None_on_failure: bool = True
    ) -> List[int]:
    return await_awaitable(
        multicall_decimals_async(addresses, block=block, return_None_on_failure=return_None_on_failure)
    )

#yLazyLogger(logger)
async def multicall_decimals_async(
    addresses: Iterable[AddressOrContract], 
    block: Optional[Block] = None,
    return_None_on_failure: bool = True
    ) -> List[int]:

    try: 
        return await 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)

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


#yLazyLogger(logger)
def multicall_totalSupply(
    addresses: Iterable[AddressOrContract], 
    block: Optional[Block] = None,
    return_None_on_failure: bool = True
    ) -> List[int]:
    return await_awaitable(
        multicall_totalSupply_async(addresses, block=block, return_None_on_failure=return_None_on_failure)
    )

#yLazyLogger(logger)
async def multicall_totalSupply_async(
    addresses: Iterable[AddressOrContract], 
    block: Optional[Block] = None,
    return_None_on_failure: bool = True
    ) -> List[int]:

    try:
        return await multicall_same_func_no_input_async(addresses, 'totalSupply()(uint256)', block=block)
    except (CannotHandleRequest,InsufficientDataBytes):
        pass
        
    return [await _totalSupply(address,block=block,return_None_on_failure=return_None_on_failure) for address in addresses] 


#yLazyLogger(logger)
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 = web3.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 = multicall2.tryAggregate.call(
            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 multicall_matrix(contracts, params, block="latest"):
    matrix = list(product(contracts, params))
    calls = [[contract, param] for contract, param in matrix]

    results = fetch_multicall(*calls, block=block)

    output = defaultdict(dict)
    for (contract, param), value in zip(matrix, results):
        output[contract][param] = value

    return dict(output)


#yLazyLogger(logger)
def batch_call(calls: Tuple[Contract,str,List[Any],Block]) -> List[Any]:
    """
    Similar interface but block height as last param. Uses JSON-RPC batch.
    [[contract, 'func', arg, block_identifier]]
    """
    jsonrpc_batch = []
    fn_list = []
    ids = count()

    for contract, fn_name, *fn_inputs, block in calls:
        fn = getattr(contract, fn_name)
        if hasattr(fn, "_get_fn_from_args"):
            fn = fn._get_fn_from_args(fn_inputs)
        fn_list.append(fn)

        jsonrpc_batch.append(
            {
                'jsonrpc': '2.0',
                'id': next(ids),
                'method': 'eth_call',
                'params': [
                    {'to': str(contract), 'data': fn.encode_input(*fn_inputs)},
                    block,
                ],
            }
        )

    response = requests.post(web3.provider.endpoint_uri, json=jsonrpc_batch).json()
    return [
        fn.decode_output(res['result'])
        for res in sorted(response, key=itemgetter('id'))
    ]


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

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