"""
MinIO-backed DataFrame payload for Temporal workflows.

Data is never stored as a pandas ``DataFrame`` field on the dataclass.
Instead, the DataFrame is only provided as an input to:
`from_dataframe` / `from_dataframe_to_dict`.

At build time, the DataFrame is evaluated for its serialized size; if it exceeds
the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
Otherwise, it is inlined as a Temporal-friendly ``dict``.
"""

import pickle
import re
from collections.abc import Hashable
from dataclasses import dataclass
from datetime import datetime
from io import BytesIO
from os import getenv
from typing import Any, Literal

from pandas import DataFrame, read_parquet
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now

# Keys that are part of the serialized wire format (not arbitrary metadata).
_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})

_OBJECT_TIMESTAMP_PATTERN = re.compile(
    r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
)

OFFLOAD_THRESHOLD_BYTES = int(
    float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
)

# Relative prefix used for storing offloaded training datasets in MinIO.
# It is also the root directory for retention cleanup listing.
TRAINING_DATASETS_PREFIX = 'training_datasets'

OperationKind = Literal['initial', 'transform', 'predict']


def _build_object_key(
    model_name: str, operation: OperationKind, timestamp: str
) -> tuple[str, str | None]:
    """
    Build the MinIO object key and the directory prefix used for retention listing.

    Args:
        model_name: Registered model name used in the pipeline.
        operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
        timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.

    Return:
        tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
    """
    # Naming convention:
    #   - Directory is always `training_datasets/<model_name>`
    #   - Filename follows the retention-parsing pattern
    basename = f'{model_name}-{operation}-{timestamp}.parquet'
    model_dir = model_name.strip().strip('/')
    prefix = f'{TRAINING_DATASETS_PREFIX}/{model_dir}'
    return f'{prefix}/{basename}', prefix


@dataclass
class MinioDataFramePayload:
    """
    Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.

    Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
    Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
    """

    last_timestamp: str
    status: dict[str, Any] | None = None
    data: dict[Hashable, Any] | None = None
    bucket: str | None = None
    object_key: str | None = None
    object_prefix: str | None = None
    uri: str | None = None

    @classmethod
    def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload':
        """
        Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.

        Temporal converts dataclass return values into plain dicts when crossing
        workflow/activity boundaries. This method rebuilds the typed instance so
        that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
        available on the receiving side.

        If the argument is already a MinioDataFramePayload, it is returned as-is.

        Args:
            raw: Dict with keys matching the dataclass fields
                (last_timestamp, status, data, bucket, object_key, object_prefix, uri),
                or an existing MinioDataFramePayload instance.

        Return:
            MinioDataFramePayload: Reconstructed (or original) instance.
        """
        if isinstance(raw, MinioDataFramePayload):
            return raw
        return cls(
            last_timestamp=raw['last_timestamp'],
            status=raw.get('status'),
            data=raw.get('data'),
            bucket=raw.get('bucket'),
            object_key=raw.get('object_key'),
            object_prefix=raw.get('object_prefix'),
            uri=raw.get('uri'),
        )

    @staticmethod
    def estimate_size_bytes(df: DataFrame) -> int:
        """
        Approximate serialized size of the DataFrame as the default-orient dict.

        Args:
            df: DataFrame whose tabular content size is estimated.

        Return:
            int: Estimated size in bytes (pickle of dict representation).
        """
        try:
            return len(pickle.dumps(df.to_dict()))
        except Exception:
            return len(pickle.dumps(df))

    @staticmethod
    def parse_object_timestamp(object_key: str) -> datetime | None:
        """
        Parse the timestamp embedded in the object key basename (before .parquet).

        Args:
            object_key: S3/MinIO object key whose basename follows
                ``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.

        Return:
            datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
        """
        basename = object_key.rsplit('/', 1)[-1]
        match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
        if not match:
            return None
        try:
            return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
        except ValueError:
            return None

    def cleanup_prefix(self) -> str | None:
        """
        Return True if cleanup is enabled for this payload.
        """
        if self.object_key is not None and self.data is None:
            return self.object_prefix
        return None

    def has_data(self) -> bool:
        """
        Return True if the payload has some data internally or in MinIO.
        """
        return (self.data is not None and not self.data != {}) or self.object_key is not None

    @classmethod
    async def from_dataframe(
        cls,
        dataframe: DataFrame | None,
        minio_repo: MinioRepository,
        model_name: str,
        operation: OperationKind,
        status: dict[str, Any] | None = None,
        workflow_metadata: dict | None = None,
        last_timestamp: str | None = None,
    ) -> 'MinioDataFramePayload':
        """
        Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.

        The DataFrame is not stored on the returned instance.

        Args:
            dataframe: Tabular data to evaluate and persist (inline or MinIO).
            metadata: Small metadata dict merged into the payload (e.g. success, message).
            minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
            workflow_metadata: Metadata passed to MinIO store for logging/metrics.
            model_name: Registered model name used in the object basename.
            operation: Either ``initial`` (query load) or ``transform`` (post-transform).
            key_prefix: Backward-compatible parameter (currently ignored for object naming).
            size_threshold_bytes: Byte limit before offload. When None, the module-level
                environment-derived default is used.

        Return:
            MinioDataFramePayload: Instance with data and/or MinIO fields set.
        """

        if dataframe is None or dataframe.empty:
            return cls(
                data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
            )


        if last_timestamp is None:
            last_timestamp = max(dataframe['timestamp'].values.tolist())

        if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES:
            return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp)

        timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
        object_key, object_prefix = _build_object_key(model_name, operation, timestamp)

        # Upload using the relative object key. The upstream repository will
        # prefix it internally under its MinIO namespace.
        parquet_buffer = BytesIO()
        dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
        file_bytes = parquet_buffer.getvalue()

        upload_result = await minio_repo.upload_file(
            file_bytes=file_bytes,
            relative_key=object_key,
            metadata=workflow_metadata,
        )

        bucket = minio_repo.bucket
        object_key_full = upload_result.get('minio_object_name', object_key)
        uri = f's3://{bucket}/{object_key_full}' if bucket else None

        return cls(
            data=None,
            bucket=bucket,
            object_key=object_key_full,
            object_prefix=object_prefix,
            uri=uri,
            last_timestamp=last_timestamp,
        )

    async def retrieve(
        self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | None = None
    ) -> DataFrame:
        """
        Load parquet from MinIO when object_key is set and populate inline data.

        Args:
            minio_repo: sientia_do MinioRepository (or compatible) with download_file().
            workflow_metadata: Metadata passed to MinIO read for logging/metrics.

        Return:
            dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
        """
        if self.data is not None:
            return DataFrame(self.data)

        if not self.has_data():
            return DataFrame()

        file_bytes = await minio_repo.download_file(
            object_name=self.object_key, metadata=workflow_metadata
        )
        df = read_parquet(BytesIO(file_bytes))
        return df