"""Contract interaction layer: ABI parsing, function/event encoding/decoding,
and Multicall3 helpers.
"""

from dataclasses import dataclass, field
from typing import Any

from eth_abi import decode as abi_decode, encode as abi_encode
from eth_abi.exceptions import DecodingError
from eth_utils import keccak
from loguru import logger


# === ABI descriptors ===

@dataclass(frozen=True)
class ABIFunction:
    """Immutable, hashable descriptor for a single ABI function entry.

    The selector is computed from name + inputs in __post_init__ and stored
    as a plain attribute, avoiding the @property/@lru_cache anti-pattern on
    frozen dataclasses.
    """

    name: str
    inputs: tuple[str, ...]
    outputs: tuple[str, ...]
    state_mutability: str = "nonpayable"
    selector: str = field(init=False, repr=False, compare=False, hash=False)

    def __post_init__(self) -> None:
        signature = f"{self.name}({','.join(self.inputs)})"
        object.__setattr__(self, "selector", "0x" + keccak(text=signature).hex()[:8])

    @property
    def signature(self) -> str:
        """Full function signature, e.g. 'transfer(address,uint256)'."""
        return f"{self.name}({','.join(self.inputs)})"

    def encode_input(self, *args: Any) -> bytes:
        """ABI-encode function arguments into raw bytes.

        Handles hex-string → bytes conversion for bytes<M> parameter types.

        Raises:
            ValueError: If the number of arguments does not match the ABI.
        """
        if len(args) != len(self.inputs):
            raise ValueError(
                f"{self.name} expects {len(self.inputs)} args, got {len(args)}"
            )
        processed: list[Any] = []
        for typ, arg in zip(self.inputs, args):
            if typ.startswith("bytes") and isinstance(arg, str):
                processed.append(
                    bytes.fromhex(arg[2:]) if arg.startswith("0x") else arg.encode()
                )
            else:
                processed.append(arg)
        return abi_encode(list(self.inputs), processed)

    def decode_output(self, data: bytes) -> tuple[Any, ...]:
        """ABI-decode raw response bytes into a tuple of return values."""
        if not self.outputs:
            return ()
        return abi_decode(list(self.outputs), data)


@dataclass(frozen=True)
class ABIEvent:
    """Immutable, hashable descriptor for a single ABI event entry.

    topic0 is computed in __post_init__ and stored as a plain attribute.
    """

    name: str
    inputs: tuple[tuple[str, str, bool], ...]  # (field_name, abi_type, indexed)
    anonymous: bool = False
    topic0: str = field(init=False, repr=False, compare=False, hash=False)

    def __post_init__(self) -> None:
        if self.anonymous:
            object.__setattr__(self, "topic0", "")
        else:
            types = [inp[1] for inp in self.inputs]
            sig = f"{self.name}({','.join(types)})"
            object.__setattr__(self, "topic0", "0x" + keccak(text=sig).hex())

    @property
    def signature(self) -> str:
        """Full event signature, e.g. 'Transfer(address,address,uint256)'."""
        types = [inp[1] for inp in self.inputs]
        return f"{self.name}({','.join(types)})"

    @property
    def indexed_inputs(self) -> list[tuple[str, str]]:
        """Return (name, type) pairs for all indexed parameters."""
        return [(name, typ) for name, typ, indexed in self.inputs if indexed]

    @property
    def data_inputs(self) -> list[tuple[str, str]]:
        """Return (name, type) pairs for all non-indexed parameters."""
        return [(name, typ) for name, typ, indexed in self.inputs if not indexed]

    def decode_log(self, topics: list[str], data: str) -> dict[str, Any]:
        """Decode a raw event log into a named-argument dict.

        Raises:
            ValueError: If the non-indexed data cannot be ABI-decoded.
        """
        result: dict[str, Any] = {}

        topic_start = 0 if self.anonymous else 1
        for i, (name, typ) in enumerate(self.indexed_inputs):
            if topic_start + i >= len(topics):
                break
            topic_bytes = bytes.fromhex(topics[topic_start + i][2:])
            result[name] = self._decode_topic(typ, topic_bytes)

        if self.data_inputs:
            data_bytes = bytes.fromhex(data[2:] if data.startswith("0x") else data)
            types = [typ for _, typ in self.data_inputs]
            try:
                decoded = abi_decode(types, data_bytes)
                for (name, _), value in zip(self.data_inputs, decoded):
                    result[name] = value
            except DecodingError as exc:
                raise ValueError(f"Failed to decode {self.name} event data: {exc}") from exc

        return result

    @staticmethod
    def _decode_topic(typ: str, data: bytes) -> Any:
        """Decode a single 32-byte indexed topic value to the appropriate Python type."""
        if typ == "address":
            return "0x" + data[-20:].hex()
        if typ == "bool":
            return bool(int.from_bytes(data, "big"))
        if typ.startswith("uint"):
            return int.from_bytes(data, "big")
        if typ.startswith("int"):
            return int.from_bytes(data, "big", signed=True)
        # bytes<M> and complex reference types: return as hex
        return "0x" + data.hex()


