from actions_logging.app_logging import logger
import pymongo


class MongoClientManager:
    """
    Context manager for managing MongoDB client connections.
    """
    def __init__(self, mongo_uri: str):
        self.mongo_uri = mongo_uri
        self.client = None

    def __enter__(self):
        try:
            self.client = pymongo.MongoClient(self.mongo_uri)
            self.client.admin.command('ping')  # Test the connection
            logger.info_green("MongoDB connection established.")
            return self.client
        except pymongo.errors.PyMongoError as e:
            raise Exception(f"Failed to connect to MongoDB: {e}")

    def __exit__(self, exc_type, exc_value, traceback):
        if self.client:
            self.client.close()
            logger.info("MongoDB connection closed.")


def mongo_query(client, query: dict, db_name: str, collection_name: str) -> list:
    """
    Execute a MongoDB query and return the results.
    """
    try:
        db = client[db_name]
        collection = db[collection_name]
        results = collection.find(query)
        return list(results)
    except pymongo.errors.PyMongoError as e:
        raise Exception(f"MongoDB query failed: {e}")
    except Exception as e:
        raise Exception(f"An unexpected error occurred: {e}")


def count_documents(client, query: dict, db_name: str, collection_name: str) -> int:
    """
    Execute a MongoDB query and return the count of matching documents.
    """
    try:
        db = client[db_name]
        collection = db[collection_name]
        count = collection.count_documents(query)
        return count
    except pymongo.errors.PyMongoError as e:
        raise Exception(f"MongoDB query failed: {e}")
    except Exception as e:
        raise Exception(f"An unexpected error occurred: {e}")