from __future__ import annotations

from io import BytesIO
from typing import Any

import pandas as pd

from xfintech.serde.common.dataformat import DataFormat
from xfintech.serde.common.error import (
    SerialiserFailedError,
    SerialiserImportError,
    SerialiserInputTypeError,
    SerialiserNotSupportedError,
)
from xfintech.serde.common.serialiserlike import SerialiserLike


class PandasSerialiser(SerialiserLike):
    """
    描述:
    - PandasSerialiser序列化模块。
    - DataFrame -> BYTES 的转换，支持多种常见数据格式（Parquet、CSV、JSON）。

    方法:
    - serialise(data, format, **kwargs): 通用序列化方法，支持多种格式。
    - to_parquet(data, **kwargs): 将 DataFrame 序列化为 Parquet 格式的字节流。
    - to_csv(data, **kwargs): 将 DataFrame 序列化为 CSV 格式的字节流。
    - to_json(data, **kwargs): 将 DataFrame 序列化为 JSON 格式的字节流。

    例子:
    ```python
        from xfintech.serde.serialiser.pandas import PandasSerialiser
        from xfintech.serde.common.dataformat import DataFormat
        import pandas as pd

        # 使用通用序列化方法（推荐）
        bytes = PandasSerialiser.serialise(
            df,
            format=DataFormat.PARQUET,
        )
    ```
    """

    @staticmethod
    def serialise(
        data: pd.DataFrame,
        format: DataFormat | str,
        **kwargs: Any,
    ) -> bytes:
        if not isinstance(data, pd.DataFrame):
            msg = f"PandasSerialiser.serialise expects pd.DataFrame, got {type(data)}"
            raise SerialiserInputTypeError(msg)

        format = DataFormat.from_any(format)
        if format == DataFormat.PARQUET:
            return PandasSerialiser.to_parquet(data, **kwargs)
        elif format == DataFormat.CSV:
            return PandasSerialiser.to_csv(data, **kwargs)
        elif format == DataFormat.JSON:
            return PandasSerialiser.to_json(data, **kwargs)
        else:
            msg = f"Unsupported DataFormat for serialization: {format}"
            raise SerialiserNotSupportedError(msg)

    @staticmethod
    def to_parquet(
        data: pd.DataFrame,
        **kwargs: Any,
    ) -> bytes:
        kwargs.setdefault("index", False)
        kwargs.setdefault("engine", "pyarrow")
        kwargs.setdefault("coerce_timestamps", "us")
        kwargs.setdefault("allow_truncated_timestamps", True)
        buffer = BytesIO()
        try:
            data.to_parquet(buffer, **kwargs)
            return buffer.getvalue()
        except ImportError as e:
            msg = "pyarrow is required for parquet serialization. install via 'pip install pyarrow'."
            raise SerialiserImportError(msg) from e
        except Exception as e:
            msg = f"Failed to serialize data: {e}"
            raise SerialiserFailedError(msg) from e

    @staticmethod
    def to_csv(
        data: pd.DataFrame,
        **kwargs: Any,
    ) -> bytes:
        kwargs.setdefault("index", False)
        encoding = kwargs.pop("encoding", "utf-8")
        try:
            text = data.to_csv(**kwargs)
            return text.encode(encoding)
        except Exception as e:
            msg = f"Failed to serialize data: {e}"
            raise SerialiserFailedError(msg) from e

    @staticmethod
    def to_json(
        data: pd.DataFrame,
        **kwargs: Any,
    ) -> bytes:
        encoding = kwargs.pop("encoding", "utf-8")
        orient = kwargs.pop("orient", "records")
        lines = kwargs.pop("lines", False)
        try:
            text = data.to_json(
                orient=orient,
                lines=lines,
                **kwargs,
            )
            return text.encode(encoding)
        except Exception as e:
            msg = f"Failed to serialize data: {e}"
            raise SerialiserFailedError(msg) from e
