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

import brownie
import requests
from brownie import Contract as _Contract
from brownie import chain, convert, web3
from brownie.convert.datatypes import EthAddress
from eth_abi.exceptions import InsufficientDataBytes
from eth_typing.evm import Address, BlockNumber
from web3.exceptions import CannotHandleRequest
from y.contracts import Contract, contract_creation_block
from y.decorators import log
from y.exceptions import continue_if_call_reverted
from y.networks import Network
from ypricemagic.interfaces.multicall2 import MULTICALL2_ABI
from ypricemagic.utils.raw_calls import _decimals, _totalSupply

from multicall import Call, Multicall

logger = logging.getLogger(__name__)

SUPPORTED_INPUT_TYPES = str, Address, EthAddress, brownie.Contract, Contract

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

multicall = None
multicall2 = _Contract.from_abi("Multicall2",MULTICALL2, MULTICALL2_ABI) if chain.id == Network.Harmony else Contract(MULTICALL2)

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

multicall_deploy_block = contract_creation_block(multicall2.address)


@log(logger)
def multicall_same_func_no_input(
    addresses: list,
    method: str, 
    block = None,
    apply_func = None
    ):

    addresses = _clean_addresses(addresses)
    calls = [Call(address, [method], [[address,apply_func]]) for address in addresses]
    return [result for result in Multicall(calls, block_id=block, _w3=web3)().values()]


@log(logger)
def multicall_same_func_different_contracts_same_input(
    addresses: list, 
    method: str, 
    input = None, 
    block: Union[int, BlockNumber, None] = None,
    apply_func = None
    ):

    assert input
    addresses = _clean_addresses(addresses)
    calls = [Call(address, [method, input], [[address,apply_func]]) for address in addresses]
    return [result for result in Multicall(calls, block_id=block, _w3=web3)().values()]


@log(logger)
def multicall_same_func_same_contract_different_inputs(
    address: Union[str, Address, brownie.Contract, Contract], 
    method: str, 
    inputs: Union[List, Tuple],  
    block: Union[int, BlockNumber, None] = None,
    apply_func = None,
    return_None_on_failure: bool = False
    ):

    assert inputs
    address = _clean_address(address)
    calls = [Call(address, [method, input], [[input,apply_func]]) for input in inputs]
    return [result for result in Multicall(calls, block_id=block, _w3=web3, require_success = not return_None_on_failure)().values()]


@log(logger)
def multicall_decimals(
    addresses: List[str], 
    block: Union[int, BlockNumber, None] = None,
    return_None_on_failure: bool = True
    ):

    try: 
        return multicall_same_func_no_input(addresses, 'decimals()(uint256)', block=block)
    except (CannotHandleRequest,InsufficientDataBytes):
        pass # TODO investigate these
    except Exception as e:
        continue_if_call_reverted(e)

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


@log(logger)
def multicall_totalSupply(
    addresses: List[str], 
    block: Union[int, BlockNumber, None] = None,
    return_None_on_failure: bool = True
    ):

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


@log(logger)
def multicall_balanceOf(
    token_addresses: List[str], 
    hodler_address: str, 
    block: Union[int, BlockNumber, None] = None,
    return_None_on_failure: bool = True # TODO: implement this kwarg
    ):
    return multicall_same_func_different_contracts_same_input(token_addresses, 'balanceOf(address)(uint)', input=hodler_address, block=block)


@log(logger)
def fetch_multicall(*calls, block=None):
    # 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


@log(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)


@log(logger)
def batch_call(calls):
    """
    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'))
    ]


@log(logger)
def _clean_addresses(
    addresses: Sequence[Any]
):

    return [_clean_address(address) for address in addresses]


@log(logger)
def _clean_address(
    address: Any
    ):

    if type(address) not in SUPPORTED_INPUT_TYPES:
        raise TypeError(f'Unsupported input type: {type(address)}')
    
    if type(address) != str: 
        address = str(address)

    return convert.to_address(address)
