from collections.abc import Iterator
from functools import singledispatchmethod

from ape.api.query import ContractCreation, ContractCreationQuery, QueryAPI, QueryType
from ape.exceptions import APINotImplementedError, ProviderError, QueryEngineError
from ape.types.address import AddressType


class EthereumQueryProvider(QueryAPI):
    """
    Implements more advanced queries specific to Ethereum clients.
    """

    def __init__(self):
        self.supports_contract_creation = None  # will be set after we try for the first time

    @singledispatchmethod
    def estimate_query(self, query: QueryType) -> int | None:  # type: ignore[override]
        return None

    @singledispatchmethod
    def perform_query(self, query: QueryType) -> Iterator:  # type: ignore[override]
        raise QueryEngineError(f"Cannot handle '{type(query)}'.")

    @estimate_query.register
    def estimate_contract_creation_query(self, query: ContractCreationQuery) -> int | None:
        # NOTE: Extremely expensive query, involves binary search of all blocks in a chain
        #       Very loose estimate of 5s per transaction for this query.
        if self.supports_contract_creation is False:
            return None
        return 5000

    @perform_query.register
    def perform_contract_creation_query(
        self, query: ContractCreationQuery
    ) -> Iterator[ContractCreation]:
        """
        Find when a contract was deployed using binary search and block tracing.
        """
        # skip the search if there is still no code at address at head
        if not self.chain_manager.get_code(query.contract):
            return None

        def find_creation_block(lo, hi):
            # perform a binary search to find the block when the contract was deployed.
            # takes log2(height), doesn't work with contracts that have been reinit.
            while hi - lo > 1:
                mid = (lo + hi) // 2
                code = self.chain_manager.get_code(query.contract, block_id=mid)
                if not code:
                    lo = mid
                else:
                    hi = mid

            if self.chain_manager.get_code(query.contract, block_id=hi):
                return hi

            return None

        try:
            block = find_creation_block(0, self.chain_manager.blocks.height)
        except ProviderError:
            self.supports_contract_creation = False
            return None

        # iterate over block transaction traces to find the deployment call
        # this method also supports contracts created by factories
        try:
            if "geth" in self.provider.client_version.lower():
                yield from self._find_creation_in_block_via_geth(block, query.contract)
            else:
                yield from self._find_creation_in_block_via_parity(block, query.contract)
        except (ProviderError, APINotImplementedError):
            self.supports_contract_creation = False
            return None

        self.supports_contract_creation = True

    def _find_creation_in_block_via_parity(self, block, contract_address):
        # NOTE requires `trace_` namespace
        traces = self.provider.make_request("trace_replayBlockTransactions", [block, ["trace"]])

        for tx in traces:
            for trace in tx["trace"]:
                if (
                    "error" not in trace
                    and trace["type"] == "create"
                    and trace["result"]["address"] == contract_address.lower()
                ):
                    receipt = self.chain_manager.get_receipt(tx["transactionHash"])
                    creator = self.conversion_manager.convert(trace["action"]["from"], AddressType)
                    yield ContractCreation(
                        txn_hash=tx["transactionHash"],
                        block=block,
                        deployer=receipt.sender,
                        factory=creator if creator != receipt.sender else None,
                    )

    def _find_creation_in_block_via_geth(self, block, contract_address):
        # NOTE requires `debug_` namespace
        traces = self.provider.make_request(
            "debug_traceBlockByNumber", [hex(block), {"tracer": "callTracer"}]
        )

        def flatten(call):
            if call["type"] in ["CREATE", "CREATE2"]:
                yield call["from"], call["to"]

            if "error" in call or "calls" not in call:
                return

            for sub in call["calls"]:
                if sub["type"] in ["CREATE", "CREATE2"]:
                    yield sub["from"], sub["to"]
                else:
                    yield from flatten(sub)

        for tx in traces:
            call = tx["result"]
            sender = call["from"]
            for factory, contract in flatten(call):
                if contract == contract_address.lower():
                    yield ContractCreation(
                        txn_hash=tx["txHash"],
                        block=block,
                        deployer=self.conversion_manager.convert(sender, AddressType),
                        factory=(
                            self.conversion_manager.convert(factory, AddressType)
                            if factory != sender
                            else None
                        ),
                    )
