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.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

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",
}.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)


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()]


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()]


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()]


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]

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] 


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)


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


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)


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'))
    ]


def _clean_addresses(
    addresses: Sequence[Any]
):

    return [_clean_address(address) for address in addresses]


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)
