import logging
import re
from time import time
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Coroutine,
    DefaultDict,
    Dict,
    Iterable,
    Iterator,
    List,
    Literal,
    Mapping,
    NewType,
    NoReturn,
    Optional,
    Set,
    Tuple,
    Type,
    TypedDict,
    TypeVar,
    Union,
    overload,
)

import evmspec
from dictstruct import DictStruct
from eth_typing import ChecksumAddress, HexStr
from eth_utils.toolz import valmap
from evmspec.data import Address, BlockNumber, ChainId, Wei, uint, _decode_hook
from evmspec.structs.block import BaseBlock, Block, MinedBlock, ShanghaiCapellaBlock
from evmspec.structs.log import Log
from hexbytes import HexBytes
from msgspec import UNSET, Raw, ValidationError, json
from web3.datastructures import AttributeDict
from web3.types import RPCEndpoint, RPCResponse

from dank_mids import constants, stats
from dank_mids._exceptions import (
    BadResponse,
    ChainstackRateLimited,
    ExceedsMaxBatchSize,
    OutOfGas,
    PayloadTooLarge,
)

if TYPE_CHECKING:
    from dank_mids._requests import Multicall


T = TypeVar("T")


BlockId = NewType("BlockId", str)
"""A type representing the identifier for a specific block in the blockchain."""

BatchId = Union[int, str]
"""A type representing the identifier for a batch of operations, which can be either an integer or a string."""

Multicalls = Dict[BlockId, "Multicall"]
"""
A dictionary mapping BlockId to Multicall objects.

See Also:
    :class:`dank_mids._requests.Multicall`: The Multicall class used in this mapping.
"""

eth_callParams = TypedDict("eth_callParams", {"to": ChecksumAddress, "data": str})
"""
A typed dictionary representing the parameters for an eth_call.

See Also:
    :meth:`web3.eth.Eth.call`: Web3's method to perform a call without creating a transaction.
"""

OverrideParams = TypedDict("OverrideParams", {"code": str})
"""A typed dictionary representing override parameters."""

JsonrpcParams = List[Union[eth_callParams, BlockId, OverrideParams]]
"""A list of parameters for JSON-RPC calls, which should be eth_callParams, BlockId, and OverrideParams, in that order."""


# This type alias was introduced in web3 v5.28.0 but we like loose deps here so we recreate instead of import.
AsyncMiddleware = Callable[[RPCEndpoint, Any], Coroutine[Any, Any, RPCResponse]]
"""A type alias for asynchronous middleware functions."""

_list_of_stuff = List[Union[str, None, dict, list]]
"""A type alias for a list that can contain strings, None, dictionaries, or lists."""

_dict_of_stuff = Dict[str, Union[str, None, _list_of_stuff, Dict[str, Optional[Any]]]]
"""A type alias for a dictionary with string keys and values that can be strings, None, :obj:`_list_of_stuff`, or dictionaries with optional values."""

_nested_dict_of_stuff = Dict[str, Union[str, None, _list_of_stuff, _dict_of_stuff]]
"""A type alias for a nested dictionary structure."""


class PartialRequest(DictStruct, frozen=True, omit_defaults=True, repr_omit_defaults=True):  # type: ignore [call-arg]
    """
    Represents a partial JSON-RPC request.

    While technially part of a request, we can successfully make requests to many nodes without including the `jsonrpc` field.

    This class leaves off the `jsonrpc` field reduce encoding burden and web traffic.

    This works with many but not all nodes.
    """

    method: str
    """The RPC method to be called."""

    id: Union[str, int]
    """The request identifier."""

    params: Optional[list] = None
    """The parameters for the RPC method call."""

    @property
    def data(self) -> bytes:
        return json.encode(self, enc_hook=_encode_hook)


class Request(PartialRequest):
    """
    Represents a complete JSON-RPC request.

    Inherits from PartialRequest and adds the JSON-RPC version.
    """

    jsonrpc: Literal["2.0"] = "2.0"
    """The JSON-RPC version, always set to "2.0"."""

    @property
    def data(self) -> bytes:
        # we have to do some hacky stuff because omit_defaults kwarg on PartialRequest
        # is preventing jsonrpc field from being included in the encoded bytes
        encoded_omit_defaults = json.encode(self, enc_hook=_encode_hook)
        return encoded_omit_defaults[:-1] + b',"jsonrpc":"2.0"}'


_getattribute = object.__getattribute__


class Error(DictStruct, frozen=True, omit_defaults=True, repr_omit_defaults=True):  # type: ignore [call-arg]
    """
    Represents an error in a JSON-RPC response.
    """

    code: int
    """The error code."""

    message: str
    """The error message."""

    data: Optional[Any] = UNSET
    """Additional error data, if any."""

    def to_dict(self) -> Dict[str, Any]:
        d = {"code": self.code, "message": self.message}
        data = _getattribute(self, "data")
        if data is not UNSET:
            d["data"] = data
        return d


