from __future__ import annotations

from io import BytesIO, StringIO
from typing import Any

import pandas as pd

from xfintech.serde.common.dataformat import DataFormat
from xfintech.serde.common.deserialiserlike import DeserialiserLike
from xfintech.serde.common.error import (
    DeserialiserFailedError,
    DeserialiserImportError,
    DeserialiserInputTypeError,
    DeserialiserNotSupportedError,
)


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

    方法:
    - deserialise(data, format, **kwargs): 通用反序列化方法，支持多种格式。
    - from_parquet(data, **kwargs): 从 Parquet 格式的字节流反序列化为 DataFrame。
    - from_csv(data, **kwargs): 从 CSV 格式的字节流反序列化为 DataFrame。
    - from_json(data, **kwargs): 从 JSON 格式的字节流反序列化为 DataFrame。

    例子:
    ```python
        from xfintech.serde.deserialiser.pandas import PandasDeserialiser
        from xfintech.serde.common.dataformat import DataFormat

        # 使用通用反序列化方法（推荐）
        df = PandasDeserialiser.deserialise(
            bytes,
            format=DataFormat.PARQUET,
        )

    ```
    """

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

        format = DataFormat.from_any(format)
        if format == DataFormat.PARQUET:
            return PandasDeserialiser.from_parquet(data, **kwargs)
        elif format == DataFormat.CSV:
            return PandasDeserialiser.from_csv(data, **kwargs)
        elif format == DataFormat.JSON:
            return PandasDeserialiser.from_json(data, **kwargs)
        else:
            msg = f"Unsupported DataFormat for deserialization: {format}"
            raise DeserialiserNotSupportedError(msg)

    @staticmethod
    def from_parquet(
        data: bytes,
        **kwargs: Any,
    ) -> pd.DataFrame:
        buffer = BytesIO(data)
        try:
            output = pd.read_parquet(buffer, **kwargs)
            return output
        except ImportError as e:
            msg = "pyarrow is required for parquet deserialization. install via 'pip install pyarrow'."
            raise DeserialiserImportError(msg) from e
        except Exception as e:
            msg = f"Failed to deserialize data: {e}"
            raise DeserialiserFailedError(msg) from e

    @staticmethod
    def from_csv(
        data: bytes,
        **kwargs: Any,
    ) -> pd.DataFrame:
        encoding = kwargs.pop("encoding", "utf-8")
        try:
            buffer = StringIO(data.decode(encoding))
            return pd.read_csv(buffer, **kwargs)
        except Exception as e:
            msg = f"Failed to deserialize data: {e}"
            raise DeserialiserFailedError(msg) from e

    @staticmethod
    def from_json(
        data: bytes,
        **kwargs: Any,
    ) -> pd.DataFrame:
        encoding = kwargs.pop("encoding", "utf-8")
        orient = kwargs.pop("orient", "records")
        lines = kwargs.pop("lines", False)
        try:
            text = data.decode(encoding)
            buf = StringIO(text)
            return pd.read_json(
                buf,
                orient=orient,
                lines=lines,
                **kwargs,
            )
        except Exception as e:
            msg = f"Failed to deserialize data: {e}"
            raise DeserialiserFailedError(msg) from e
