"""Optimised async JSON-RPC client for Ethereum-compatible blockchains.

Features: connection pooling, request batching, exponential-backoff retries,
TTL result cache.
"""

import asyncio
import itertools
import random
from collections.abc import Callable
from dataclasses import dataclass, field
from functools import lru_cache
from typing import Any

import aiohttp
import orjson
from aiohttp import ClientTimeout
from cachetools import TTLCache
from eth_abi import decode as abi_decode, encode as abi_encode
from eth_utils import function_signature_to_4byte_selector, keccak
from loguru import logger

# === Request ID counter ===

_request_id: itertools.count = itertools.count(1)


# === Data types ===

@dataclass
class RPCConfig:
    """Configuration for RPCClient.

    Attributes:
        rpc_url: HTTP(S) endpoint of the JSON-RPC node.
        max_connections: Total TCP connection pool size.
        connections_per_host: Per-host TCP connection limit.
        connect_timeout: Seconds to wait for a new TCP connection.
        read_timeout: Seconds to wait for a response body.
        max_retries: Number of retry attempts on network failure.
        retry_base_delay: Base delay in seconds for exponential back-off.
        retry_jitter: Maximum random jitter added to each retry delay.
        enable_cache: Enable TTL result cache.
        cache_ttl: Cache entry lifetime in seconds.
        cache_maxsize: Maximum number of cached entries.
    """

    rpc_url: str
    max_connections: int = 200
    connections_per_host: int = 10
    connect_timeout: int = 10
    read_timeout: int = 60
    max_retries: int = 5
    retry_base_delay: float = 0.15
    retry_jitter: float = 0.15
    enable_cache: bool = True
    cache_ttl: int = 300
    cache_maxsize: int = 2048


@dataclass
class CallRequest:
    """Single JSON-RPC call descriptor used when building batch requests."""

    method: str
    params: list[Any]
    id: int = field(default_factory=lambda: next(_request_id))


class RPCError(Exception):
    """Raised when the RPC node returns an error object in its response."""

    def __init__(self, code: int, message: str, data: Any = None) -> None:
        self.code = code
        self.message = message
        self.data = data
        super().__init__(f"RPC Error {code}: {message}")


# === Client ===

