"""

    AUTHOR:

        Khushal P Soonderji
        Bhushan Thakkar

    DATE:

        Friday, 19th April, 2024

    OBJECTIVE:

        To provide an easy interface to work with Firebase.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""

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


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

# For async operations:
# ---
import asyncio

# For system-level activities:
# ---
import os

# My async utils:
# ---
import async_json_utils

# Firebase:
# ---
import firebase_admin
import firebase_admin.firestore_async
import firebase_admin.auth
import firebase_admin.db

# For debugging and logging:
# ---
from icecream import IceCreamDebugger


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class AsyncFirebase:

    def __init__(
            self,
            credentials_json_path,
            app_name,
            max_connections = 5,
            debug = True
    ):

        # Debugging print:
        # ---
        self.__printer = IceCreamDebugger(prefix = f"FBase ({app_name[:8]}) | ", includeContext = True)
        if not debug: self.__printer.disable()
        alert = f"Firebase session starting with max. {max_connections} connections."
        self.__printer(alert)

        # Initialize the instance:
        # ---
        self.__semaphore = asyncio.Semaphore(max_connections)
        credentials = firebase_admin.credentials.Certificate(credentials_json_path)
        self.__firebase_app = firebase_admin.initialize_app(credentials, name = app_name)
        self.__firestore = firebase_admin.firestore_async.client(self.__firebase_app)

    def __del__(self):
        firebase_admin.delete_app(self.__firebase_app)
        self.__firestore.close()
        alert = "Firebase session ended."
        self.__printer(alert)

    def __user_to_json(self, firebase_user):

        user_json = {
            "uid": firebase_user.uid,
            "email": firebase_user.email,
            "emailVerified": firebase_user.email_verified,
            "displayName": firebase_user.display_name,
            "phoneNo": firebase_user.phone_number,
            "photoUrl": firebase_user.photo_url,
            "customClaims": firebase_user.custom_claims,
            "disabled": firebase_user.disabled,
            "providerId": firebase_user.provider_id,
            "providerData": firebase_user.provider_data,
            "tenantId": firebase_user.tenant_id
        }

        return user_json

    async def create_custom_token(self, uid):

        async with self.__semaphore:

            try:
                custom_token = firebase_admin.auth.create_custom_token(
                    uid = uid,
                    app = self.__firebase_app
                ).decode("utf-8")
                return custom_token

            except Exception as exception:
                self.__printer(exception)
                return None

    async def authenticate_token(self, token):

        """
        To authenticate the given session token.
        :param token: The session token generated by Firebase on a successful sign-in.
        :return: Either the retrieved user information or a blank dictionary.
        """

        try:
            user_info = firebase_admin.auth.verify_id_token(token, app = self.__firebase_app)
            return user_info
        except Exception as exception:
            self.__printer(exception)
            return {}

    async def create_user(
            self,
            uid,
            display_name = None,
            password = None,
            email = None,
            phone_number = None
    ):

        """
        Create a new user.
        :param uid: The id to identify the user by.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                firebase_admin.auth.create_user(
                    uid = uid,
                    app = self.__firebase_app
                )
                return True

            except firebase_admin.auth.UidAlreadyExistsError as excp:
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def delete_user(self, uid):

        async with self.__semaphore:

            try:
                firebase_admin.auth.delete_user(
                    uid = uid,
                    app = self.__firebase_app
                )
                return True

            except firebase_admin.auth.UserNotFoundError as excp:
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def get_user(self, uid):

        async with self.__semaphore:

            try:
                user = firebase_admin.auth.get_user(
                    uid = uid,
                    app = self.__firebase_app
                )
                return self.__user_to_json(user)

            except Exception as exception:
                self.__printer(exception)
                return None

    async def list_users(self, users_per_page, page_token = None):

        """
        To get a list of users. Firebase has a limit of 1000 per call of this API. So we use page-tokens to fetch next
        pages of users.
        :param users_per_page: How many users you want to list in this API call (Max. 1,000).
        :param page_token: To be used in case of pagination.
        :return: The list of users and the page token to be used for the next call. Will be None in case of failure.
        """

        async with self.__semaphore:

            try:
                users = firebase_admin.auth.list_users(
                    page_token = page_token,
                    max_results = min(users_per_page, 1000),
                    app = self.__firebase_app
                )
                next_page_token = users.next_page_token if users.has_next_page else None
                users = users.users
                users_json = [self.__user_to_json(user) for user in users]
                return users_json, next_page_token

            except Exception as exception:
                self.__printer(exception)
                return None, None

    async def create_document(self, collection_path, document_name, document_data = None):

        """
        To create a new document with the specified data in an existing collection.
        :param collection_path: The path of the collection (can be a sub-collection).
        :param document_name: The name of the document you want to create.
        :param document_data: The data that you want to populate in the document.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.collection(
                    collection_path
                ).add(
                    document_id = document_name,
                    document_data = document_data or {}
                )
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def delete_document(self, path):

        """
        To delete a document.
        :param path: The path of the document.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.document(path).delete()
                return True

            except Exception as excp:
                print("FIRESTORE DOCUMENT DELETION EXCEPTION:", excp)
                return False

    async def set_document(self, path, data):

        """
        To overwrite the data in a document.
        :param path: The path of the document.
        :param data: The data you want to update as a dictionary.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.document(path).set(data)
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def update_document(self, path, data):

        """
        To update the data in a document. Provide only the fields that you want to update.
        :param path: The path of the document.
        :param data: The data you want to update as a dictionary.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.document(path).update(data)
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def get_document(self, path):

        """
        To get the data in a document.
        :param path: The path of the document.
        :return: The dictionary of data as found in the path specified or None if the operation failed.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.document(path).get()
                return snapshot.to_dict()

            except Exception as exception:
                self.__printer(exception)
                return None

    async def get_document_fields(self, path, fields):

        """
        To get only specific fields (keys) in a document. Like how projections are used in Mongo.
        :param path: The path of the document.
        :param fields: The list of fields (keys) of the document that you want.
        :return: The dictionary of data as found in the path specified.
        """

        async with self.__semaphore:

            try:
                if type(fields) is not list: fields = [fields]
                snapshot = await self.__firestore.document(path).get(fields)
                return snapshot.to_dict()

            except Exception as exception:
                self.__printer(exception)
                return None

    async def create_collection(self, collection_name, document_name, document_data):

        """
        Creates a new collection in the root of the database. Note that Firebase doesn't allow creating new empty
        collections, so we must add one first document in it.
        :param collection_name: The name of the collection you want to create.
        :param document_name: The name of the first document you want to put in the collection.
        :param document_data: The data that you want to put in the first document of the new collection.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.collection(
                    collection_name
                ).add(
                    document_id = document_name,
                    document_data = document_data
                )
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def create_sub_collection(
            self,
            document_path,
            sub_collection_name,
            sub_document_name = None,
            sub_document_data = None
    ):

        """
        Firebase allows you to create collections inside documents. This method is built for that. Note that Firebase
        doesn't allow creating new empty collections, so we must add one first document in it.
        :param document_path: The path of the document in which you want to create a new collection.
        :param sub_collection_name: The name of the collection you want to create.
        :param sub_document_name: The name of the first document you want to put in the collection.
        :param sub_document_data: The data that you want to put in the first document of the new collection.
        :return: True or False based on the success of the operation.
        """

        async with self.__semaphore:

            try:
                snapshot = await self.__firestore.document(
                    document_path
                ).collection(
                    sub_collection_name
                ).add(
                    document_id = sub_document_name,
                    document_data = sub_document_data or {}
                )
                return True

            except Exception as exception:
                self.__printer(exception)
                return False

    async def get_collection(self, path):

        """
        Get a whole collection's data.
        :param path: The path of the collection.
        :return: The collection.
        """

        async with self.__semaphore:

            try:
                snapshots = await self.__firestore.collection(path).get()
                return {snapshot.id: snapshot.to_dict() for snapshot in snapshots}

            except Exception as exception:
                self.__printer(exception)
                return None


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