# some devving tools that will go away eventually
_dict_responses: Set[str] = set()
_str_responses: Set[str] = set()


_RETURN_TYPES = {
    "eth_call": HexBytes,
    "eth_chainId": ChainId,
    "eth_getCode": HexBytes,
    "eth_getLogs": List[Log],
    "eth_getBalance": Wei,
    "eth_blockNumber": BlockNumber,
    "eth_accounts": List[Address],
    "eth_getBlockByNumber": Block,
    "eth_getTransactionCount": uint,
    "eth_getTransactionByHash": evmspec.Transaction,
    "eth_getTransactionReceipt": evmspec.FullTransactionReceipt,
    "erigon_getHeaderByNumber": evmspec.ErigonBlockHeader,
    "trace_filter": List[evmspec.FilterTrace],
    "trace_transaction": List[evmspec.FilterTrace],
}
"""
A dictionary mapping RPC method names to their expected return types.
Used to enable more efficient decoding and validation of RPC responses.
"""


_CHAINSTACK_429_ERR_MSG = "You've exceeded the RPS limit available on the current plan."


class PartialResponse(DictStruct, frozen=True, omit_defaults=True, repr_omit_defaults=True):
    """
    Represents a partial JSON-RPC response.

    We use these to more efficiently decode responses from the node.
    """

    result: Raw = None  # type: ignore
    """The result of the RPC call, if successful."""

    error: Optional[Error] = None
    """The error object, if the call failed."""

    @property
    def exception(self) -> Exception:
        """If the rpc response contains an 'error' field, returns a specialized exception for the specified rpc error."""
        if self.error is None:
            raise AttributeError(f"{self} did not error.")

        message = self.error.message
        if "out of gas" in message:
            return OutOfGas(self)
        if self.payload_too_large:
            return PayloadTooLarge(self)
        if re.search(r"batch limit (\d+) exceeded", message):
            return ExceedsMaxBatchSize(self)
        if message == "invalid argument 1: hex string without 0x prefix":
            return TypeError(
                message,
                "DANKMIDS NOTE: You're probably passing what should be an integer type as a string type. "
                "The usual culprit is a block number.",
            )
        # chainstack doesnt return a response with status code 429 when we reach rate limits, so we need to handle it specifically here instead of in the usual place
        if _CHAINSTACK_429_ERR_MSG in message:
            return ChainstackRateLimited(self)
        return BadResponse(self)

    @property
    def payload_too_large(self) -> bool:
        return any(map(self.error.message.__contains__, constants.TOO_MUCH_DATA_ERRS))

    def to_dict(self, method: Optional[RPCEndpoint] = None) -> RPCResponse:  # type: ignore [override]
        """Returns a complete dictionary representation of this response ``Struct``."""
        data: RPCResponse = {
            key: self.decode_result(method=method, caller=self) if key == "result" else value
            for key, value in self.items()
            if value is not None
        }
        return data

    def decode_result(
        self, method: Optional[RPCEndpoint] = None, *, caller=None
    ) -> Union[HexBytes, Wei, uint, ChainId, BlockNumber, AttributeDict]:
        # NOTE: These must be added to the `_RETURN_TYPES` constant above manually
        if method and (typ := _RETURN_TYPES.get(method)):
            if method in [
                "eth_call",
                "eth_blockNumber",
                "eth_getCode",
                "eth_getBlockByNumber",
                "eth_getTransactionReceipt",
                "eth_getTransactionCount",
                "eth_getBalance",
                "eth_chainId",
                "erigon_getHeaderByNumber",
            ]:
                try:
                    return better_decode(
                        self.result, type=typ, dec_hook=_decode_hook, method=method
                    )
                except Exception as e:
                    if typ is not Block:
                        raise

                    if e.args[0] == "Object contains unknown field `totalDifficulty`":
                        try:
                            # NOTE should we do this??
                            # _RETURN_TYPES[method] = MinedBlock
                            return better_decode(
                                self.result, type=MinedBlock, dec_hook=_decode_hook, method=method
                            )
                        except ValidationError as e2:
                            if e2.args[0] != "Object contains unknown field `baseFeePerGas`":
                                raise
                            result = better_decode(
                                self.result, type=BaseBlock, dec_hook=_decode_hook, method=method
                            )
                            _RETURN_TYPES[method] = BaseBlock  # all blocks on base are BaseBlocks
                            return result

                    elif e.args[0] == "Object contains unknown field `withdrawals`":
                        return better_decode(
                            self.result,
                            type=ShanghaiCapellaBlock,
                            dec_hook=_decode_hook,
                            method=method,
                        )
                    else:
                        raise

                    return better_decode(
                        self.result, type=typ, dec_hook=_decode_hook, method=method
                    )

            start = time()
            try:
                decoded = better_decode(self.result, type=typ, dec_hook=_decode_hook, method=method)
            except (ValidationError, TypeError) as e:
                stats.logger.log_validation_error(self, e)
                raise

            if caller:
                stats.log_duration(f"decoding {type(caller)} {method}", start)
            return decoded

        # We have some semi-smart logic for providing decoder hints even if method not in `_RETURN_TYPES`
        if method:
            try:
                if method in _dict_responses:
                    decoded = AttributeDict(json.decode(self.result, type=_nested_dict_of_stuff))
                    stats.logger.log_types(method, decoded)
                    return decoded
                elif method in _str_responses:
                    # TODO: finish adding methods and get rid of this
                    stats.logger.devhint("Must add `%s: str` to `_RETURN_TYPES`", method)
                    return json.decode(self.result, type=str)
            except (ValidationError, TypeError) as e:
                stats.logger.log_validation_error(method, e)

        # In this case we can provide no hints, let's let the decoder figure it out
        decoded = json.decode(self.result)
        if isinstance(decoded, str):
            if method:
                _str_responses.add(method)
            return decoded
        elif isinstance(decoded, dict):
            if method:
                _dict_responses.add(method)
            return AttributeDict(decoded)
        elif isinstance(decoded, list):
            if method is None:
                return decoded
        raise TypeError(
            f"type {type(decoded)} is not supported.  method: {method}  decoded: {decoded}"
        )