# === Contract ===

class Contract:
    """High-performance interface to an on-chain smart contract.

    Parses the raw ABI dict list on first access and caches all
    functions/events for fast repeated lookups.

    Usage::

        erc20 = Contract("0xA0b...eB48", ERC20_ABI)
        calldata = erc20.encode_function_data("transfer", recipient, amount)
        decoded = erc20.decode_function_result("balanceOf", raw_hex)
    """

    def __init__(self, address: str, abi: list[dict[str, Any]]) -> None:
        self.address: str = address.lower()
        self._abi = abi
        self._functions: dict[str, ABIFunction] | None = None
        self._events: dict[str, ABIEvent] | None = None
        self._event_by_topic: dict[str, ABIEvent] | None = None

    @property
    def functions(self) -> dict[str, ABIFunction]:
        """All functions from the ABI, keyed by name (lazy, cached)."""
        if self._functions is None:
            self._functions = {}
            for entry in self._abi:
                if entry.get("type") == "function" and "name" in entry:
                    func = ABIFunction(
                        name=entry["name"],
                        inputs=tuple(inp["type"] for inp in entry.get("inputs", [])),
                        outputs=tuple(out["type"] for out in entry.get("outputs", [])),
                        state_mutability=entry.get("stateMutability", "nonpayable"),
                    )
                    self._functions[func.name] = func
        return self._functions

    @property
    def events(self) -> dict[str, ABIEvent]:
        """All events from the ABI, keyed by name (lazy, cached)."""
        if self._events is None:
            self._events = {}
            for entry in self._abi:
                if entry.get("type") == "event" and "name" in entry:
                    event = ABIEvent(
                        name=entry["name"],
                        inputs=tuple(
                            (
                                inp.get("name", f"arg{i}"),
                                inp["type"],
                                inp.get("indexed", False),
                            )
                            for i, inp in enumerate(entry.get("inputs", []))
                        ),
                        anonymous=entry.get("anonymous", False),
                    )
                    self._events[event.name] = event
        return self._events

    @property
    def events_by_topic(self) -> dict[str, ABIEvent]:
        """All events indexed by topic0 (lazy, cached)."""
        if self._event_by_topic is None:
            self._event_by_topic = {
                event.topic0: event
                for event in self.events.values()
                if event.topic0
            }
        return self._event_by_topic

    def get_function(self, name: str) -> ABIFunction:
        """Return function descriptor by name.

        Raises:
            ValueError: If no function with this name exists in the ABI.
        """
        func = self.functions.get(name)
        if func is None:
            raise ValueError(f"Function '{name}' not found in contract {self.address}")
        return func

    def get_event(self, name: str) -> ABIEvent:
        """Return event descriptor by name.

        Raises:
            ValueError: If no event with this name exists in the ABI.
        """
        event = self.events.get(name)
        if event is None:
            raise ValueError(f"Event '{name}' not found in contract {self.address}")
        return event

    def get_event_by_topic(self, topic0: str) -> "ABIEvent | None":
        """Return event descriptor matching the given topic0, or None."""
        return self.events_by_topic.get(topic0.lower())

    def encode_function_data(self, function_name: str, *args: Any) -> str:
        """Return hex calldata (4-byte selector + ABI-encoded args) for the given function."""
        func = self.get_function(function_name)
        encoded = func.encode_input(*args)
        return func.selector + encoded.hex()

    def decode_function_result(self, function_name: str, data: str) -> Any:
        """ABI-decode the raw hex result of an eth_call for the given function.

        Returns:
            The single return value directly, a tuple for multi-value functions,
            or None for void functions.
        """
        func = self.get_function(function_name)
        data_bytes = bytes.fromhex(data[2:] if data.startswith("0x") else data)
        decoded = func.decode_output(data_bytes)
        if not decoded:
            return None
        return decoded[0] if len(decoded) == 1 else decoded

    def decode_log(self, topics: list[str], data: str) -> "dict[str, Any] | None":
        """Decode a raw event log into a named-argument dict.

        Returns None if the topic0 is not recognised by this contract's ABI.
        The returned dict always includes a ``'_event'`` key with the event name.
        """
        if not topics:
            return None
        event = self.get_event_by_topic(topics[0])
        if event is None:
            return None
        decoded = event.decode_log(topics, data)
        decoded["_event"] = event.name
        return decoded

    def has_function(self, name: str) -> bool:
        """Return True if a function with this name exists in the ABI."""
        return name in self.functions

    def has_event(self, name: str) -> bool:
        """Return True if an event with this name exists in the ABI."""
        return name in self.events

    def list_functions(self) -> list[str]:
        """Return the names of all functions defined in the ABI."""
        return list(self.functions.keys())

    def list_events(self) -> list[str]:
        """Return the names of all events defined in the ABI."""
        return list(self.events.keys())


# === Multicall helpers ===

