from typing import Any, List, Optional, Dict
from pymongo import MongoClient
from pymongo.errors import OperationFailure, ServerSelectionTimeoutError, ConnectionFailure, ConfigurationError
from analogai.foundation.infrastructure.storage.storage_adapter import StorageAdapter
from analogai.foundation.settings import config
from analogai.foundation.common.logging.app_logger import app_logger
import uuid
from typeguard import typechecked
import atexit
from analogai.foundation.common.design_patterns.singleton import Singleton
import time

class MongoDBStorageAdapterImpl(StorageAdapter, metaclass=Singleton):
    def __init__(self):
        app_logger.info("initializing mongo")
        self._validate_configuration()
        self.uri = config.get("MONGODB_URI")
        self.database_name = config.get("MONGODB_DATABASE")
        self.client: Optional[MongoClient] = None
        self.db = None
        self.max_retries = 3
        self.retry_delay = 2
        self._initialize_connection()
        atexit.register(self.close)
    
    def _validate_configuration(self):

        missing_configs = []
        
        if not config.get("MONGODB_URI"):
            missing_configs.append("MONGODB_URI")
        if not config.get("MONGODB_DATABASE"):
            missing_configs.append("MONGODB_DATABASE")
        
        if missing_configs:
            error_msg = f"MongoDB configuration is missing: {', '.join(missing_configs)}. Please check your environment variables."
            app_logger.error(error_msg)
            raise ValueError(error_msg)
        
        if not str(config.get("MONGODB_URI", "")).startswith(('mongodb://', 'mongodb+srv://')):
            error_msg = f"Invalid MongoDB URI format: {config.get('MONGODB_URI')}. Must start with mongodb:// or mongodb+srv://"
            app_logger.error(error_msg)
            raise ValueError(error_msg)
    
    def _initialize_connection(self):

        app_logger.debug(f"Attempting MongoDB connection with URI: {self.uri}")
        
        for attempt in range(self.max_retries):
            try:
                self._create_client()
                self._test_connection()
                app_logger.info(f"MongoDB connection established successfully at {self.uri}")
                return
                
            except ServerSelectionTimeoutError as e:
                error_msg = self._format_connection_error(e, attempt)
                if attempt < self.max_retries - 1:
                    app_logger.warning(f"{error_msg} Retrying in {self.retry_delay} seconds... (attempt {attempt + 1}/{self.max_retries})")
                    time.sleep(self.retry_delay)
                else:
                    app_logger.error(f"{error_msg} All retry attempts failed.")
                    raise ConnectionError(f"Failed to connect to MongoDB after {self.max_retries} attempts: {error_msg}")
                    
            except Exception as e:
                if "authentication" in str(e).lower() or "auth" in str(e).lower():
                    error_msg = f"MongoDB authentication failed. Check username and password. Error: {str(e)}"
                    app_logger.error(error_msg)
                    raise ConnectionError(error_msg)
                
            except ConfigurationError as e:
                error_msg = f"MongoDB configuration error: {str(e)}"
                app_logger.error(error_msg)
                raise ValueError(error_msg)
                
            except ConnectionFailure as e:
                error_msg = f"MongoDB connection failed: {str(e)}"
                app_logger.error(error_msg)
                if attempt < self.max_retries - 1:
                    app_logger.warning(f"Retrying in {self.retry_delay} seconds... (attempt {attempt + 1}/{self.max_retries})")
                    time.sleep(self.retry_delay)
                else:
                    raise ConnectionError(error_msg)
                    
            except Exception as e:
                error_msg = f"Unexpected error connecting to MongoDB: {str(e)}"
                app_logger.error(error_msg)
                if attempt < self.max_retries - 1:
                    app_logger.warning(f"Retrying in {self.retry_delay} seconds... (attempt {attempt + 1}/{self.max_retries})")
                    time.sleep(self.retry_delay)
                else:
                    raise ConnectionError(error_msg)
    
    def _create_client(self):

        try:
            self.client = MongoClient(
                self.uri,
                serverSelectionTimeoutMS=5000,
                connectTimeoutMS=5000,
                socketTimeoutMS=5000,
                maxPoolSize=50,
                retryWrites=True
            )
            self.db = self.client[self.database_name]
        except Exception as e:
            app_logger.error(f"Failed to create MongoDB client: {str(e)}")
            raise
    
    def _test_connection(self):

        try:
            self.client.admin.command('ping')
            app_logger.debug("MongoDB connection test successful")
        except Exception as e:
            app_logger.error(f"MongoDB connection test failed: {str(e)}")
            raise
    
    def _format_connection_error(self, error: Exception, attempt: int) -> str:

        error_str = str(error)
        
        if "ConnectionRefusedError" in error_str or "actively refused" in error_str:
            return f"MongoDB server is not running or not accessible at {self.uri}. Please ensure MongoDB is started and accessible."
        elif "timeout" in error_str.lower():
            return f"MongoDB connection timeout. The server can be overloaded or network issues exist. URI: {self.uri}"
        elif "authentication" in error_str.lower():
            return f"MongoDB authentication failed. Check credentials in URI: {self.uri}"
        else:
            return f"MongoDB connection failed: {error_str}"

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def close(self) -> None:

        if hasattr(self, 'client') and self.client:
            try:
                self.client.close()
                app_logger.info("MongoDB client connection closed successfully")
                self.client = None
                self.db = None
                if MongoDBStorageAdapterImpl in Singleton._instances:
                    Singleton._instances.pop(MongoDBStorageAdapterImpl, None)
            except Exception as e:
                app_logger.error(f"Error closing MongoDB client: {str(e)}")
                self.client = None
                self.db = None

    def is_connected(self) -> bool:

        if self.client is None:
            return False
        
        try:
            self.client.admin.command('ping')
            return True
        except Exception:
            return False
    
    def get_connection_info(self) -> dict:

        return {
            "uri": self.uri,
            "database": self.database_name,
            "is_connected": self.is_connected(),
            "client_active": self.client is not None
        }

    def __del__(self):

        try:
            self.close()
        except Exception:
            pass  # Ignore errors during cleanup

    def _normalize_document(self, doc: dict) -> dict:
        if not doc:
            return doc
        
        normalized = {}
        for key, value in doc.items():
            if key == '_id':
                continue
            
            if isinstance(value, int) and key.endswith('_id'):
                normalized[key] = str(value)
            elif isinstance(value, list):
                normalized[key] = [self._normalize_value(item) for item in value]
            elif isinstance(value, dict):
                normalized[key] = self._normalize_value(value)
            else:
                normalized[key] = value
        
        return normalized

    def _normalize_value(self, value):
        if isinstance(value, dict):
            normalized = {}
            for k, v in value.items():
                if isinstance(v, int) and k.endswith('_id'):
                    normalized[k] = str(v)
                elif isinstance(v, list):
                    normalized[k] = [self._normalize_value(item) for item in v]
                elif isinstance(v, dict):
                    normalized[k] = self._normalize_value(v)
                else:
                    normalized[k] = v
            return normalized
        elif isinstance(value, list):
            return [self._normalize_value(item) for item in value]
        else:
            return value

    @typechecked
    def store(self, collection: str, data: dict) -> dict:

        if not collection:
            raise ValueError("Collection name cannot be empty")
        if not data:
            raise ValueError("Data cannot be empty")
        
        try:
            self._ensure_connection()
            result = self.db[collection].replace_one({'id': data['id']}, data, upsert=True)
            app_logger.debug(f"Successfully stored document in mongodb collection {collection} (upserted: {result.upserted_id is not None})")
            stored = self.db[collection].find_one({'id': data['id']})
            return stored
        except OperationFailure as e:
            error_msg = f"Failed to store document in mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
        except Exception as e:
            error_msg = f"Unexpected error storing document in mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)

    @typechecked
    def retrieve_by_id(self, collection: str, id: str) -> Optional[Any]:
        if not collection:
            raise ValueError("Collection name cannot be empty")
            
        app_logger.debug(f"Retrieving document with id: {id} from mongodb collection: {collection}")
        try:
            self._ensure_connection()
            result = self.db[collection].find_one({'id': id})
            app_logger.debug(f"Raw result from mongodb: {result}")
            if result:
                result = self._normalize_document(result)
            return result if result else None
        except OperationFailure as e:
            error_msg = f"Failed to retrieve document from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
        except Exception as e:
            error_msg = f"Unexpected error retrieving document from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)

    @typechecked
    def retrieve_all(self, collection: str, query: Optional[dict] = None) -> List[Any]:

        if not collection:
            raise ValueError("Collection name cannot be empty")
            
        find_query = {}
        if query:
            find_query.update(query)
        try:
            self._ensure_connection()
            return [self._normalize_document(item) for item in self.db[collection].find(find_query)]
        except OperationFailure as e:
            error_msg = f"Failed to retrieve all documents from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
        except Exception as e:
            error_msg = f"Unexpected error retrieving all documents from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)

    @typechecked
    def retrieve_by_attributes(self, collection: str, query: Dict[str, str]) -> List[Any]:

        if not collection:
            raise ValueError("Collection name cannot be empty")
        if not query:
            raise ValueError("Query cannot be empty")
            
        try:
            self._ensure_connection()
            return [self._normalize_document(item) for item in self.db[collection].find(query)]
        except OperationFailure as e:
            error_msg = f"Failed to retrieve documents by attributes from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
        except Exception as e:
            error_msg = f"Unexpected error retrieving documents by attributes from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)

    @typechecked
    def delete(self, collection: str, id: str) -> bool:
        if not collection:
            raise ValueError("Collection name cannot be empty")
            
        try:
            self._ensure_connection()
            result = self.db[collection].delete_one({'id': id})
            if result.deleted_count == 0:
                app_logger.warning(f"Document with id {id} not found in collection {collection}")
                return False
            app_logger.debug(f"Successfully deleted document with id {id} from mongodb collection {collection}")
            return True
        except Exception as e:
            error_msg = f"Unexpected error deleting document from mongodb collection {collection}: {str(e)}"
            app_logger.error(error_msg)
            return False

    def clear(self) -> None:

        try:
            self._ensure_connection()
            collections = self.db.list_collection_names()
            if not collections:
                app_logger.info("No collections found to clear")
                return
                
            for collection in collections:
                result = self.db[collection].delete_many({})
                app_logger.debug(f"Cleared {result.deleted_count} documents from collection {collection}")
            app_logger.info(f"Successfully cleared all collections: {collections}")
        except OperationFailure as e:
            error_msg = f"Failed to clear mongodb collections: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
        except Exception as e:
            error_msg = f"Unexpected error clearing mongodb collections: {str(e)}"
            app_logger.error(error_msg)
            raise ConnectionError(error_msg)
    
    def _ensure_connection(self):

        if self.client is None or not self.is_connected():
            app_logger.warning("MongoDB connection lost, attempting to reconnect")
            self._initialize_connection()

