"""Serialization utilities for Protocol Buffer-like messages.

This module provides msgpack-based serialization for proto dataclasses,
with custom handling for datetime objects and hash computation for
reproducibility tracking.
"""

from __future__ import annotations

import hashlib
from dataclasses import asdict, is_dataclass
from datetime import datetime, timezone
from enum import IntEnum
from typing import Any, TypeVar

import msgpack

T = TypeVar("T")


def _encode_value(obj: Any) -> Any:
    """Recursively encode a value for msgpack serialization.

    Args:
        obj: Value to encode.

    Returns:
        Msgpack-compatible value.
    """
    if isinstance(obj, datetime):
        return {"__datetime__": True, "iso": obj.isoformat()}
    elif isinstance(obj, IntEnum):
        return int(obj)
    elif is_dataclass(obj) and not isinstance(obj, type):
        return {k: _encode_value(v) for k, v in asdict(obj).items()}
    elif isinstance(obj, dict):
        return {k: _encode_value(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [_encode_value(item) for item in obj]
    else:
        return obj


def _decode_value(obj: Any, expected_type: type | None = None) -> Any:
    """Recursively decode a value from msgpack deserialization.

    Args:
        obj: Value to decode.
        expected_type: Optional expected type for reconstruction.

    Returns:
        Decoded value.
    """
    if isinstance(obj, dict):
        if obj.get("__datetime__"):
            iso_str = obj["iso"]
            # Handle both aware and naive datetimes
            dt = datetime.fromisoformat(iso_str)
            # Ensure timezone-aware
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        return {k: _decode_value(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_decode_value(item) for item in obj]
    else:
        return obj


def serialize_proto(proto: Any) -> bytes:
    """Serialize a proto dataclass to msgpack bytes.

    Args:
        proto: A dataclass instance to serialize.

    Returns:
        Msgpack-encoded bytes.

    Raises:
        TypeError: If proto is not a dataclass instance.
    """
    if not is_dataclass(proto) or isinstance(proto, type):
        raise TypeError(f"Expected dataclass instance, got {type(proto).__name__}")

    encoded = _encode_value(proto)
    return msgpack.packb(encoded, use_bin_type=True)


def deserialize_proto(data: bytes, proto_class: type[T]) -> T:
    """Deserialize msgpack bytes to a proto dataclass.

    Args:
        data: Msgpack-encoded bytes.
        proto_class: The dataclass type to instantiate.

    Returns:
        Instance of proto_class populated with decoded data.

    Raises:
        TypeError: If proto_class is not a dataclass.
        ValueError: If data cannot be deserialized.
    """
    if not is_dataclass(proto_class):
        raise TypeError(f"Expected dataclass type, got {proto_class}")

    try:
        decoded = msgpack.unpackb(data, raw=False)
    except Exception as e:
        raise ValueError(f"Failed to decode msgpack data: {e}") from e

    # Recursively decode values
    decoded = _decode_value(decoded)

    # Reconstruct the dataclass
    return _reconstruct_dataclass(decoded, proto_class)


def _reconstruct_dataclass(data: dict, proto_class: type[T]) -> T:
    """Reconstruct a dataclass from decoded data.

    Args:
        data: Decoded dictionary data.
        proto_class: The dataclass type to instantiate.

    Returns:
        Instance of proto_class.
    """
    import dataclasses
    import types
    import typing

    if not isinstance(data, dict):
        raise ValueError(f"Expected dict, got {type(data).__name__}")

    # Get field types for nested reconstruction, with type hints resolved
    type_hints = typing.get_type_hints(proto_class)
    fields = {f.name: f for f in dataclasses.fields(proto_class)}
    kwargs = {}

    for name, value in data.items():
        if name not in fields:
            continue

        # Use resolved type hints instead of field.type to handle forward refs
        field_type = type_hints.get(name, fields[name].type)

        # Handle string type annotations that couldn't be resolved
        if isinstance(field_type, str):
            kwargs[name] = value
            continue

        # Handle Optional types and Union types (X | None)
        origin = typing.get_origin(field_type)
        args = typing.get_args(field_type)

        # Check for Union types (typing.Union or types.UnionType for X | Y syntax)
        is_union = origin is typing.Union or isinstance(field_type, types.UnionType)
        if is_union:
            # Get the actual type from Optional[X] or X | None
            # A Union always has args, and at least one is not NoneType
            # (`None | None` collapses to NoneType, which is not a Union).
            if args:  # pragma: no branch
                # Filter out NoneType
                non_none_types = [t for t in args if t is not type(None)]
                if non_none_types:  # pragma: no branch
                    field_type = non_none_types[0]
                    # Re-evaluate origin for the actual type
                    origin = typing.get_origin(field_type)
                    args = typing.get_args(field_type)

        # Recursively reconstruct nested dataclasses
        if value is not None and is_dataclass(field_type) and isinstance(value, dict):
            value = _reconstruct_dataclass(value, field_type)
        elif origin is list and isinstance(value, list) and value:
            # Check if it's a list of dataclasses
            if args:
                list_item_type = args[0]
                if is_dataclass(list_item_type):
                    value = [
                        _reconstruct_dataclass(item, list_item_type)
                        if isinstance(item, dict)
                        else item
                        for item in value
                    ]

        kwargs[name] = value

    return proto_class(**kwargs)


def compute_proto_hash(proto: Any) -> str:
    """Compute a SHA256 hash of a serialized proto.

    The hash is computed from the msgpack-serialized bytes, providing
    a deterministic identifier for the proto contents.

    Args:
        proto: A dataclass instance to hash.

    Returns:
        Hex-encoded SHA256 hash string.

    Raises:
        TypeError: If proto is not a dataclass instance.
    """
    serialized = serialize_proto(proto)
    return hashlib.sha256(serialized).hexdigest()