class RPCClient:
    """High-performance async JSON-RPC client.

    Supports automatic retries with exponential back-off, request batching,
    TCP connection pooling, and a TTL result cache.

    Usage::

        async with RPCClient("https://rpc.example.com") as client:
            block = await client.get_block_number()
            balance = await client.get_balance("0xAbc...")
    """

    def __init__(self, config: str | RPCConfig) -> None:
        if isinstance(config, str):
            config = RPCConfig(rpc_url=config)

        self.config = config
        self._session: aiohttp.ClientSession | None = None
        self.chain_id: int | None = None
        self._cache: TTLCache = TTLCache(
            maxsize=config.cache_maxsize,
            ttl=config.cache_ttl,
        )
        self._lock = asyncio.Lock()

        self._connector = aiohttp.TCPConnector(
            limit=config.max_connections,
            limit_per_host=config.connections_per_host,
            enable_cleanup_closed=True,
            ttl_dns_cache=config.cache_ttl,
        )
        self._timeout = ClientTimeout(
            total=None,
            connect=config.connect_timeout,
            sock_connect=config.connect_timeout,
            sock_read=config.read_timeout,
        )
        self._headers = {
            "Accept": "application/json",
            "Accept-Encoding": "gzip",
            "Connection": "keep-alive",
            "Content-Type": "application/json",
        }

    async def __aenter__(self) -> "RPCClient":
        await self._ensure_session()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()

    # === Session management ===

    async def _ensure_session(self) -> aiohttp.ClientSession:
        """Return the active session, creating one if needed (double-checked lock)."""
        if self._session is None or self._session.closed:
            async with self._lock:
                if self._session is None or self._session.closed:
                    self._session = aiohttp.ClientSession(
                        connector=self._connector,
                        timeout=self._timeout,
                        headers=self._headers,
                        raise_for_status=False,
                    )
        return self._session

    async def close(self) -> None:
        """Close the HTTP session and TCP connector, releasing all connections."""
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None
        if not self._connector.closed:
            await self._connector.close()

    async def close_session(self) -> None:
        """Alias for close(). Kept for backward compatibility."""
        await self.close()

    # === Cache ===

    def _get_cache_key(self, method: str, params: list[Any]) -> str:
        return f"{method}:{orjson.dumps(params).decode()}"

    def _get_cached(self, key: str) -> Any | None:
        if not self.config.enable_cache:
            return None
        return self._cache.get(key)

    def _set_cache(self, key: str, value: Any) -> None:
        if self.config.enable_cache:
            self._cache[key] = value

    # === Retry ===

    async def _retry_request(self, coro: Callable, *args: Any, **kwargs: Any) -> Any:
        """Execute a coroutine with exponential back-off on network errors."""
        last_error: Exception | None = None
        for attempt in range(self.config.max_retries):
            try:
                return await coro(*args, **kwargs)
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                last_error = e
                if attempt < self.config.max_retries - 1:
                    delay = self.config.retry_base_delay * (2 ** attempt)
                    delay += random.random() * self.config.retry_jitter
                    logger.debug(
                        f"Retry {attempt + 1}/{self.config.max_retries}"
                        f" after {delay:.2f}s: {e}"
                    )
                    await asyncio.sleep(delay)
        raise last_error  # type: ignore[misc]

    # === Request dispatch ===

    async def _send_request(
        self,
        method: str,
        params: list[Any],
        cacheable: bool = False,
    ) -> Any:
        """Send a single JSON-RPC request, optionally reading/writing the TTL cache."""
        cache_key = self._get_cache_key(method, params) if cacheable else None

        if cache_key:
            cached = self._get_cached(cache_key)
            if cached is not None:
                return cached

        payload = {
            "jsonrpc": "2.0",
            "id": next(_request_id),
            "method": method,
            "params": params,
        }

        async def _do_request() -> Any:
            session = await self._ensure_session()
            async with session.post(self.config.rpc_url, data=orjson.dumps(payload)) as resp:
                if resp.status != 200:
                    raise aiohttp.ClientError(f"HTTP {resp.status}: {await resp.text()}")
                data = orjson.loads(await resp.read())
                if "error" in data:
                    err = data["error"]
                    raise RPCError(
                        code=err.get("code", -1),
                        message=err.get("message", "Unknown error"),
                        data=err.get("data"),
                    )
                return data["result"]

        result = await self._retry_request(_do_request)
        if cache_key:
            self._set_cache(cache_key, result)
        return result

    async def _send_batch(self, requests: list[CallRequest]) -> list[Any]:
        """Send a JSON-RPC batch request and return results in original call order."""
        if not requests:
            return []

        payloads = [
            {
                "jsonrpc": "2.0",
                "id": req.id,
                "method": req.method,
                "params": req.params,
            }
            for req in requests
        ]

        async def _do_batch() -> list[dict[str, Any]]:
            session = await self._ensure_session()
            async with session.post(self.config.rpc_url, data=orjson.dumps(payloads)) as resp:
                if resp.status != 200:
                    raise aiohttp.ClientError(f"HTTP {resp.status}: {await resp.text()}")
                return orjson.loads(await resp.read())

        responses: list[dict[str, Any]] = await self._retry_request(_do_batch)
        id_to_response = {r["id"]: r for r in responses}
        results: list[Any] = []

        for req in requests:
            resp = id_to_response.get(req.id)
            if not resp:
                results.append(None)
            elif "error" in resp:
                if resp["error"]["message"] != "execution reverted":
                    logger.error(f"Batch error for {req.method}: {resp['error']}")
                results.append(None)
            else:
                results.append(resp.get("result"))

        return results

    # === Basic RPC methods ===

    async def get_chain_id(self) -> int:
        """Return network chain ID. Fetched once per client instance, then stored."""
        if self.chain_id is None:
            result = await self._send_request("eth_chainId", [], cacheable=True)
            self.chain_id = int(result, 16)
        return self.chain_id

    async def get_block_number(self) -> int:
        """Return the current latest block number."""
        result = await self._send_request("eth_blockNumber", [])
        return int(result, 16)

    async def get_balance(self, address: str, block: str = "latest") -> int:
        """Return account balance in wei."""
        result = await self._send_request("eth_getBalance", [address, block])
        return int(result, 16)

    async def get_nonce(self, address: str, block: str = "pending") -> int:
        """Return the transaction count (nonce) for the given address."""
        result = await self._send_request("eth_getTransactionCount", [address, block])
        return int(result, 16)

    async def get_bytecode(self, address: str, block: str = "latest") -> str:
        """Return contract bytecode at the given block. Cached when block='latest'."""
        cacheable = block == "latest"
        return await self._send_request("eth_getCode", [address, block], cacheable=cacheable)

    async def get_gas_price(self) -> int:
        """Return the current gas price in wei."""
        result = await self._send_request("eth_gasPrice", [])
        return int(result, 16)

    async def call(
        self,
        to: str,
        data: str,
        from_address: str | None = None,
        block: str = "latest",
    ) -> str:
        """Execute eth_call and return the raw hex result."""
        tx: dict[str, str] = {"to": to, "data": data}
        if from_address:
            tx["from"] = from_address
        return await self._send_request("eth_call", [tx, block])

    async def estimate_gas(self, tx: dict[str, Any]) -> int:
        """Estimate gas usage for a transaction."""
        result = await self._send_request("eth_estimateGas", [tx])
        return int(result, 16)

    async def send_raw_transaction(self, signed_tx: str) -> str:
        """Broadcast a signed transaction and return its hash."""
        return await self._send_request("eth_sendRawTransaction", [signed_tx])

    # === Block methods ===

    async def get_block_by_number(
        self,
        block_number: int,
        full_transactions: bool = True,
    ) -> dict[str, Any] | None:
        """Return block data by number with timestamp and number converted from hex."""
        result = await self._send_request(
            "eth_getBlockByNumber",
            [hex(block_number), full_transactions],
        )
        if result:
            result["timestamp"] = int(result["timestamp"], 16)
            result["number"] = int(result["number"], 16)
        return result

    async def batch_get_blocks(
        self,
        block_numbers: list[int],
    ) -> list[dict[str, Any] | None]:
        """Fetch multiple blocks in a single JSON-RPC batch request."""
        requests = [
            CallRequest("eth_getBlockByNumber", [hex(bn), True])
            for bn in block_numbers
        ]
        return await self._send_batch(requests)

    # === Transaction methods ===

    async def get_transaction(self, tx_hash: str) -> dict[str, Any] | None:
        """Return transaction data by hash."""
        return await self._send_request("eth_getTransactionByHash", [tx_hash])

    async def get_transaction_receipt(self, tx_hash: str) -> dict[str, Any] | None:
        """Return transaction receipt by hash."""
        return await self._send_request("eth_getTransactionReceipt", [tx_hash])

    async def batch_get_transactions(self, tx_hashes: list[str]) -> dict[str, Any]:
        """Fetch multiple transactions in one batch request, deduplicating hashes."""
        unique_hashes = list(dict.fromkeys(tx_hashes))
        requests = [CallRequest("eth_getTransactionByHash", [h]) for h in unique_hashes]
        results = await self._send_batch(requests)
        return {h: r for h, r in zip(unique_hashes, results) if r is not None}

    async def batch_get_receipts(self, tx_hashes: list[str]) -> dict[str, Any]:
        """Fetch multiple receipts in one batch request, deduplicating hashes."""
        unique_hashes = list(dict.fromkeys(tx_hashes))
        requests = [CallRequest("eth_getTransactionReceipt", [h]) for h in unique_hashes]
        results = await self._send_batch(requests)
        return {h: r for h, r in zip(unique_hashes, results) if r is not None}

    # === EIP-1559 methods ===

    async def get_fee_history(
        self,
        block_count: int,
        newest_block: str = "latest",
        reward_percentiles: list[float] | None = None,
    ) -> dict[str, Any]:
        """Return fee history data for EIP-1559 fee estimation."""
        params: list[Any] = [hex(block_count), newest_block]
        if reward_percentiles:
            params.append(reward_percentiles)
        return await self._send_request("eth_feeHistory", params)

    async def get_max_priority_fee(self) -> int:
        """Return the suggested max priority fee per gas in wei."""
        result = await self._send_request("eth_maxPriorityFeePerGas", [])
        return int(result, 16)

    async def get_eip1559_fees(self) -> dict[str, int]:
        """Return current EIP-1559 fee parameters: base_fee, max_priority_fee, max_fee."""
        requests = [
            CallRequest("eth_getBlockByNumber", ["latest", False]),
            CallRequest("eth_maxPriorityFeePerGas", []),
        ]
        results = await self._send_batch(requests)

        block, max_priority_hex = results
        base_fee = int(block["baseFeePerGas"], 16) if block else 0
        max_priority = int(max_priority_hex, 16) if max_priority_hex else 0

        return {
            "base_fee": base_fee,
            "max_priority_fee": max_priority,
            "max_fee": base_fee * 2 + max_priority,
        }

    async def get_params_for_eip1559_transaction(self, address: str) -> dict[str, int]:
        """Return nonce + EIP-1559 fee parameters in a single batch call."""
        requests = [
            CallRequest("eth_getTransactionCount", [address, "pending"]),
            CallRequest("eth_getBlockByNumber", ["latest", False]),
            CallRequest("eth_maxPriorityFeePerGas", []),
        ]
        results = await self._send_batch(requests)

        nonce_hex, block, max_priority_hex = results
        nonce = int(nonce_hex, 16) if nonce_hex else 0
        base_fee = int(block["baseFeePerGas"], 16) if block else 0
        max_priority = int(max_priority_hex, 16) if max_priority_hex else 0

        return {
            "nonce": nonce,
            "base_fee_per_gas": base_fee,
            "max_priority_fee_per_gas": max_priority,
        }

    # === Multicall helpers ===

    async def multicall(
        self,
        multicall_contract: Any,
        calls: list[tuple[Any, str, tuple[Any, ...]]],
        require_success: bool = False,
    ) -> list[Any]:
        """Execute Multicall3.tryAggregate in a single eth_call.

        Args:
            multicall_contract: Contract instance with the Multicall3 ABI and address.
            calls: List of (contract, function_name, args).
            require_success: If True, the on-chain call reverts on any single failure.

        Returns:
            Decoded results in the same order as calls; None for failed/empty calls.
        """
        call_tuples: list[tuple[str, bytes]] = []
        for contract, fn_name, args in calls:
            calldata = contract.encode_function_data(fn_name, *args)
            call_tuples.append((contract.address, bytes.fromhex(calldata[2:])))

        encoded = abi_encode(["bool", "(address,bytes)[]"], [require_success, call_tuples])
        selector = "0xbce38bd7"  # tryAggregate(bool,(address,bytes)[])
        calldata_hex = selector + encoded.hex()

        raw = await self.call(multicall_contract.address, calldata_hex)
        result_bytes = bytes.fromhex(raw[2:] if raw.startswith("0x") else raw)
        decoded_raw = abi_decode(["(bool,bytes)[]"], result_bytes)[0]

        results: list[Any] = []
        for i, (success, return_data) in enumerate(decoded_raw):
            if not success or not return_data:
                results.append(None)
                continue
            contract, fn_name, _ = calls[i]
            try:
                decoded = contract.decode_function_result(fn_name, "0x" + return_data.hex())
                results.append(decoded)
            except Exception as e:
                logger.debug(f"Failed to decode multicall result {i} ({fn_name}): {e}")
                results.append(None)

        return results

    async def call_function(self, contract: Any, function_name: str, *args: Any) -> Any:
        """Call a contract function and return the decoded result.

        Args:
            contract: Contract instance with ABI.
            function_name: Name of the function to call.
            *args: Positional arguments forwarded to the function.

        Returns:
            Decoded return value (unwrapped if single), or None for void functions.
        """
        calldata = contract.encode_function_data(function_name, *args)
        raw = await self.call(contract.address, calldata)
        return contract.decode_function_result(function_name, raw)

    # === Debug / trace methods ===

    async def debug_trace_transaction(
        self,
        tx_hash: str,
        tracer: str = "callTracer",
        timeout: str = "30s",
    ) -> dict[str, Any]:
        """Trace a transaction using the debug_traceTransaction RPC method."""
        return await self._send_request(
            "debug_traceTransaction",
            [tx_hash, {"tracer": tracer, "timeout": timeout}],
        )


