from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from typing import Any
    import traceback
    from logging import Logger
    from datetime import datetime
    from pymongo import MongoClient
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.temporal.activities.base import BaseActivity


def clear_mongo_id(docs: list) -> list:
    """
    Remove the MongoDB internal `_id` field from the document.

    Args:
        docs (list): The document to clear.

    Returns:
        list: The documents without the `_id` field.
    """
    for doc in docs:
        if isinstance(doc, list):
            clear_mongo_id(doc)

        elif isinstance(doc, dict):
            if "_id" in doc:
                del doc["_id"]

            for key, value in doc.items():
                if isinstance(value, list):
                    clear_mongo_id(value)
                elif isinstance(value, dict):
                    clear_mongo_id([value])

    return docs


class MongoDB(BaseActivity):
    def __init__(self, connection_string: str, database_name: str,
                 logger: Logger,
                 notification_handler: NotificationHandler):
        self.connection_string = connection_string
        self.database_name = database_name

        self.client = MongoClient(
            self.connection_string, serverSelectionTimeoutMS=5000)
        self.client.server_info()  # Trigger an exception if connection fails

        self.database = self.client[self.database_name]

        # Initialize MongoDB client here (omitted for brevity)
        logger.info("MongoDB connection initialized")

        BaseActivity.__init__(self,
                              logger=logger,
                              notification_handler=notification_handler)

    def shutdown(self):
        """
        Close the MongoDB client connection.
        """
        try:
            if self.client:
                self.logger.info("Closing MongoDB connection...")
                self.client.close()
                self.logger.info("MongoDB connection closed successfully")
        except Exception as e:
            self.logger.error(f"Failed to close MongoDB connection: {e}")

    def __del__(self):
        """
        Destructor to ensure MongoDB client is closed when the object is deleted.
        """
        self.shutdown()

    @activity.defn(name="load_latest_data")
    async def load_latest_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Loads the latest data from MongoDB.
        """
        metadata = input_data['metadata']
        collection_name = input_data['collection_name']
        last_data_timestamp = input_data['last_data_timestamp']

        self.debug(
            f"Loading data from MongoDB: {input_data}",
            metadata=metadata
        )

        try:

            data_filter = {
                "inserted_at": {
                    "$gt": datetime.strptime(last_data_timestamp, "%Y-%m-%d %H:%M:%S.%f")
                }
            }

            data = self.database[collection_name].find(data_filter, {"_id": 0})

            data = clear_mongo_id(data)

            for item in data:
                item['inserted_at'] = item['inserted_at'].isoformat()

            self.info(
                f"Loaded {len(data)} documents from MongoDB",
                metadata=metadata
            )

            self.debug(
                f"Loaded data: {data}",
                metadata=metadata
            )

            return data
        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                metadata=metadata,
                notification_id="MONGO_LOAD_ERROR",
                message=f"Error loading data from MongoDB: {e}",
                block="load_latest_data",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            raise e