def build_try_aggregate_call(
    calls: list[tuple[str, Contract, str, tuple[Any, ...]]],
    require_success: bool = False,
) -> str:
    """Build calldata for Multicall3.tryAggregate(bool,(address,bytes)[]).

    Args:
        calls: List of (target_address, contract, function_name, args).
        require_success: Forwarded to tryAggregate — if True, the on-chain
            call reverts when any individual call fails.

    Returns:
        '0x'-prefixed hex calldata string.
    """
    call_tuples: list[tuple[str, bytes]] = [
        (address, bytes.fromhex(contract.encode_function_data(fn, *args)[2:]))
        for address, contract, fn, args in calls
    ]
    encoded = abi_encode(["bool", "(address,bytes)[]"], [require_success, call_tuples])
    return "0xbce38bd7" + encoded.hex()  # tryAggregate(bool,(address,bytes)[])


def decode_results(
    calls: list[tuple[str, Contract, str, tuple[Any, ...]]],
    return_data: list[bytes],
) -> dict[str, Any]:
    """Decode raw return_data from a Multicall3 response.

    Args:
        calls: Original call list passed to build_try_aggregate_call.
        return_data: Raw bytes list, one entry per call.

    Returns:
        Dict keyed by contract address. When multiple calls target the same
        address, later results overwrite earlier ones — use
        MulticallHelper.execute_multicall with explicit result_key for
        uniqueness.
    """
    results: dict[str, Any] = {}
    for i, (address, contract, fn_name, _) in enumerate(calls):
        if i >= len(return_data):
            results[address] = None
            continue
        try:
            results[address] = contract.decode_function_result(
                fn_name, "0x" + return_data[i].hex()
            )
        except Exception as exc:
            logger.warning(f"Failed to decode result for {address}.{fn_name}: {exc}")
            results[address] = None
    return results


# Type alias for the full call entry used by MulticallHelper
_CallEntry = tuple[str, Any, str, tuple[Any, ...], "str | None"]


class MulticallHelper:
    """Batching helper for Multicall3.tryAggregate.

    Splits large call lists into fixed-size batches, issues one eth_call per
    batch, and reassembles decoded results under caller-supplied keys.

    Usage::

        helper = MulticallHelper(rpc_client)
        results = await helper.execute_multicall([
            (token_addr, token_contract, "balanceOf", (wallet,), "balance"),
            (pair_addr,  pair_contract,  "getReserves", (),        "reserves"),
        ])
    """

    def __init__(
        self,
        rpc_client: Any,
        multicall_address: str = "0xcA11bde05977b3631167028862bE2a173976CA11",
    ) -> None:
        self.rpc_client = rpc_client
        self.multicall_address = multicall_address

    async def execute_multicall(
        self,
        calls: list[_CallEntry],
        batch_size: int = 100,
    ) -> dict[str, Any]:
        """Execute multicall requests with automatic batching.

        Args:
            calls: List of
                (contract_address, contract, function_name, args, result_key).
                result_key is an optional unique identifier for the result
                dict; falls back to contract_address when None.
            batch_size: Maximum number of calls per on-chain multicall.

        Returns:
            Dict mapping result_key → decoded value (None on failure).
        """
        results: dict[str, Any] = {}

        for batch_start in range(0, len(calls), batch_size):
            batch = calls[batch_start : batch_start + batch_size]

            multicall_calls: list[tuple[str, Any, str, tuple[Any, ...]]] = [
                (contract_address, contract, function_name, args)
                for contract_address, contract, function_name, args, _ in batch
            ]

            calldata = build_try_aggregate_call(multicall_calls)
            raw = await self.rpc_client.call(self.multicall_address, calldata)

            raw_bytes = bytes.fromhex(raw[2:] if raw.startswith("0x") else raw)
            decoded_raw = abi_decode(["(bool,bytes)[]"], raw_bytes)[0]

            for j, (contract_address, contract, function_name, _, result_key) in enumerate(batch):
                key = result_key if result_key is not None else contract_address

                if j >= len(decoded_raw):
                    results[key] = None
                    continue

                success, return_data = decoded_raw[j]
                if not success or not return_data:
                    results[key] = None
                    continue

                try:
                    results[key] = contract.decode_function_result(
                        function_name,
                        "0x" + return_data.hex(),
                    )
                except Exception as exc:
                    logger.warning(f"Failed to decode multicall result for '{key}': {exc}")
                    results[key] = None

        return results

    async def call_fn_on_many(
        self,
        addresses: list[str],
        abi: list[dict[str, Any]],
        function_name: str,
        batch_size: int = 100,
    ) -> dict[str, Any]:
        """Call the same zero-argument function on many contracts in one or more Multicall3 batches.

        Args:
            addresses: List of contract addresses to call.
            abi: ABI containing at least the target function definition.
            function_name: Name of the function to call (must take no arguments).
            batch_size: Maximum number of calls per on-chain multicall.

        Returns:
            Dict mapping address → decoded return value. Addresses where the
            call failed or returned empty data are omitted.
        """
        calls: list[_CallEntry] = [
            (addr, Contract(addr, abi), function_name, (), addr)
            for addr in addresses
        ]
        results = await self.execute_multicall(calls, batch_size=batch_size)
        return {k: v for k, v in results.items() if v is not None}