if __name__ == "__main__":

    import time
    from shared import constants

    user_data = {
        "firebaseUid": None,
        "roles": "ClientAdmin",
        "sessionToken": "siddhesh_20240419"
    }

    async def test():
        my_firebase = AsyncFirebase(
            f"{constants.PROJECT_DIRECTORY}/utils/cred/firebase_certs.json",
            "myFire",
            max_connections = 5
        )
        print("\n\n---\n\n")
        results = await my_firebase.get_document("testCollection/testDoc")
        print(async_json_utils.to_json_string(results))
        print("\n\n---\n\n")
        results = await my_firebase.get_document_fields("testCollection/testDoc", ["sampleMap"])
        print(async_json_utils.to_json_string(results))
        # await my_firebase.create_sub_collection(
        #     "myDeepCollection/deepDocId",
        #     "subCollection3",
        #     "subDoc",
        #     {"sub_key": "sub_val"}
        # )
        # await my_firebase.create_collection(
        #     "rootCollection",
        #     "subDoc",
        #     {"sub_key": "sub_val"}
        # )
        # await my_firebase.delete_document("activeSessions/9d402f1502dfd55a4326fa7fc8e6cb7d")
        # print(async_json_utils.to_json_string(await my_firebase.get_collection("myDeepCollection")))

    start_time = time.time()
    asyncio.run(test())
    print(f"FINISHED IN {time.time() - start_time} SECONDS.")
