from __future__ import annotations
from web3 import Web3

import typing as _t
from typing import get_args, get_origin, ClassVar
from typing import Any, Dict, Type

from eth_abi import encode as abi_encode
from eth_utils import keccak
from pydantic import BaseModel

from .eip712_struct import EIP712Struct


def _struct_string(cls: Type["EIP712Struct"]) -> str:
    sol_types = getattr(cls, "SOL_TYPES", {})
    if not sol_types:
        return cls.TYPE_STRING

    hints = _t.get_type_hints(cls)
    parts: list[str] = []

    for field, t in hints.items():
        if get_origin(t) is ClassVar:
            continue
        if isinstance(t, type) and issubclass(t, EIP712Struct):
            parts.append(f"{sol_types.get(field, t.SOL_TYPE_NAME)} {field}")
        else:
            parts.append(f"{sol_types[field]} {field}")

    return f"{cls.SOL_TYPE_NAME}(" + ",".join(parts) + ")"


def _gather_deps(
    cls: Type[EIP712Struct], seen: set[Type[EIP712Struct]]
) -> list[Type[EIP712Struct]]:
    """Depth-first collect all user-defined nested structs (no duplicates)."""
    deps: list[Type[EIP712Struct]] = []
    for t in _t.get_type_hints(cls).values():
        if get_origin(t) is ClassVar:
            continue
        if isinstance(t, type) and issubclass(t, EIP712Struct) and t not in seen:
            seen.add(t)
            deps.extend(_gather_deps(t, seen))  # recurse first
            deps.append(t)  # then this struct itself
    return deps


def _canonical_type_string(cls: Type[EIP712Struct]) -> str:
    """Root encodeType + ASCII-sorted dependencies (EIP-712 protocol definition)."""
    deps = sorted(_gather_deps(cls, set()), key=lambda c: c.SOL_TYPE_NAME)
    return _struct_string(cls) + "".join(_struct_string(d) for d in deps)


class _AutoMeta(type(BaseModel)):
    """Populate TYPE_STRING / TYPE_STRUCT at **class-definition** time."""

    def __init__(cls, name, bases, ns, **kw):
        super().__init__(name, bases, ns, **kw)

        if cls.__name__ == "EIP712AutoStruct":
            return

        cls.TYPE_STRING = _canonical_type_string(cls)  # type: ignore[attr-defined]

        hints = _t.get_type_hints(cls)
        struct: list[Dict[str, str]] = []
        for field, t in hints.items():
            if get_origin(t) is ClassVar:
                continue
            if isinstance(t, type) and issubclass(t, EIP712Struct):
                struct.append(
                    {
                        "name": field,
                        "type": getattr(cls, "SOL_TYPES", {}).get(field, t.SOL_TYPE_NAME),
                    }
                )
            else:
                struct.append({"name": field, "type": cls.SOL_TYPES[field]})
        cls.TYPE_STRUCT = struct  # type: ignore[attr-defined]


def _coerce_bytes(val: Any) -> bytes:
    """Accept raw bytes *or* '0x…' str/bytes and return binary bytes."""
    if isinstance(val, (bytes, bytearray)):
        b = bytes(val)
        if b.startswith(b"0x"):
            return Web3.to_bytes(hexstr=b.decode())
        return b
    if isinstance(val, str):
        if val.startswith("0x"):
            return Web3.to_bytes(hexstr=val)
        return val.encode()
    raise TypeError("Expected bytes-like or 0x-string")


def _hash_fixed_bytes_array(n: int, seq: list[bytes]) -> bytes:
    """Hash a ``bytesN[]`` value-type array (N ≤ 32)."""
    packed = bytearray(32 * len(seq))
    for i, raw in enumerate(seq):
        b = _coerce_bytes(raw)
        if len(b) != n:
            raise ValueError(
                f"bytes{n} element length mismatch: Expected {n} bytes, got {len(b)}, value: {b.hex()}"
            )
        packed[i * 32 : i * 32 + n] = b
    return keccak(bytes(packed))


def _hash_dynamic_bytes_array(seq: list[bytes]) -> bytes:
    """Hash a ``bytes[]`` (or ``string[]``) array."""
    packed = bytearray(32 * len(seq))
    for i, raw in enumerate(seq):
        packed[i * 32 : (i + 1) * 32] = keccak(_coerce_bytes(raw))
    return keccak(bytes(packed))


def _hash_value_type_array(base: str, seq: list[Any]) -> bytes:
    packed = b"".join(abi_encode([base], [x]) for x in seq)
    return keccak(packed)


def _hash_array(sol_array_type: str, seq: list[Any]) -> bytes:
    base = sol_array_type[:-2]

    if base in ("bytes", "string"):
        return _hash_dynamic_bytes_array(seq)

    if base.startswith("bytes") and base != "bytes":
        n = int(base[5:])
        return _hash_fixed_bytes_array(n, seq)

    return _hash_value_type_array(base, seq)


class EIP712AutoStruct(EIP712Struct, metaclass=_AutoMeta):
    SOL_TYPES: ClassVar[Dict[str, str]] = {}
    model_config = {"arbitrary_types_allowed": True}

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "EIP712AutoStruct":
        field_types = _t.get_type_hints(cls)
        kwargs = {}

        for field in cls.__annotations__.keys():
            if get_origin(field_types.get(field)) is ClassVar:
                continue

            if field not in data:
                continue

            field_value = data[field]
            field_type = field_types.get(field)

            if isinstance(field_type, type) and issubclass(field_type, EIP712Struct):
                kwargs[field] = field_type.from_dict(field_value)
            else:
                if (
                    field_type is bytes
                    and isinstance(field_value, str)
                    and field_value.startswith("0x")
                ):
                    kwargs[field] = Web3.to_bytes(hexstr=field_value)
                elif (
                    field_type is str
                    and cls.SOL_TYPES.get(field) == "address"
                    and isinstance(field_value, str)
                ):
                    kwargs[field] = Web3.to_checksum_address(field_value)
                elif (
                    get_origin(field_type) in (list, _t.List)
                    and get_args(field_type)[0] is bytes
                ):
                    cleaned = []
                    for elem in field_value:
                        if isinstance(elem, str) and elem.startswith("0x"):
                            cleaned.append(Web3.to_bytes(hexstr=elem))
                        else:
                            cleaned.append(elem)
                    kwargs[field] = cleaned
                else:
                    kwargs[field] = field_value

        return cls(**kwargs)

    def hash(self) -> bytes:
        abi_types: list[str] = ["bytes32"]
        abi_vals: list[Any] = [self.type_hash]

        for field in self.__class__.__annotations__.keys():
            sol_type = self.SOL_TYPES[field]
            val = getattr(self, field)

            if isinstance(val, EIP712Struct):
                abi_types.append("bytes32")
                abi_vals.append(val.hash())
                continue

            if sol_type in ("bytes", "string"):
                abi_types.append("bytes32")
                abi_vals.append(keccak(_coerce_bytes(val)))
                continue

            if sol_type.endswith("[]"):
                abi_types.append("bytes32")
                abi_vals.append(_hash_array(sol_type, val))
                continue

            abi_types.append(sol_type)
            abi_vals.append(val)

        return keccak(abi_encode(abi_types, abi_vals))
