from collections import defaultdict
from dataclasses import dataclass
from threading import Lock, get_ident
from types import FunctionType, TracebackType
from typing import Any

from lazy_object_proxy import Proxy
from wrapt import ObjectProxy

from brownie._c_constants import ujson_loads
from brownie._config import BROWNIE_FOLDER, CONFIG
from brownie.exceptions import ContractNotFound
from brownie.network import accounts, web3
from brownie.network.contract import Contract, ContractCall
from brownie.project import compile_source
from brownie.utils import color

DATA_DIR = BROWNIE_FOLDER.joinpath("data")
MULTICALL2_ABI = ujson_loads(DATA_DIR.joinpath("interfaces", "Multicall2.json").read_text())
MULTICALL2_SOURCE = DATA_DIR.joinpath("contracts", "Multicall2.sol").read_text()


@dataclass
class Call:

    calldata: tuple[str, bytes]
    decoder: FunctionType
    readable: str


class Result(ObjectProxy):
    """A proxy object to be updated with the result of a multicall."""

    def __repr__(self) -> str:
        return repr(self.__wrapped__)


class LazyResult(Proxy):
    """A proxy object to be updated with the result of a multicall."""

    def __repr__(self) -> str:
        return repr(self.__wrapped__)


class Multicall:
    """Context manager for batching multiple calls to constant contract functions."""

    _lock = Lock()

    def __init__(self) -> None:
        self.address = None
        self.default_verbose = False
        self._block_number = defaultdict(lambda: None)
        self._verbose = defaultdict(lambda: None)
        self._contract = None
        self._pending_calls: dict[int, list[Call]] = defaultdict(list)

        setattr(ContractCall, "__original_call_code", ContractCall.__call__.__code__)
        setattr(ContractCall, "__proxy_call_code", self._proxy_call.__code__)
        setattr(ContractCall, "__multicall", defaultdict(lambda: None))
        ContractCall.__call__.__code__ = self._proxy_call.__code__

    @property
    def block_number(self) -> int:
        return self._block_number[get_ident()]

    def __call__(
        self,
        address: str | None = None,
        block_identifier: str | bytes | int | None = None,
        verbose: bool | None = None,
    ) -> "Multicall":
        self.address = address
        self._block_number[get_ident()] = block_identifier
        self._verbose[get_ident()] = verbose if verbose is not None else self.default_verbose
        return self

    def _flush(self, future_result: Result = None) -> Any:
        pending_calls = self._pending_calls[get_ident()]
        self._pending_calls[get_ident()] = []

        if not pending_calls:
            # either all calls have already been made
            # or this result has already been retrieved
            return future_result
        with self._lock:
            if self._verbose.get(get_ident(), self.default_verbose):
                message = (
                    "Multicall:"
                    f"\n  Thread ID: {get_ident()}"
                    f"\n  Block number: {self._block_number[get_ident()]}"
                    f"\n  Calls: {len(pending_calls)}"
                )
                for c, item in enumerate(pending_calls, start=1):
                    u = "\u2514" if c == len(pending_calls) else "\u251c"
                    message = f"{message}\n    {u}\u2500{item.readable}"
                print(color.highlight(f"{message}\n"))

            ContractCall.__call__.__code__ = getattr(ContractCall, "__original_call_code")
            try:
                results = self._contract.tryAggregate(
                    False,
                    [_call.calldata for _call in pending_calls],
                    block_identifier=self._block_number[get_ident()],
                )
            finally:
                ContractCall.__call__.__code__ = getattr(ContractCall, "__proxy_call_code")

        for _call, result in zip(pending_calls, results):
            _call.__wrapped__ = _call.decoder(result[1]) if result[0] else None

        return future_result

    def flush(self) -> Any:
        """Flush the pending queue of calls, retrieving all the results."""
        return self._flush()

    def _call_contract(self, call: ContractCall, *args: Any, **kwargs: Any) -> Proxy:
        """Add a call to the buffer of calls to be made"""
        calldata = (call._address, call.encode_input(*args, **kwargs))
        readable = f"{call._name}({', '.join(str(i) for i in args)})"
        call_obj = Call(calldata, call.decode_output, readable)
        # future result
        result = Result(call_obj)
        self._pending_calls[get_ident()].append(result)

        return LazyResult(lambda: self._flush(result))

    @staticmethod
    def _proxy_call(*args: Any, **kwargs: Any) -> Any:
        """Proxy code which substitutes `ContractCall.__call__"""
        if self := getattr(ContractCall, "__multicall", {}).get(get_ident()):
            return self._call_contract(*args, **kwargs)

        # standard call we let pass through
        ContractCall.__call__.__code__ = getattr(ContractCall, "__original_call_code")
        try:
            result = ContractCall.__call__(*args, **kwargs)
        finally:
            ContractCall.__call__.__code__ = getattr(ContractCall, "__proxy_call_code")
        return result

    def __enter__(self) -> "Multicall":
        """Enter the Context Manager and substitute `ContractCall.__call__`"""
        # we set the code objects on ContractCall class so we can grab them later

        active_network = CONFIG.active_network

        if "multicall2" in active_network:
            self.address = active_network["multicall2"]
        elif "cmd" in active_network:
            deployment = self.deploy({"from": accounts[0]})
            self.address = deployment.address
            self._block_number[get_ident()] = deployment.tx.block_number

        self._block_number[get_ident()] = (
            self._block_number[get_ident()] or web3.eth.get_block_number()
        )

        if self.address is None:
            raise ContractNotFound(
                "Must set Multicall address via `brownie.multicall(address=...)`"
            )
        elif not web3.eth.get_code(self.address, block_identifier=self.block_number):
            raise ContractNotFound(
                f"Multicall at address {self.address} does not exist at block {self.block_number}"
            )

        self._contract = Contract.from_abi("Multicall", self.address, MULTICALL2_ABI)
        getattr(ContractCall, "__multicall")[get_ident()] = self

    def __exit__(self, exc_type: Exception, exc_val: Any, exc_tb: TracebackType) -> None:
        """Exit the Context Manager and reattach original `ContractCall.__call__` code"""
        self.flush()
        getattr(ContractCall, "__multicall")[get_ident()] = None

        self._block_number.pop(get_ident(), None)
        self._verbose.pop(get_ident(), None)

    @staticmethod
    def deploy(tx_params: dict) -> Contract:
        """Deploy an instance of the `Multicall2` contract.

        Args:
            tx_params: parameters passed to the `deploy` method of the `Multicall2` contract
                container.
        """
        project = compile_source(MULTICALL2_SOURCE)
        deployment = project.Multicall2.deploy(tx_params)
        CONFIG.active_network["multicall2"] = deployment.address
        return deployment
