"""

    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 system-level activity:
import io

# For async behaviour:
import asyncio

# MongoDB for File Storage:
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorGridFSBucket
from bson.objectid import ObjectId

# For debugging:
from icecream import IceCreamDebugger


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


# --- Nothing Yet


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class AsyncMongoStorage:

    __db = None
    __db_name = None
    __client = None
    __store = None

    def __init__(
            self,
            connection_string = None,
            max_connections = 5,
            host_name = "localhost",
            port = 27017,
            database_name = "fileStore",
            debug = True,
            debug_prefix = "GridFS (M) | ",
            debug_only_errors = True
    ):

        # Database Initialization:
        self.__host_name = host_name
        self.__port = port
        self.__db_name = database_name
        self.__connection_string = connection_string
        self.__max_connections = max_connections

        # Debugging:
        self.__debug_only_errors = debug_only_errors
        self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: 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):

        """
        Initialize the database connection.
        :return: Nothing.
        """

        if self.__connection_string is None:
            self.__client = AsyncIOMotorClient(
                self.__host_name,
                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)
        self.__store = AsyncIOMotorGridFSBucket(self.__db)

    @property
    def fs(self):

        """
        To access the features that have not been wrapped in this reportlab directly.
        This could include things like streaming files chunk-by-chunk.
        :return: The file-store instance.
        """

        return self.__store

    async def write_from_memory(self, file_name, file_data, metadata_json = None):

        """
        Save a file (from RAM) to Mongo. Suitable for smaller files (a few MBs max.).
        :param file_name: The name of the file.
        :param file_data: The data of the file (held in RAM).
        :param metadata_json: A JSON of metadata information that can later be used to search files (RECOMMENDED).
        :return: The file's id as a string (if it gets saved) or None.
        """

        if self.__store is None: await self.connect()

        file_data.seek(0)
        file_size = file_data.__sizeof__()

        file_id = None
        try: file_id = await self.__store.upload_from_stream(file_name, file_data, metadata = metadata_json)
        except Exception as exception: self.__printer(exception, file_name, file_size, file_id)
        if not self.__debug_only_errors: self.__printer(file_name, file_size, file_id)
        return str(file_id)

    async def read_to_memory(self, file_identifier, by_id = True):

        """
        To retrieve a file (in RAM) based on the provided identifier.
        Suitable for smaller files (a few MBs max.).
        :param file_identifier: Either the name or the "_id" of the file.
        :param by_id: Set to True if you are fetching by the "_id" of the file.
        :return: Either the file (in RAM) or None.
        """

        if self.__store is None: await self.connect()

        file_data = None

        try:
            if by_id: grid_out = await self.__store.open_download_stream(ObjectId(file_identifier))
            else: grid_out = await self.__store.open_download_stream_by_name(file_identifier)
            file_data = io.BytesIO(await grid_out.read())
            file_data.seek(0)

        except Exception as exception:
            file_data = None
            self.__printer(exception, file_identifier, by_id)

        if not self.__debug_only_errors: self.__printer(file_identifier, by_id)
        return file_data

    async def delete_file_by_id(self, file_id):

        """
        Tries to delete one file by the id.
        :param file_id: The id of the file in the database.
        :return: True or False based on the success of the operation.
        """

        if self.__store is None: await self.connect()

        deleted = False

        try:
            response = await self.__store.delete(file_id = ObjectId(file_id))
            deleted = True
        except Exception as exception:
            self.__printer(exception, file_id, deleted)

        return deleted

    def __format_metadata_json(self, metadata_json):

        """
        NOTE: ONLY USE WHEN SEARCHING FILES BY METADATA.
        MongoDB expects dot-notation while searching for files by the metadata. We are making a function to search
        files assuming that the conditions are to be applied to the metadata itself. So this function add the
        dot-notation to the right places to conduct a successful search.
        :param metadata_json: The JSON to format.
        :return: The formatted JSON that has the right dot-notation.
        """

        formatted_metadata_json = {}

        for key, value in metadata_json.items():
            if not key.startswith("$"): key = f"metadata.{key}"
            else:
                if type(value) is dict: value = self.__format_metadata_json(value)
                if type(value) is list: value = [self.__format_metadata_json(item) for item in value]
            formatted_metadata_json[key] = value

        return formatted_metadata_json

    async def find_file_by_metadata(self, metadata_json, limit = None, skip = None, sort = None):

        """
        This method only lists the files that match the criteria mentioned in the metadata JSON.
        :param metadata_json: The JSON that describes what you want to find.
        :param limit: Max. no. of records to retrieve.
        :param skip: No. of starting results to skip. Useful for pagination.
        :param sort: The sorting conditions to follow.
        :return: A list of (JSONs of) files that match the conditions. The list can be empty.
        """

        if self.__store is None: await self.connect()

        files_list = []

        try:
            limit = limit or 10
            skip = skip or 0
            sort = {"_id": -1} if not isinstance(sort, dict) else sort
            formatted_metadata_json = self.__format_metadata_json(metadata_json)
            return await self.__store.find(
                formatted_metadata_json
            ).sort(sort).skip(skip).limit(limit).to_list(None)

        except Exception as exception: self.__printer(exception, metadata_json, len(files_list))

        if not self.__debug_only_errors: self.__printer(metadata_json, len(files_list))
        return files_list

    async def find_file_by_id(self, file_id):

        """
        This method allows you to get the file's info from the id of the file.
        :param file_id: The id that was assigned by Mongo during upload.
        :return: The file's info or None if the file doesn't exist.
        """

        if self.__store is None: await self.connect()

        file_info = None

        try:
            formatted_metadata_json = {"_id": ObjectId(file_id)}
            file_info = (await self.__store.find(formatted_metadata_json).to_list(1))[0]

        except Exception as exception: self.__printer(exception, file_id, file_info)

        return file_info

    async def get_file_name(self, file_id):

        """
        Returns the file name if the id of the file is known.
        :param file_id: The id of the file as assigned by MongoDB when the file was stored.
        :return: The file's name (if it exists), or None.
        """

        # Ensure that we are connected:
        if self.__store is None: await self.connect()

        # Ensure that the input given is of 'ObjectId' type:
        if type(file_id) is not ObjectId: file_id = ObjectId(str(file_id))

        # Fetch and return the file name:
        files_list = await self.__store.find(
            {"_id": file_id},
            {"filename": True}
        ).sort({"_id": -1}).limit(1).to_list(None)
        try: file_name = files_list[0]["filename"]
        except: file_name = None
        return file_name

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

        """
        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.
        :return: The array of matching records or null if there was an exception.
        """

        if self.__store is None: await self.connect()

        results = None
        try: results = await self.__db[collection].aggregate(pipeline).to_list(limit)
        except Exception as exception: self.__printer(exception)
        return results


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


if __name__ == "__main__":

    pass