class Response(PartialResponse, omit_defaults=True, repr_omit_defaults=True):  # type: ignore [call-arg]
    """
    Represents a complete JSON-RPC response.

    Inherits from PartialResponse and adds the response identifier and JSON-RPC version.
    """

    id: Optional[Union[str, int]] = None
    """The response identifier, matching the request."""

    jsonrpc: Literal["2.0"] = "2.0"
    """The JSON-RPC version, always set to "2.0"."""


class RawResponse:
    """
    Wraps a Raw object that we know represents a Response with a `decode` helper method.
    A `RawResponse` is a properly shaped response for one rpc call, received back from a jsonrpc batch request.
    They represent either a successful or a failed response, stored as pre-decoded bytes.
    """

    def __init__(self, raw: Raw) -> None:
        self._raw = raw
        """The :class:`Raw` object wrapped by this wrapper."""

    @overload
    def decode(self, partial: Literal[True]) -> PartialResponse: ...
    @overload
    def decode(self, partial: Literal[False] = False) -> Response: ...
    def decode(self, partial: bool = False) -> Union[Response, PartialResponse]:
        """Decode the wrapped :class:`Raw` object into a :class:`Response` or a :class:`PartialResponse`."""
        return better_decode(self._raw, type=PartialResponse if partial else Response)


JSONRPCBatchRequest = List[Request]
# NOTE: A PartialResponse result implies a failure response from the rpc.
JSONRPCBatchResponse = Union[List[RawResponse], PartialResponse]
# We need this for proper decoding.
JSONRPCBatchResponseRaw = Union[List[Raw], PartialResponse]


StrEncodable = Union[ChecksumAddress, HexStr]
Encodable = Union[int, StrEncodable, HexBytes, bytes]

RpcThing = Union[HexStr, List[HexStr], Dict[str, HexStr]]


def _encode_hook(obj: Encodable) -> RpcThing:
    """
    A hook function for encoding objects during JSON serialization.

    Args:
        obj: The object to encode.

    Returns:
        The encoded object.

    Raises:
        NotImplementedError: If the object type is not supported for encoding.
    """
    try:
        # We just assume `obj` is an int subclass instead of performing if checks because it usually is.
        return hex(int(obj))  # type: ignore [return-value]
    except TypeError as e:
        # I put this here for AttributeDicts which come from eth_getLogs params
        # but I check for mapping so it can work with user custom classes
        if not isinstance(obj, Mapping):
            raise TypeError(obj, type(obj)) from e
        return valmap(_rudimentary_encode_dict_value, obj)
    except ValueError as e:
        # NOTE: The error is probably this if `obj` is a string:
        # ValueError: invalid literal for int() with base 10:"""
        if not isinstance(obj, HexBytes):
            e.args = *e.args, obj, type(obj)
            raise ValueError(obj, type(obj)) from e
        return obj.hex()  # type: ignore [return-value]


def _rudimentary_encode_dict_value(value):
    # I dont think this needs to be robust, time will tell
    return hex(value) if isinstance(value, int) else value


def better_decode(
    data: Raw,
    *,
    type: Optional[Type[T]] = None,
    dec_hook: Optional[Callable[[Type, object], T]] = None,
    method: Optional[str] = None,
) -> T:
    try:
        return json.decode(data, type=type, dec_hook=dec_hook)
    except (ValidationError, TypeError) as e:
        extra_args = [
            f"type: {type.__module__}.{type.__qualname__}",
            f"result: {json.decode(data)}",
        ]
        if method:
            extra_args.insert(0, f"method: {method}")
        e.args = (*e.args, *extra_args)
        raise
