from temporalio import activity, workflow

from laborious.utils.repository.minio_repository import MinioRepository


with workflow.unsafe.imports_passed_through():
    # Extend the Temporal Postgres activities for convenient query -> MinIO export
    from sientia_do.temporal.activities.postgres import Postgres
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import Logger
    from sientia_do.temporal.constants import now
    from sientia_do.notifications.models import NotificationLevel
    from typing import Any
    import traceback
    import pandas as pd

DATETIME_FILENAME_FORMAT = "%Y-%m-%d_%H-%M-%S"


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

    def __init__(self,
                 host: str,
                 port: int,
                 user: str,
                 password: str,
                 dbname: str,
                 min_connections: int,
                 max_connections: int,
                 minio_config: dict[str, Any],
                 logger: Logger,
                 notification_handler: NotificationHandler):
        super().__init__(host=host,
                         port=port,
                         user=user,
                         password=password,
                         dbname=dbname,
                         min_connections=min_connections,
                         max_connections=max_connections,
                         logger=logger,
                         notification_handler=notification_handler)

        if not hasattr(self, 'minio_repository'):
            self.minio_repository = MinioRepository(
                logger=logger,
                notification_handler=notification_handler,
                minio_endpoint_url=minio_config['endpoint_url'],
                minio_access_key=minio_config['access_key'],
                minio_secret_key=minio_config['secret_key'],
                minio_region_name=minio_config['region_name'],
                minio_default_bucket=minio_config['default_bucket'])

        if self.minio_repository is None:
            self.minio_repository = MinioRepository(
                logger=logger,
                notification_handler=notification_handler,
                minio_endpoint_url=minio_config['endpoint_url'],
                minio_access_key=minio_config['access_key'],
                minio_secret_key=minio_config['secret_key'],
                minio_region_name=minio_config['region_name'],
                minio_default_bucket=minio_config['default_bucket'])

    @activity.defn(name='query_to_minio')
    async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Execute SQL query, write result as Parquet to MinIO, and return object name.

        Args (input_data):
            metadata (dict): Workflow metadata
            query (str): SQL query
            model_name (str): Model name for object naming
            object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain)

        Returns:
            dict: { success: bool, object_name: str, uri: str }
        """

        metadata = input_data.get('metadata', {})
        object_prefix = input_data.get('object_prefix', 'datasets/retrain')

        timestamp = now().strftime(DATETIME_FILENAME_FORMAT)
        object_name = f"{object_prefix}_{timestamp}.parquet"
        uri = f"s3://{self.minio_repository.minio_bucket}/{object_name}"

        try:
            data = await self.load_custom_query(input_data)
            if not data:
                self.error(
                    f"query_to_minio failed: No data returned from query", metadata)
                return {"success": False, "message": "No data returned from query"}

            # Ensure we have a DataFrame
            data = pd.DataFrame(data)

            # Write parquet to memory and upload via persistent client
            self.minio_repository.store_dataframe_as_parquet(
                dataframe=data,
                uri=uri,
                object_name=object_name,
                metadata=metadata
            )

            return {"success": True, "object_key": object_name, "uri": uri}
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id="ERROR_LOADING_CUSTOM_QUERY",
                message=f"Error fetching data from query: {e}",
                block="load_custom_query",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )

            self.error(trace, metadata)

            return {"success": False, "message": str(e)}

    def close(self) -> None:
        """Close Storage resources (MinIO client and Postgres engine)."""
        try:
            if hasattr(self, 's3_client') and self.s3_client is not None:
                try:
                    self.s3_client.close()
                finally:
                    self.s3_client = None
        finally:
            # Ensure Postgres resources are disposed as well
            try:
                super().close()
            except Exception:
                pass

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass