from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    # Extend the Temporal Postgres activities for convenient query -> MinIO export
    import traceback
    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.temporal.activities.postgres import Postgres
    from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now

    from laborious.utils.repository.minio_repository import MinioRepository


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 | None = None

        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 }
        """

        if self.minio_repository is None:
            raise ValueError('Minio repository not initialized')

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

        timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
        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('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_STORING_QUERY_TO_MINIO',
                message=f'Error storing query to MinIO: {e}',
                block='query_to_minio',
                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, 'minio_repository') and self.minio_repository is not None:
                try:
                    self.minio_repository.close()
                finally:
                    self.minio_repository = None
        finally:
            # Ensure Postgres resources are disposed as well
            try:
                super().close()
            except Exception:
                self.logger.error('Error closing Postgres resources')

    def __del__(self):
        self.close()