from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    # Extend the Temporal Postgres activities for convenient query -> MinIO export
    import traceback
    from datetime import timedelta
    from typing import Any

    import pandas as pd
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.repository.minio_repository_sync import MinioRepository
    from sientia_do.temporal.activities.postgres_sync import Postgres
    from sientia_do.temporal.constants import now

    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload

_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})


class Storage(Postgres, SientiaMonitoring):
    """
    Extensions for Postgres activities with a helper to export query results
    directly to MinIO as Parquet and return the object name.
    """

    minio_repository: MinioRepository | None = None

    def __init__(
        self,
        host: str,
        port: int,
        user: str,
        password: str,
        dbname: str,
        min_connections: int,
        max_connections: int,
        retention_hours: int = 24,
        minio_repository: MinioRepository | None = None,
        logger: Logger | None = None,
        notification_handler: NotificationHandler | None = None,
        metrics_controller: MetricsController | None = None,
    ):
        self.retention_hours = retention_hours
        Postgres.__init__(
            self,
            host=host,
            port=port,
            user=user,
            password=password,
            dbname=dbname,
            min_connections=min_connections,
            max_connections=max_connections,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        self.minio_repository = minio_repository
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

    @activity.defn(name='load_query_with_minio_offload')
    def load_query_with_minio_offload(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
        """
        Run the custom SQL load, then return a MinIO-aware dataframe wire dict.

        Args (input_data):
            metadata (dict): Workflow metadata (same as load_custom_query).
            query (str): SQL query.
            datetime_columns (list[str], optional): Datetime column names.
            model_name (str): Model name for object key basename.
            key_prefix (str, optional): Directory prefix inside the bucket.
            size_threshold_bytes (int, optional): Override env offload threshold.

        Returns:
            dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
        """
        if self.minio_repository is None:
            raise ValueError('Minio repository not initialized')

        metadata: dict = input_data.get('metadata', {})
        model_name = input_data['model_name']

        rows = self.load_custom_query(
            input_data,
        )
        if not rows:
            self.error(
                'load_query_with_minio_offload failed: No data returned from query', metadata
            )
            dataframe = None
        else:
            dataframe = pd.DataFrame(rows)

        return MinioDataFramePayload.from_dataframe(
            dataframe,
            minio_repo=self.minio_repository,
            workflow_metadata=metadata,
            model_name=model_name,
            operation='initial',
            logger=self.logger,
        )

    @activity.defn(name='export_payload_to_postgres')
    def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
        """
        Resolve a MinIO-aware payload into a DataFrame and persist it into PostgreSQL.

        This activity accepts the serialized payload produced by previous steps
        (inline dict or MinIO object reference), reconstructs the tabular data,
        and delegates the final write to ``export_data_to_postgres`` using the
        same input contract expected by the Postgres activity mixin.

        Args:
         - input_data (dict[str, Any]): Activity input containing ``data`` as a
           ``MinioDataFramePayload``-compatible dict plus database write options
           (schema/table/on_conflict/metadata and related fields).

        Return:
         dict: Result dictionary returned by ``export_data_to_postgres``, including
         success status and optional write diagnostics.
        """
        metadata = input_data.get('metadata')
        payload = MinioDataFramePayload.from_dict(input_data['data'])
        data = payload.retrieve(self.minio_repository, metadata)

        return self.export_data_to_postgres(
            {
                **input_data,
                'data': data,
            }
        )

    @activity.defn(name='cleanup_minio_objects_expired')
    def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Delete objects under the given prefixes that are older than the retention window.

        Args (input_data):
            metadata (dict): Workflow metadata for logging and metrics.
            prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).

        Returns:
            dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
        """
        if self.minio_repository is None:
            raise ValueError('Minio repository not initialized')

        metadata = input_data.get('metadata', {})
        payload = MinioDataFramePayload.from_dict(input_data['data'])
        prefix = payload.cleanup_prefix()
        base = now()
        cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
            hours=self.retention_hours
        )

        report: dict[str, Any] = {
            'failed': {},
            'deleted': {},
            'failed_count': 0,
            'deleted_count': 0,
        }
        try:
            keys = self.minio_repository.list_objects(
                prefix=prefix,
                recursive=True,
                metadata=metadata,
            )
            for key in keys:
                try:
                    ts = MinioDataFramePayload.parse_object_timestamp(key)
                    if ts is None:
                        continue
                    if ts >= cutoff:
                        continue
                    self.minio_repository.delete_file(
                        object_name=key,
                        metadata=metadata,
                    )
                except Exception as e:
                    report['failed'][key] = {
                        'success': False,
                        'message': str(e),
                    }
                    report['failed_count'] += 1
                    continue
                report['deleted'][key] = {
                    'success': True,
                    'message': 'Deleted',
                }
                report['deleted_count'] += 1
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
                message=f'Error cleaning up MinIO objects: {e}',
                block='cleanup_minio_objects_expired',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )
            self.error(trace, metadata)
        else:
            # Cleanup success is expected in normal flow; avoid noisy INFO notifications
            # that do not impact behavior and can flood observability in test runs.
            self.info('MinIO objects cleaned up successfully', metadata)

        return report

    def close(self) -> None:
        """
        Shutdown Storage resources in deterministic order.

        The method first closes Postgres resources via ``Postgres.close`` (engine,
        sessions, and monitoring hooks), then closes the optional MinIO repository
        and clears the local reference to avoid accidental reuse after shutdown.
        """
        Postgres.close(self)
        if self.minio_repository is not None:
            try:
                self.minio_repository.close()
            finally:
                self.minio_repository = None