import asyncio
from time import time
from typing import Any, Dict, List, Optional, Tuple, Union

import aiohttp
import requests
from web3 import Web3

from mucalla import Call
from mucalla.constants import (
    GAS_LIMIT,
    MULTICALL2_ADDRESSES,
    MULTICALL2_BYTECODE,
    MULTICALL_ADDRESSES,
    w3,
)
from mucalla.loggers import setup_logger
from mucalla.utils import (
    await_awaitable,
    chain_id,
    gather,
    run_in_subprocess,
    state_override_supported,
)

logger = setup_logger(__name__)

CallResponse = Tuple[Union[None, bool], bytes]


def get_args(
    calls: List[Call], require_success: bool = True
) -> List[Union[bool, List[List[Any]]]]:
    if require_success is True:
        return [[[call.target, call.data] for call in calls]]
    return [require_success, [[call.target, call.data] for call in calls]]


def unpack_aggregate_outputs(outputs: Any) -> Tuple[CallResponse, ...]:
    return tuple((None, output) for output in outputs)


def unpack_batch_results(batch_results: List[List[CallResponse]]) -> List[CallResponse]:
    return [result for batch in batch_results for result in batch]


class Multicall:
    def __init__(
        self,
        calls: List[Call],
        block_id: Optional[int] = None,
        require_success: bool = True,
        gas_limit: int = GAS_LIMIT,
        _w3: Web3 = w3,
    ) -> None:
        self.calls = calls
        self.block_id = block_id
        self.require_success = require_success
        self.gas_limit = gas_limit
        self.w3 = _w3
        self.chainid = chain_id(self.w3)
        if require_success is True:
            multicall_map = (
                MULTICALL_ADDRESSES
                if self.chainid in MULTICALL_ADDRESSES
                else MULTICALL2_ADDRESSES
            )
            self.multicall_sig = "aggregate((address,bytes)[])(uint256,bytes[])"
        else:
            multicall_map = MULTICALL2_ADDRESSES
            self.multicall_sig = "tryBlockAndAggregate(bool,(address,bytes)[])(uint256,uint256,(bool,bytes)[])"
        self.multicall_address = multicall_map[self.chainid]

    def __call__(self) -> Dict[str, Any]:
        start = time()
        response = await_awaitable(self.coroutine())
        logger.debug(f"Multicall took {time() - start}s")
        return response

    async def coroutine(self) -> Dict[str, Any]:
        batches = await gather(
            [
                self.fetch_outputs(batch, id=str(i))
                for i, batch in enumerate(batcher.batch_calls(self.calls, batcher.step))
            ]
        )
        outputs = await run_in_subprocess(unpack_batch_results, batches)

        return {name: result for output in outputs for name, result in output.items()}

    async def fetch_outputs(
        self, calls: List[Call], ConnErr_retries: int = 0, id: str = ""
    ) -> List[CallResponse]:
        logger.debug(f"coroutine {id} started")

        if calls is None:
            calls = self.calls

        try:
            args = await run_in_subprocess(get_args, calls, self.require_success)
            if self.require_success is True:
                block_number, outputs = await self.aggregate.coroutine(args)
                outputs = await run_in_subprocess(unpack_aggregate_outputs, outputs)
            else:
                _, _, outputs = await self.aggregate.coroutine(args)
            outputs = await gather(
                [
                    run_in_subprocess(
                        Call.decode_output,
                        output,
                        call.signature,
                        call.returns,
                        success,
                    )
                    for call, (success, output) in zip(calls, outputs)
                ]
            )
            logger.debug(f"coroutine {id} finished")

            outputs.append({"block_number": block_number})
            return outputs
        except Exception as e:
            _raise_or_proceed(e, len(calls), ConnErr_retries=ConnErr_retries)

        # Failed, we need to rebatch the calls and try again.
        batch_results = await gather(
            [
                self.fetch_outputs(chunk, ConnErr_retries + 1, f"{id}_{i}")
                for i, chunk in enumerate(await batcher.rebatch(calls))
            ]
        )

        return_val = await run_in_subprocess(unpack_batch_results, batch_results)
        logger.debug(f"coroutine {id} finished")
        return return_val

    @property
    def aggregate(self) -> Call:
        if state_override_supported(self.w3):
            return Call(
                self.multicall_address,
                self.multicall_sig,
                returns=None,
                _w3=self.w3,
                block_id=self.block_id,
                gas_limit=self.gas_limit,
                state_override_code=MULTICALL2_BYTECODE,
            )

        # If state override is not supported, we simply skip it.
        # This will mean you're unable to access full historical data on chains without state override support.
        return Call(
            self.multicall_address,
            self.multicall_sig,
            returns=None,
            _w3=self.w3,
            block_id=self.block_id,
            gas_limit=self.gas_limit,
        )


class NotSoBrightBatcher:
    """
    This class helps with processing a large volume of large multicalls.
    It's not so bright, but should quickly bring the batch size down to something reasonable for your node.
    """

    def __init__(self) -> None:
        self.step = 10000

    def batch_calls(self, calls: List[Call], step: int) -> List[List[Call]]:
        """
        Batch calls into chunks of size `self.step`.
        """
        batches = []
        start = 0
        done = len(calls) - 1
        while True:
            end = start + step
            batches.append(calls[start:end])
            if end >= done:
                return batches
            start = end

    def split_calls(
        self, calls: List[Call], unused: None = None
    ) -> Tuple[List[Call], List[Call]]:
        """
        Split calls into 2 batches in case request is too large.
        We do this to help us find optimal `self.step` value.
        """
        center = len(calls) // 2
        chunk_1 = calls[:center]
        chunk_2 = calls[center:]
        return chunk_1, chunk_2

    async def rebatch(self, calls):
        # If a separate coroutine changed `step` after calls were last batched, we will use the new `step` for rebatching.
        if self.step <= len(calls) // 2:
            return await run_in_subprocess(self.batch_calls, calls, self.step)

        # Otherwise we will split calls in half.
        if self.step >= len(calls):
            new_step = round(len(calls) * 0.99) if len(calls) >= 100 else len(calls) - 1
            logger.warning(
                f"Multicall batch size reduced from {self.step} to {new_step}. The failed batch had {len(calls)} calls."
            )
            self.step = new_step
        return await run_in_subprocess(self.split_calls, calls, self.step)


batcher = NotSoBrightBatcher()


def _raise_or_proceed(e: Exception, ct_calls: int, ConnErr_retries: int) -> None:
    """Depending on the exception, either raises or ignores and allows `batcher` to rebatch."""
    if isinstance(e, aiohttp.ClientOSError):
        if "broken pipe" not in str(e).lower():
            raise e
        logger.warning(e)
    elif isinstance(e, aiohttp.ClientResponseError):
        strings = ["request entity too large", "connection reset by peer"]
        if not any([string in str(e).lower() for string in strings]):
            raise e
        logger.warning(e)
    elif isinstance(e, requests.ConnectionError):
        if (
            "('Connection aborted.', ConnectionResetError(104, 'Connection reset by peer'))"
            not in str(e)
            or ConnErr_retries > 5
        ):
            raise e
    elif isinstance(e, requests.HTTPError):
        strings = (
            "request entity too large",
            "payload too large",
            "time-out",
            "520 server error",
        )
        if not any([string in str(e).lower() for string in strings]):
            raise e
        logger.warning(e)
    elif isinstance(e, asyncio.TimeoutError):
        pass
    elif isinstance(e, ValueError):
        if "out of gas" not in str(e).lower():
            raise e
        if ct_calls == 1:
            raise e
        logger.warning(e)
    else:
        raise e
