"""
Pydantic payload converter for Temporal
Adapted from: https://github.com/temporalio/samples-python/blob/main/pydantic_converter_v1/converter.py
"""

import json
from typing import Any, Optional

from pydantic.json import pydantic_encoder
from temporalio.api.common.v1 import Payload
from temporalio.converter import (
    CompositePayloadConverter,
    DataConverter,
    DefaultPayloadConverter,
    JSONPlainPayloadConverter,
)


class PydanticJSONPayloadConverter(JSONPlainPayloadConverter):
    """
    Pydantic JSON payload converter.

    This extends the JSONPlainPayloadConverter to override to_payload
    using the Pydantic encoder, enabling automatic serialization of
    Pydantic models with their validation and custom serialization logic.
    """

    def to_payload(self, value: Any) -> Optional[Payload]:
        """
        Convert all values with Pydantic encoder or fail.

        Like the base class, we fail if we cannot convert. This payload
        converter is expected to be the last in the chain, so it can fail
        if unable to convert.

        Args:
            value: Any Python object to serialize

        Returns:
            Temporal Payload with JSON-encoded data
        """
        # We let JSON conversion errors be thrown to caller
        return Payload(
            metadata={"encoding": self.encoding.encode()},
            data=json.dumps(
                value, separators=(",", ":"), sort_keys=True, default=pydantic_encoder
            ).encode(),
        )


class PydanticPayloadConverter(CompositePayloadConverter):
    """
    Payload converter that replaces Temporal JSON conversion with Pydantic JSON conversion.

    This converter maintains all default Temporal converters but replaces
    the JSON converter with PydanticJSONPayloadConverter, enabling seamless
    Pydantic model serialization.
    """

    def __init__(self) -> None:
        """Initialize with default converters, replacing JSON with Pydantic."""
        super().__init__(
            *(
                (
                    c
                    if not isinstance(c, JSONPlainPayloadConverter)
                    else PydanticJSONPayloadConverter()
                )
                for c in DefaultPayloadConverter.default_encoding_payload_converters
            )
        )


pydantic_data_converter = DataConverter(payload_converter_class=PydanticPayloadConverter)
