"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Tuesday, 28th May, 2024

    OBJECTIVE:

        To have one central place from where all async database connectivity happens.

    REFERENCES:

        01. https://motor.readthedocs.io/en/stable/

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# For async behaviour:
import asyncio

# For datetime handling:
import pytz
import datetime

# MongoDB:
from motor.motor_asyncio import AsyncIOMotorClient
from bson.objectid import ObjectId
from bson.json_util import dumps, loads

# For debugging:
from icecream import IceCreamDebugger


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                EXCEPTIONS                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class MongoFindException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo find operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


# ---------------------------------------------------------------------------------------------------------------------


class MongoInsertException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo insert operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


# ---------------------------------------------------------------------------------------------------------------------


class MongoUpdateException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo update operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


# ---------------------------------------------------------------------------------------------------------------------


class MongoReplaceException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo replace operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


# ---------------------------------------------------------------------------------------------------------------------


class MongoDeleteException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo delete operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


# ---------------------------------------------------------------------------------------------------------------------


class MongoException(Exception):

    def __init__(self, hint = None, origin = None):
        self.__hint = hint
        self.__origin = origin

    def __str__(self):
        message = "mongo operation failed"
        if self.__hint is not None: message = f"{message} ({self.__hint})"
        if self.__origin is not None: message = f"{self.__origin} --> {message}"
        return message


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class AsyncMongo:

    __db = None
    __db_name = None
    __client = None

    def __init__(
            self,
            database_name = "myDb",
            max_connections = 5,
            debug = True,
            debug_only_errors = True,
            host = "localhost",
            port = 27017,
            connection_string = None
    ):

        # Basic config:
        self.__max_connections = max_connections
        self.__db_name = database_name
        self.__host = host,
        self.__port = port
        self.__connection_string = connection_string

        # For debugging:
        self.__debug_only_errors = debug_only_errors
        self.__printer = IceCreamDebugger(prefix = f"Mongo ({self.__db_name}) | ", includeContext = True)
        if not debug: self.__printer.disable()

        # Rate/access control:
        self.__exclusive_semaphore = asyncio.Semaphore(1)

    def enable_debug(self):
        self.__printer.enable()

    def disable_debug(self):
        self.__printer.disable()

    @staticmethod
    def generate_id():

        """
        Just generates an '_id' in MongoDB style.
        :return: The '_id' in MongoDB style.
        """

        return str(ObjectId())

    async def connect(self):
        try:
            if self.__connection_string is None:
                self.__client = AsyncIOMotorClient(
                    self.__host,
                    self.__port,
                    maxPoolSize = self.__max_connections,
                    minPoolSize = self.__max_connections
                )
            else:
                self.__client = AsyncIOMotorClient(
                    self.__connection_string,
                    maxPoolSize = self.__max_connections,
                    minPoolSize = self.__max_connections
                )
            self.__db = self.__client.get_database(self.__db_name)
            if not self.__debug_only_errors:
                server_info = await self.__client.server_info()
                self.__printer(server_info)
        except Exception as exception: self.__printer(exception)

    async def ensure_connection(self):
        if self.__db is None:
            async with self.__exclusive_semaphore:
                await self.connect()

    @property
    async def client(self):
        await self.ensure_connection()
        return self.__client

    @staticmethod
    def dict_to_dot_notation(input_dict, parent_key = "", separator = "."):
        items = []
        for k, v in input_dict.items():
            new_key = f"{parent_key}{separator}{k}" if parent_key else k
            if isinstance(v, dict) and v:
                items.extend(AsyncMongo.dict_to_dot_notation(v, new_key, separator = separator).items())
            else:
                items.append((new_key, v))
        return dict(items)

    @staticmethod
    def normalize_date_time(document):

        """
        MongoDB doesn't support timezones. A good strategy would be to convert everything to UTC format and store it.
        This method does exactly that. Any datetime object is converted to UTC timezone. If the datetime object was
        timezone naive, UTC timezone will be applied to it without changing the time value.
        :param document: The document that you want to normalize the date-time in.
        :return: The document with normalized datetime.
        """

        if isinstance(document, datetime.datetime):
            utc_tz = pytz.timezone("UTC")
            if document.tzinfo is None: document = utc_tz.localize(document)
            else: document = document.astimezone(utc_tz)

        if type(document) is list:
            document = [AsyncMongo.normalize_date_time(item) for item in document]

        if type(document) is dict:
            document = {
                AsyncMongo.normalize_date_time(k): AsyncMongo.normalize_date_time(v)
                for k, v in document.items()
            }

        return document

    @staticmethod
    def __from_json_string(json_data):
        return loads(json_data)

    @staticmethod
    def __to_json_string(python_data, indent = 4, default = None):
        return dumps(python_data, indent = indent, default = default)

    async def list_indexes(
            self,
            collection,
            session = None,
            raise_exception = False
    ):

        """
        Lists out the indexes of a collection.
        :param collection: The collection whose indexes you want to list out,
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The list of indexes or None if the action fails.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        indexes = None

        # Try to list the indexes:
        try:
            responses = await self.__db[collection].list_indexes(session = session).to_list(None)
            indexes = [{key: value for key, value in response.items()} for response in responses]
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not indexes: raise MongoException(hint = "list indexes")
        return indexes

    async def create_index(
            self,
            collection,
            keys,
            options = None,
            session = None,
            raise_exception = False
    ):

        """
        Creates an index on a collection.
        :param collection: The collection to create the index on.
        :param keys: The keys (and sorting) to implement the index on.
        :param options: Additional config.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the execution.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        success = False

        # Try to make the insertion:
        try:
            options = options or {}
            keys = [(k, v) for k, v in keys.items()]
            response = await self.__db[collection].create_index(keys, session = session, **options)
            if response: success = True
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not success: raise MongoException(hint = "create index")
        return success

    async def insert_one(
            self,
            collection,
            document,
            session = None,
            raise_exception = False
    ):

        """
        Insert data into a collection.
        :param collection: The collection you want to feed the data into.
        :param document: The data to be stored.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The id of the inserted data, or null if the action fails.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        inserted_id = None

        # Try to make the insertion:
        try:
            response = await self.__db[collection].insert_one(document.copy(), session = session)
            inserted_id = response.inserted_id
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not inserted_id: raise MongoInsertException(hint = f"{collection}")
        return inserted_id

    async def insert_many(
            self,
            collection,
            documents,
            session = None,
            raise_exception = False
    ):

        """
        Insert a lot of data into a collection.
        :param collection: The collection you want to feed the data into.
        :param documents: The data to be stored.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The id of the inserted data, or null if the action fails.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        inserted_ids = []

        # Try to make the insertion:
        try:
            response = await self.__db[collection].insert_many(documents, session = session)
            inserted_ids = response.inserted_ids
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not inserted_ids: raise MongoInsertException(hint = f"{collection}")
        return inserted_ids

    async def update_one(
            self,
            collection,
            filter_json = None,
            update_json = None,
            upsert = False,
            session = None,
            raise_exception = False
    ):

        """
        Update one document.
        :param collection: The collection you want to update.
        :param filter_json: The selection criteria to locate the document to update.
        :param update_json: The values you want to update.
        :param upsert: If you want to insert if the document doesn't already exist.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the operation.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        success = False

        # Try to make the insertion:
        try:
            response = await self.__db[collection].update_one(
                filter_json,
                update_json,
                upsert = upsert,
                session = session
            )
            success = False if response.modified_count == 0 else True
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not success: raise MongoUpdateException(hint = f"{collection}")
        return success

    async def update_many(
            self,
            collection,
            filter_json = None,
            update_json = None,
            upsert = False,
            session = None,
            raise_exception = False
    ):

        """
        Update many documents.
        :param collection: The collection you want to update.
        :param filter_json: The selection criteria to locate the document to update.
        :param update_json: The values you want to update.
        :param upsert: If you want to insert if the document doesn't already exist.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the operation.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        update_count = 0

        # Try to make the insertion:
        try:
            response = await self.__db[collection].update_many(
                filter_json,
                update_json,
                upsert = upsert,
                session = session
            )
            update_count = response.modified_count
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not update_count: raise MongoUpdateException(hint = f"{collection}")
        return update_count

    async def count(
            self,
            collection,
            filter_json = None
    ):

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        count = 0

        # try to query the data:
        try:
            if filter_json is None: filter_json = {}
            count = await self.__db[collection].count_documents(filter_json)
        except Exception as exception: self.__printer(exception)

        # Done here:
        return count

    async def bulk_write(
            self,
            collection,
            requests,
            session = None,
            raise_exception = False
    ):

        """
        To perform various individual operations in one go. You will have to import individual actions like "UpdateOne"
        and "InsertMany" from PyMongo and pass them as an array of requests (operations) to this method.
        :param collection: The collection you want to run the requests on.
        :param requests: The array of requests (operations) to be performed.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return:
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        count = 0

        # try to query the data:
        try:
            response = await self.__db[collection].bulk_write(requests, session = session)
            count = response.modified_count + response.inserted_count + response.upserted_count + response.deleted_count
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not count: raise MongoException(hint = "bulk write")
        return count

    async def find_many(
            self,
            collection,
            filter_json,
            projections = None,
            skip = 0,
            limit = None,
            sort = None,
            session = None,
            as_json_string = False,
            indent = 4,
            raise_exception = False
    ):

        """
        Finds one or more records that match the given conditions.
        :param collection: The name of the collection to perform the search in.
        :param filter_json: The filter criteria.
        :param projections: What parts of the matching data you want to fetch.
        :param skip: The no. of records to skip before picking next ones. Needed for pagination.
        :param limit: The max. no. of records you want to fetch.
        :param sort: The sorting rules to apply.
        :param session: The session if you need to do this in a transaction.
        :param as_json_string: Whether you want it as a JSON string or a Python dict/list.
        :param indent: The indentation to use if you want it as a JSON string.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The array of matching records or null if there was an exception.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        results = None

        # Assume defaults:
        if sort is None: sort = {"_id": -1}
        if limit is None: limit = 10

        # try to query the data:
        try:
            results = await self.__db[collection].find(
                filter_json,
                projections,
                session = session
            ).sort(sort).skip(skip).limit(limit).to_list(None)
            if as_json_string: results = self.__to_json_string(results, indent = indent, default = str)
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not results: raise MongoFindException(hint = f"{collection}")
        return results

    async def find_one(
            self,
            collection,
            filter_json,
            projections = None,
            session = None,
            as_json_string = False,
            indent = 4,
            raise_exception = False
    ):

        """
        Finds one record that matches the given conditions.
        :param collection: The name of the collection to perform the search in.
        :param filter_json: The filter criteria.
        :param projections: What parts of the matching data you want to fetch.
        :param session: The session if you need to do this in a transaction.
        :param as_json_string: Whether you want it as a JSON string or a Python dict/list.
        :param indent: The indentation to use if you want it as a JSON string.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The array of matching records or null if there was an exception.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        results = None

        # try to query the data:
        try:
            results = await self.__db[collection].find_one(filter_json, projections, session = session)
            if as_json_string: results = self.__to_json_string(results, indent = indent, default = str)
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not results: raise MongoFindException(hint = f"{collection}")
        return results

    async def replace_one(
            self,
            collection,
            filter_json,
            replacement_json,
            upsert = False,
            session = None,
            raise_exception = False
    ):

        """
        To delete one document from a collection.
        :param collection: The collection from which you want to delete many records.
        :param filter_json: The filter criteria.
        :param replacement_json: The data to put in place of the existing document.
        :param upsert: If you want to insert if the document doesn't already exist.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the operation.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        success = False

        # try to query the data:
        try:
            result = await self.__db[collection].replace_one(
                filter_json,
                replacement_json,
                upsert = upsert,
                session = session
            )
            if result.modified_count or result.upserted_id: success = True
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not success: raise MongoReplaceException(hint = f"{collection}")
        return success

    async def delete_one(
            self,
            collection,
            filter_json,
            session = None,
            raise_exception = False
    ):

        """
        To delete one document from a collection.
        :param collection: The collection from which you want to delete many records.
        :param filter_json: The filter criteria.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the operation.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        count = 0

        # try to query the data:
        try:
            result = await self.__db[collection].delete_one(filter_json, session = session)
            count = result.deleted_count
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not count: raise MongoDeleteException(hint = f"{collection}")
        return count

    async def delete_many(
            self,
            collection,
            filter_json,
            session = None,
            raise_exception = False
    ):

        """
        To delete many documents from a collection.
        WARNING: sending {} in the filter would mean deleting ALL the documents.
        :param collection: The collection from which you want to delete many records.
        :param filter_json: The filter criteria.
        :param session: The session if you need to do this in a transaction.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: True or False based on the success of the operation.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        count = 0

        # try to query the data:
        try:
            result = await self.__db[collection].delete_many(filter_json, session = session)
            count = result.deleted_count
        except Exception as exception:
            self.__printer(exception)

        # Check results and return:
        if raise_exception and not count: raise MongoDeleteException(hint = f"{collection}")
        return count

    async def aggregate(
            self,
            collection,
            pipeline,
            limit = None,
            raise_exception = False
    ):

        """
        Perform an advance query on the data.
        :param collection: The collection to perform the query on.
        :param pipeline: The pipeline of actions to take. Must be a list.
        :param limit: The max. no. of records to retrieve.
        :param raise_exception: Whether, or not, you want to raise an exception when something fails.
        :return: The array of matching records or null if there was an exception.
        """

        # Ensure you are connected:
        await self.ensure_connection()

        # Assume failure:
        results = None

        # try to perform the aggregation action:
        try: results = await self.__db[collection].aggregate(pipeline).to_list(limit)
        except Exception as exception: self.__printer(exception)

        # Check results and return:
        if raise_exception and not results: raise MongoException(hint = "aggregation")
        return results


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    pass
