from collections import defaultdict
from itertools import count, product
from operator import itemgetter
from typing import List

import requests
from brownie import chain, web3
from eth_abi.exceptions import InsufficientDataBytes
from web3.exceptions import CannotHandleRequest
from ypricemagic.utils.contracts import Contract, contract_creation_block
from ypricemagic.utils.raw_calls import _decimals, _totalSupply

from multicall import Call, Multicall

# the address doesn't matter, it just needs to have code
code = " "
multicall, multicall2 = None, None
if chain.id == 1:
    multicall2 = Contract("0x5BA1e12693Dc8F9c48aAD8770482f4739bEeD696")
if chain.id == 56:
    multicall2 = Contract("0xfF6FD90A470Aaa0c1B8A54681746b07AcdFedc9B")
if chain.id == 137:
    multicall2 = Contract("0xc8E51042792d7405184DfCa245F2d27B94D013b6")
if chain.id == 250:
    multicall2 = Contract("0xBAD2B082e2212DE4B065F636CA4e5e0717623d18")

multicall_deploy_block = contract_creation_block(multicall2.address)

def multicall_same_func_no_input(addresses: list, method: str, apply_func=None, block=None):
    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, apply_func=None, block=None):
    assert input
    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: list, method: str, inputs: List=None, apply_func=None, block=None):
    assert inputs
    calls = [Call(address, [method, input], [[input,apply_func]]) for input in inputs]
    return [result for result in Multicall(calls, block_id=block, _w3=web3)().values()]

def multicall_decimals(addresses: List[str], block=None):
    try: return multicall_same_func_no_input(addresses, 'decimals()(uint256)', block=block)
    except ValueError as e:
        if 'execution reverted' in str(e): return [_decimals(address,block,return_None_on_failure=True) for address in addresses]
        else: raise
    except (CannotHandleRequest,InsufficientDataBytes): return [_decimals(address,block,return_None_on_failure=True) for address in addresses] 

def multicall_totalSupply(addresses: List[str], block=None):
    try: return multicall_same_func_no_input(addresses, 'totalSupply()(uint256)', block=block)
    except (CannotHandleRequest,InsufficientDataBytes): return [_totalSupply(address,block) for address in addresses] 

def multicall_balanceOf(token_addresses: List[str], hodler_address: str, block=None):
    return multicall_same_func_different_contracts_same_input(token_addresses, 'balanceOf(address)(uint)', input=hodler_address)

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