# === Helper functions ===

@lru_cache(maxsize=1024)
def compute_function_selector(signature: str) -> str:
    """Return the 4-byte function selector hex string for the given ABI signature."""
    return "0x" + keccak(text=signature).hex()[:8]


@lru_cache(maxsize=1024)
def compute_event_topic(signature: str) -> str:
    """Return the 32-byte topic0 hex string for the given ABI event signature."""
    return "0x" + keccak(text=signature).hex()


def ensure_hex_prefix(value: str) -> str:
    """Return the string prefixed with '0x' if it isn't already."""
    return value if value.startswith("0x") else f"0x{value}"


def topic_from_address(address: str) -> str:
    """Encode an address as a zero-left-padded 32-byte log topic."""
    addr = address.lower().removeprefix("0x")
    return f"0x{'0' * 24}{addr}"


def topic_from_uint(value: int) -> str:
    """Encode an unsigned integer as a 32-byte log topic."""
    return f"0x{value:064x}"


def format_tx_dict(tx: dict[str, Any]) -> dict[str, Any]:
    """Convert integer transaction fields to hex strings for JSON-RPC compatibility."""
    formatted = dict(tx)
    for key in ("nonce", "value", "maxFeePerGas", "maxPriorityFeePerGas", "chainId", "type", "gas"):
        if key in formatted and isinstance(formatted[key], int):
            formatted[key] = hex(formatted[key])
    return formatted


