"""
Temporal Codec for DataFrame Serialization

This module provides a custom Temporal DataConverter that automatically
serializes pandas DataFrames to Parquet format and deserializes them back.

The codec only handles DataFrames, leaving all other types to the default
Temporal serialization mechanism.
"""

import io
from collections.abc import Sequence
from typing import Any, List, Optional, Type

from temporalio.api.common.v1 import Payload
from temporalio.converter import DataConverter, PayloadConverter
from pandas import DataFrame, read_parquet


class DataFramePayloadConverter(PayloadConverter):
    """
    Custom PayloadConverter that serializes pandas DataFrames to Parquet format.

    Only DataFrames are handled by this converter. All other types are passed
    to the default Temporal serialization mechanism.
    """

    def __init__(self, default_payload_converter: PayloadConverter):
        """
        Initialize the DataFrame payload converter.

        Args:
            default_payload_converter: The default Temporal payload converter
                to use for non-DataFrame types
        """
        self._default = default_payload_converter

    def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
        """
        Convert values to Temporal Payloads.

        If a value is a pandas DataFrame, it is serialized to Parquet format.
        Otherwise, the default converter is used.

        Args:
            values: The values to serialize

        Returns:
            List[Payload]: The serialized payloads
        """
        payloads = []
        for value in values:
            # Check if value is a DataFrame
            try:
                if isinstance(value, DataFrame):
                    buffer = io.BytesIO()
                    value.to_parquet(buffer, engine='pyarrow', index=True)
                    payloads.append(Payload(
                        metadata={"encoding": b"parquet-dataframe"},
                        data=buffer.getvalue()
                    ))
                    continue
            except ImportError:
                # pandas not available, fall through to default
                pass
            except Exception:
                # Error serializing DataFrame, fall through to default
                pass

            # Use default converter for all other types
            default_payloads = self._default.to_payloads([value])
            payloads.extend(default_payloads)
        
        return payloads

    def from_payloads(
        self,
        payloads: Sequence[Payload],
        type_hints: Optional[List[Type]] = None,
    ) -> List[Any]:
        """
        Convert Temporal Payloads back to Python values.

        If a payload metadata indicates it's a Parquet-serialized DataFrame,
        it is deserialized. Otherwise, the default converter is used.

        Args:
            payloads: The payloads to deserialize
            type_hints: Optional type hints for the expected return types

        Returns:
            List[Any]: The deserialized values
        """
        values = []
        for i, payload in enumerate(payloads):
            # Check if this is a Parquet-serialized DataFrame
            if payload.metadata.get("encoding") == b"parquet-dataframe":
                try:
                    buffer = io.BytesIO(payload.data)
                    values.append(read_parquet(buffer))
                    continue
                except ImportError:
                    # pandas not available, fall through to default
                    pass
                except Exception:
                    # Error deserializing DataFrame, fall through to default
                    pass

            # Use default converter for all other types
            # type_hints must have same length as payloads if provided
            type_hint = type_hints[i] if type_hints and i < len(type_hints) else None
            default_values = self._default.from_payloads([payload], [type_hint] if type_hint is not None else None)
            values.extend(default_values)
        
        return values


def create_dataframe_data_converter() -> DataConverter:
    """
    Create a DataConverter with DataFrame serialization support.

    Returns:
        DataConverter: A DataConverter that handles DataFrames automatically
    """
    # Get default converter to use as fallback
    # DataConverter.default is an attribute, not a method
    default_converter = DataConverter.default
    
    # Create a factory class that extends PayloadConverter
    class DataFramePayloadConverterFactory(PayloadConverter):
        def __init__(self):
            super().__init__()
            # Create instance of default payload converter to use as fallback
            self._default_converter = default_converter.payload_converter_class()
        
        def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
            return DataFramePayloadConverter(self._default_converter).to_payloads(values)
        
        def from_payloads(
            self,
            payloads: Sequence[Payload],
            type_hints: Optional[List[Type]] = None,
        ) -> List[Any]:
            return DataFramePayloadConverter(self._default_converter).from_payloads(payloads, type_hints)
    
    return DataConverter(
        payload_converter_class=DataFramePayloadConverterFactory
    )