"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 2nd Aug., 2024

    OBJECTIVE:

        To provide an easy way to log all system activities by way of managing the context of what is going on.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

    NOTES:

        N/A

"""


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


# To manage context:
import contextvars
from contextlib import contextmanager, asynccontextmanager

# To make decorators:
from functools import wraps

# For system-level activities:
import os

# My utils:
from utils_v2.string import json
from utils_v2.date_time import date_time
from utils_v2.security import sanitizers
from utils_v2.api.codes import StatusCodes, HttpCodes
from utils_v2.api.log import APILogModel
from utils_v2.api.response import ResponseModel
from utils_v2.database.async_mongo_v2 import AsyncMongo

# The needed data models:
from utils_v2.logging.model import GeneralLogModel

# To work with date and time:
import time

# For random strings:
import random

# For debugging:
import traceback
import string

# To work with datatypes:
from types import NoneType
import pandas as pd

# To work with Pydantic objects:
from pydantic import BaseModel

# For asynchronous activities:
import asyncio


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


# Chars to choose from for random strings:
ALPHANUMERIC_CHARS = string.ascii_letters + string.digits

# To capture system information:
PROCESS_ID = os.getppid()
PARENT_PROCESS_ID = os.getppid()


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


# --- Nothing Yet


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


# --- Nothing Yet


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


# --- Nothing Yet


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


def describe_exception(exc):

    """
    Describes the exception in detail. It extracts the type of exception, a brief message, and even the entire
    traceback. Useful for debugging in details without the terminal. You could either log the resultant dict or send it
    to the dev team over some service like WhatsApp/Telegram.
    :param exc: The exception that occurred.
    :return: The dict that explains the exception.
    """

    exc_desc = {
        "type": type(exc).__name__,
        "msg": str(exc),
        "tb": [str(exc_tb) for exc_tb in traceback.format_exception(exc, value = exc, tb = exc.__traceback__)]
    }

    return exc_desc


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   Classes                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class AsyncMongoLogger:

    def __init__(self, db_conn: AsyncMongo, collection: str = "logs"):

        """
        This class uses an instance of 'AsyncMongo' and makes it usable as a logger.
        :param db_conn: The instance of 'AsyncMongo' to use.
        :param collection: The collection to write the log into.
        """

        self.__db_conn = db_conn
        self.__collection = collection

    async def log(self, log_json):

        """
        Log something to the database using the connection provided when the object was made.
        :param log_json: The dict to log.
        :return: True if logged successfully, else False.
        """

        if isinstance(self.__db_conn, AsyncMongo):
            asyncio.create_task(self.__db_conn.insert_one(
                collection = self.__collection,
                document = log_json,
                raise_exception = False
            ))


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


class AsyncLoggerContext:

    # Create the context-aware variable(s):
    logger = contextvars.ContextVar("logger", default = None)
    log_chain = contextvars.ContextVar("log_chain", default = None)

    @classmethod
    @asynccontextmanager
    async def logging_context(cls, logger, log_chain = None):

        """
        This function makes the context manager that makes the value of the log chain available to everything that is
        called within the scope of the context.
        :param logger: The object which is to be used to write the log. It should have a 'log' method which should take
            in a dict as its input.
        :param log_chain: The value of the log chain to be made available within the scope.
        :return: None.
        """

        # Set the context:
        token_logger = cls.logger.set(logger)
        token_log_chain = cls.log_chain.set(log_chain)

        # Make the objects available within the context:
        try: yield

        # Release the objects when the context is over:
        finally:
            cls.logger.reset(token_logger)
            cls.log_chain.reset(token_log_chain)

    @staticmethod
    def generate_log_chain(count = 8):
        return "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(min(8, count)))

    @staticmethod
    def generate_log_id(count = 8):
        return "".join(random.choice(ALPHANUMERIC_CHARS) for _ in range(min(8, count)))

    @classmethod
    def get_logger(cls):
        return cls.logger.get()

    @classmethod
    def get_log_chain(cls):
        return cls.log_chain.get()

    @staticmethod
    def summarize(
            value,
            str_limit = 100,
            expand: bool | int = False,
            sensitive_keys: list[str] = None
    ):

        """
        To summarize an input value to capture the essence without hoarding to much data.
        :param value: Anything that you want to summarize.
        :param str_limit: The max. no. of chars of a string to retain.
        :param expand: Set to True for full expansion, False for no expansion, and an integer for a specific level of
            expansion. Applicable on iterables and dicts. The smaller this number, the more concise the summary will be,
            and vice versa.
        :param sensitive_keys: The list of keys (of a dict) to obscure when summarizing.
        :return: The summarized version of the input.
        """

        # If the input is a Pydantic class:
        if isinstance(value, BaseModel): value = value.model_dump()

        # Check the sensitive keys:
        if sensitive_keys is None: sensitive_keys = []

        # Handle datatypes that you don't want to modify:
        if isinstance(value, (int, float, bool, NoneType)): pass

        # When the value is a list or similar iterable:
        elif isinstance(value, (list, tuple, set)):
            if expand:
                if not isinstance(expand, bool): expand -= 1
                value = [AsyncLoggerContext.summarize(
                    v,
                    expand = expand,
                    sensitive_keys = sensitive_keys
                ) for v in value]
            else: value = f"array of {len(value)} item(s)"

        # If the value is a dict:
        elif isinstance(value, dict):
            if expand:
                if not isinstance(expand, bool): expand -= 1
                value = {
                    k: AsyncLoggerContext.summarize(
                        v,
                        expand = expand,
                        sensitive_keys = sensitive_keys
                    ) if k not in sensitive_keys else "********"
                    for k, v in value.items()
                }
            else: value = f"object of {len(value.keys())} field(s) [{', '.join(value.keys())}]"

        # When a dataframe is passed:
        elif isinstance(value, pd.DataFrame):
            cols = value.columns.to_list()
            value = f"table with {len(cols)} col(s) [{', '.join(cols)}] and {len(value)} row(s)"
            str_limit = 999

        # If the input is some form of non-standard object:
        else: value = str(value)

        # Handle strings:
        if isinstance(value, str):
            if len(value) > str_limit: value = value[:str_limit] + "..."

        # Done here:
        return value

    @classmethod
    def log_it(
            cls,
            api_version: str = None,
            project: str = None,
            log_type: str = None,
            operation: str = None,
            log_input: bool | int = True,
            log_output: bool | int = True,
            sensitive_keys: list = None
    ):

        """
        A decorator factor that can be used to log the results of functions automatically.
        :param api_version: A string that indicates the version code of the function being decorated.
        :param project: A hint about which project is being worked on.
        :param log_type: A hint about which module is being worked on.
        :param operation: A hint about which action in a particular module is being worked on.
        :param log_input: Set to True to capture everything that went into the function, False to capture the least
            info, and set it to an integer to capture a certain depth of the input (applicable on iterables and dicts.
        :param log_output: The same as 'log_input', but applicable to the response from the function.
        :param sensitive_keys: Keys of a dict whose values must be obscured even if that depth is being captured.
        :return: A decorator with the configuration.
        """

        def decorator(func):

            @wraps(func)
            async def wrapper(*args, **kwargs):

                # Make variables and extract available info.:
                exception = None
                response = None
                request_ts = date_time.get_current_utc_date_time()
                start_ts = time.perf_counter()
                cpu_start_ts = time.process_time()

                # Execute the function that is being wrapped:
                try: response = await func(*args, **kwargs)
                except Exception as exc: exception = exc

                # Do the next steps only if within the logging context:
                if cls.get_logger() is not None:

                    # Create the log:
                    if not args: args = None
                    if not kwargs: kwargs = None
                    func_log = GeneralLogModel(
                        pid = PROCESS_ID,
                        ppid = PARENT_PROCESS_ID,
                        project = project,
                        log = log_type or func.__name__,
                        operation = operation or func.__name__,
                        apiVer = api_version,
                        logId = cls.generate_log_id(),
                        logChain = cls.get_log_chain(),
                        ts = request_ts,
                        tat = time.perf_counter() - start_ts,
                        cpuTime = time.process_time() - cpu_start_ts,
                        func = func.__name__,
                        args = cls.summarize(args, expand = log_input, sensitive_keys = sensitive_keys),
                        kwargs = cls.summarize(kwargs, expand = log_input, sensitive_keys = sensitive_keys),
                        exception = None if exception is None else describe_exception(exception),
                        response = cls.summarize(response, expand = log_output, sensitive_keys = sensitive_keys),
                    ).model_dump()

                    # Write the log:
                    await cls.get_logger().log(func_log)

                # Done here:
                if exception is not None: raise exception
                return response

            return wrapper

        return decorator


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


if __name__ == "__main__":

    from utils_v2.database.async_mongo_v2 import AsyncMongo

    @AsyncLoggerContext.log_it(
        api_version = "0.0.1",
        project = "testProj",
        log_type = "work",
        operation = "someWork",
        log_input = True,
        log_output = True,
        sensitive_keys = None
    )
    async def some_work(*args, **kwargs):
        print("SOME WORK:", AsyncLoggerContext.get_log_chain())
        await asyncio.sleep(max(2.0 * random.random(), 1.0))
        total = sum(args)
        return total

    @AsyncLoggerContext.log_it(
        api_version = "0.0.1",
        project = "testProj",
        log_type = "work",
        operation = "moreWork",
        log_input = True,
        log_output = True,
        sensitive_keys = ["password", "sessionToken"]
    )
    async def more_work(*args, **kwargs):
        print("MORE WORK:", AsyncLoggerContext.get_log_chain())
        await asyncio.sleep(max(2.0 * random.random(), 1.0))
        return {"success": True, "sessionToken": "1234567890"}

    @AsyncLoggerContext.log_it(
        api_version = "0.0.1",
        project = "testProj",
        log_type = "work",
        operation = "moreWork",
        log_input = True,
        log_output = True,
        sensitive_keys = None
    )
    async def last_work(*args, **kwargs):
        print("LAST WORK:", AsyncLoggerContext.get_log_chain())
        await asyncio.sleep(max(2.0 * random.random(), 1.0))

    async def main(chain = None):

        # Connect to MongoDB:
        mongo = AsyncMongo(
            connection_string = r"mongodb://del.ditscentre.in:27017,wtt.ditscentre.in:27017,mum.arh.001.ditscentre.in:27017/admin?tls=true&tlsCAFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_ca.pem&tlsCertificateKeyFile=%2Fetc%2Fssl%2Fcerts%2Fmongo_data_cert.pem&replicaSet=dits_mongod_rep&readPreference=primary&authMechanism=MONGODB-X509&authSource=%24external",
            database_name = "converse",
            max_connections = 10,
            debug = True
        )

        # Convert the connection to a logger instance that can be injected
        # into the context as a dependency:
        mongo_logger = AsyncMongoLogger(
            db_conn = mongo,
            collection = "logs"
        )

        # Initialize the context:
        async with AsyncLoggerContext.logging_context(
                logger = mongo_logger,
                log_chain = chain
        ):

            # Run some functions within the context:
            await some_work(1, 2, 3, 4, 5)
            await more_work(username = "john.doe@domain.com", password = "mySecretPass")

        # Run something outside the context:
        await last_work()

    # async def multi_main():
    #     tasks = [
    #         main(chain = "kPRwXdItb1"),
    #         main(chain = "456")
    #     ]
    #     await asyncio.gather(*tasks)

    asyncio.run(main(chain = "00wGHRFYPY123"))