def mkdata(sig: str, types: list[str], params: list[Any] | tuple[Any, ...]) -> str:
    """Build calldata from a function signature string + ABI-encoded parameters.

    Args:
        sig: Full function signature, e.g. 'transfer(address,uint256)'.
        types: List of ABI parameter type strings.
        params: Parameter values matching types.

    Returns:
        '0x'-prefixed hex calldata string.
    """
    selector = function_signature_to_4byte_selector(sig)
    encoded = abi_encode(types, list(params))
    return "0x" + (selector + encoded).hex()


async def rpc_retry(
    coro: Callable,
    *args: Any,
    retries: int = 5,
    base_delay: float = 0.15,
    **kwargs: Any,
) -> Any:
    """Standalone retry wrapper for arbitrary async RPC coroutines.

    Args:
        coro: Async callable to retry.
        *args: Positional arguments forwarded to coro.
        retries: Maximum number of attempts.
        base_delay: Base delay in seconds for exponential back-off.
        **kwargs: Keyword arguments forwarded to coro.

    Returns:
        Return value of coro on success.

    Raises:
        The last exception raised by coro after all retries are exhausted.
    """
    last_error: Exception | None = None
    for attempt in range(retries):
        try:
            return await coro(*args, **kwargs)
        except Exception as e:
            last_error = e
            if attempt < retries - 1:
                delay = base_delay * (2 ** attempt) + random.random() * 0.15
                logger.debug(f"rpc_retry attempt {attempt + 1}/{retries} after {delay:.2f}s: {e}")
                await asyncio.sleep(delay)
    raise last_error  # type: ignore[misc]
