import time
import functools

from pymongo import MongoClient
from pymongo.errors import PyMongoError
from Controllers.Mongo.config import mongo_configs


def retry_operation(max_attempts=3, delay=1.0, backoff=2.0, exceptions=(PyMongoError,)):
    """
    Decorator for retrying MongoDB operations with exponential backoff.

    Args:
        max_attempts: Maximum number of retry attempts
        delay: Initial delay between retries in seconds
        backoff: Multiplier for delay after each retry
        exceptions: Tuple of exceptions to catch and retry
    """

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            mtries, mdelay = max_attempts, delay
            while mtries > 1:
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    time.sleep(mdelay)
                    mtries -= 1
                    mdelay *= backoff
            return func(*args, **kwargs)

        return wrapper

    return decorator


class MongoDBConfig:
    """
    Configuration class for MongoDB connection settings.
    """

    def __init__(
        self,
        uri: str = "mongodb://localhost:27017/",
        max_pool_size: int = 20,
        min_pool_size: int = 10,
        max_idle_time_ms: int = 30000,
        wait_queue_timeout_ms: int = 2000,
        server_selection_timeout_ms: int = 5000,
        **additional_options,
    ):
        """
        Initialize MongoDB configuration.
        """
        self.uri = uri
        self.client_options = {
            "maxPoolSize": max_pool_size,
            "minPoolSize": min_pool_size,
            "maxIdleTimeMS": max_idle_time_ms,
            "waitQueueTimeoutMS": wait_queue_timeout_ms,
            "serverSelectionTimeoutMS": server_selection_timeout_ms,
            **additional_options,
        }


class MongoDBHandler(MongoDBConfig):
    """
    A MongoDB handler that provides context manager access to specific collections
    with automatic retry capability.
    """

    _instance = None

    def __new__(cls, *args, **kwargs):
        """
        Implement singleton pattern for the handler.
        """
        if cls._instance is None:
            cls._instance = super(MongoDBHandler, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(
        self,
        uri: str,
        max_pool_size: int = 5,
        min_pool_size: int = 2,
        max_idle_time_ms: int = 10000,
        wait_queue_timeout_ms: int = 1000,
        server_selection_timeout_ms: int = 3000,
        **additional_options,
    ):
        """
        Initialize the MongoDB handler (only happens once due to singleton).
        """
        # Only initialize once
        if not hasattr(self, "_initialized") or not self._initialized:
            super().__init__(
                uri=uri,
                max_pool_size=max_pool_size,
                min_pool_size=min_pool_size,
                max_idle_time_ms=max_idle_time_ms,
                wait_queue_timeout_ms=wait_queue_timeout_ms,
                server_selection_timeout_ms=server_selection_timeout_ms,
                **additional_options,
            )
            self._initialized = True

    def collection(self, collection_name: str):
        """
        Get a context manager for a specific collection.

        Args:
            collection_name: Name of the collection to access

        Returns:
            A context manager for the specified collection
        """
        return CollectionContext(self, collection_name)


class CollectionContext:
    """
    Context manager for MongoDB collections with automatic retry capability.
    """

    def __init__(self, db_handler: MongoDBHandler, collection_name: str):
        """
        Initialize collection context.

        Args:
            db_handler: Reference to the MongoDB handler
            collection_name: Name of the collection to access
        """
        self.db_handler = db_handler
        self.collection_name = collection_name
        self.client = None
        self.collection = None

    def __enter__(self):
        """
        Enter context, establishing a new connection.

        Returns:
            The MongoDB collection object with retry capabilities
        """
        try:
            # Create a new client connection
            self.client = MongoClient(
                self.db_handler.uri, **self.db_handler.client_options
            )
            # Get database from URI
            db_name = self.client.get_database().name
            self.collection = self.client[db_name][self.collection_name]

            # Enhance collection methods with retry capabilities
            self._add_retry_capabilities()

            return self.collection
        except Exception as e:
            if self.client:
                self.client.close()
            raise

    def _add_retry_capabilities(self):
        """
        Add retry capabilities to collection methods.
        """
        # Store original methods
        original_insert_one = self.collection.insert_one
        original_insert_many = self.collection.insert_many
        original_find_one = self.collection.find_one
        original_find = self.collection.find
        original_update_one = self.collection.update_one
        original_update_many = self.collection.update_many
        original_delete_one = self.collection.delete_one
        original_delete_many = self.collection.delete_many
        original_replace_one = self.collection.replace_one
        original_count_documents = self.collection.count_documents

        # Add retry capabilities to methods
        self.collection.insert_one = retry_operation()(original_insert_one)
        self.collection.insert_many = retry_operation()(original_insert_many)
        self.collection.find_one = retry_operation()(original_find_one)
        self.collection.find = retry_operation()(original_find)
        self.collection.update_one = retry_operation()(original_update_one)
        self.collection.update_many = retry_operation()(original_update_many)
        self.collection.delete_one = retry_operation()(original_delete_one)
        self.collection.delete_many = retry_operation()(original_delete_many)
        self.collection.replace_one = retry_operation()(original_replace_one)
        self.collection.count_documents = retry_operation()(original_count_documents)

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Exit context, closing the connection.
        """
        if self.client:
            self.client.close()
            self.client = None
            self.collection = None


mongo_handler = MongoDBHandler(
    uri=mongo_configs.url,
    max_pool_size=5,
    min_pool_size=2,
    max_idle_time_ms=30000,
    wait_queue_timeout_ms=2000,
    server_selection_timeout_ms=5000,
)