"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 30th Aug., 2024

    OBJECTIVE:

        To be able to access SQL-based databases from python in a simple way.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


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

# MySQL Database:
import aiomysql
import decimal

# For data-crunching:
import pandas as pd

# For time-keeping:
import time

# OS-level operations:
import os

# My utils:
from utils_v2.string import json

# For async activities:
import asyncio

# For debugging:
from icecream import IceCreamDebugger
import traceback

# To work with datatypes:
from typing import List


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class AsyncMySQL:

    def __init__(
            self,
            pool_size,
            *args,
            **kwargs
    ):

        """
        A class to work with SQL-based databases. Originally meant to only invoke stored procedures and retrieve them as
        JSON-like structures (list or dict). The format for the results was very specific to our use case for serving
        Bicree's requirement. This may not serve your requirement at all.
        :param pool_size: The number of connections to maintain n a pool.
        :param args: Any arguments to pass. Not used.
        :param kwargs: Pass the connection configuration from here.
        """

        # Set up the variables:
        self.__args = args
        self.__kwargs = kwargs
        self.__min_pool_size = 10
        self.__max_pool_size = max(pool_size, self.__min_pool_size)
        self.__pool: aiomysql = None

        # For establishing connection:
        self.__exclusive_semaphore = asyncio.Semaphore(1)

        # Minor adjustments for backward compatibility:
        self.__kwargs["db"] = self.__kwargs.pop("database")

        # Set up the debugging tools:
        self.__printer = IceCreamDebugger(prefix = "MySQL | ", includeContext = True)

    def __del__(self):
        pass

    async def connect(self) -> bool:

        """
        Establish a connection and create a pool of connections to call from.
        :return: True if connected, else False
        """

        try:

            # Ensure that only one connection attempt is being made at one time,
            # and try to connect to the database:
            async with self.__exclusive_semaphore:
                if self.__pool is None:
                    self.__pool = await aiomysql.create_pool(
                        minsize = self.__min_pool_size,
                        maxsize = self.__max_pool_size,
                        loop = asyncio.get_event_loop(),
                        **self.__kwargs
                    )

        # In case of any exception:
        except Exception as exception:
            self.__printer(exception)
            self.__pool = None

        # Done here:
        return False if self.__pool is None else True

    async def disconnect(self) -> bool:

        try:

            # Try to close the pool and wait for the process to complete:
            self.__pool.close()
            await self.__pool.wait_closed()

        # In case of any exception:
        except Exception as exception:
            self.__printer(exception)

        # Done here:
        return False if self.__pool is None else True

    async def ensure_connection(self):

        """
        Tries to ensure that a connection is present.
        Can be called before every function to make sure that our pool is established.
        :return: None.
        """

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

    @staticmethod
    def __parse_row(row):

        """
        Converts from the custom objects of 'aiomysql' to types that are supported by Python.
        :param row: The row from the result.
        :return: The parsed row which will have types that are closer to being native to Python..
        """

        parsed_row = []
        for item in row:
            if isinstance(item, decimal.Decimal): parsed_row.append(float(item))
            else: parsed_row.append(item)
        return parsed_row

    async def fetch_all(self, cursor):

        # Make a variable to hold all the result sets.
        # Needed for when the procedure responds with many "tables":
        all_result_sets = []

        # Iterate over all result sets,
        # and process them one-by-one:
        while True:

            # Process the current result set:
            this_result_set = []
            result = await cursor.fetchall()
            if not cursor.description: break
            columns = [desc[0] for desc in cursor.description]
            for row in result: this_result_set.append(dict(zip(columns, self.__parse_row(row))))
            all_result_sets.append(this_result_set)

            # Move to the next set,
            # or break out of the loop if all done:
            if not await cursor.nextset(): break

        # Done here:
        return all_result_sets

    async def call_procedure(
            self,
            procedure_name: str,
            procedure_args: tuple,
            commit: bool = True
    ):

        """
        To call stored procedures and retrieve all the responses.
        :param procedure_name: The name of the stored procedure that must be called.
        :param procedure_args: The args to be sent to the stored procedure.
        :param commit: Whether, or not, you would like to commit the execution.
        :return: The raw result set as received from the database.
        """

        # Make sure we have a connection:
        await self.ensure_connection()

        # Make a variable to hold all the result sets.
        # Needed for when the procedure responds with many "tables":
        all_result_sets = []

        # Call the procedure and get the results:
        async with self.__pool.acquire() as connection:
            async with connection.cursor() as cursor:
                await cursor.callproc(procedure_name, procedure_args)
                all_result_sets = await self.fetch_all(cursor)
                if commit: await connection.commit()

        # Done here:
        return all_result_sets

    async def call_procedure_and_get_json(
            self,
            procedure_name,
            procedure_args,
            commit: bool = True,
            retry_count = 1,
            backoff_seconds = 0.5,
            backoff_multiplier = 1.1,
            return_exception = False
    ):

        """
        The method to call when you need to call a stored procedure and retrieve the response as a JSON-like object.
        This is custom formatting based on the structure created by Mr. bhushan Thakkar in late April (2024).
        :param procedure_name: The name of the stored procedure that must be called.
        :param procedure_args: The args to be sent to the stored procedure.
        :param commit: Whether, or not, you would like to commit the execution.
        :param retry_count: The max. number of times to try in case one or more attempts fail.
        :param backoff_seconds: The time to wait before making the next attempt if the retry count is more than 1.
        :param backoff_multiplier: The factor that dictates how much to modify the time delay by when waiting to retry.
        :param return_exception: Whether, or not, you would like to return the exception object if something goes wrong.
        :return: The formatted response and the exception (if asked for).
        """

        # Note down the start time:
        start_ts = time.time()

        # Try to get the data from the database:
        results = []
        exception = None
        for _ in range(retry_count):
            try: results = await self.call_procedure(
                procedure_name = procedure_name,
                procedure_args = procedure_args,
                commit = commit
            )
            except Exception as exc: exception = exc
            if exception is None: break
            await asyncio.sleep(backoff_seconds)
            backoff_seconds = backoff_seconds * backoff_multiplier

        # If the results are blank:
        if len(results) == 0:
            formatted_results = {
                "status": 0,
                "message": "Please contact admin (NE)" if exception is None else "Please contact admin (E)",
                "seconds": time.time() - start_ts,
                "data": {}
            }
            if return_exception: return formatted_results, exception
            else: return formatted_results

        # Extract the very basic success or failure indicators:
        formatted_results = {
            "status": results[0][0]["status"],
            "message": results[0][0].get("message", "ok"),
            "seconds": 0.0,
            "data": {}
        }

        # Handle the remaining keys of the zeroth result set:
        for key, value in results[0][0].items():
            if key not in formatted_results.keys():
                formatted_results["data"][key] = value

        # Format
        for index in range(len(results)):
            if index > 0: formatted_results["data"][f"rs{index-1}"] = results[index]

        # Note down the time taken:
        formatted_results["seconds"] = time.time() - start_ts

        # Done here:
        if return_exception: return formatted_results, exception
        else: return formatted_results

    async def execute_one(
            self,
            query: str,
            commit: bool = True,
            return_exception: bool = False
    ):

        """
        Runs one command / query in SQL.
        :param query: The query / command to run.
        :param commit: Whether, or not, you would like to commit the execution.
        :param return_exception: Whether, or not, you would like to return the exception from this function.
        :return: Either just the result or the result and the exception.
        """

        # Make sure we have a connection:
        await self.ensure_connection()

        # Start by assuming failure:
        rows_affected = None
        results = None
        excp = None

        try:

            # Get a connection and execute the command:
            async with self.__pool.acquire() as connection:
                async with connection.cursor() as cursor:
                    rows_affected = await cursor.execute(query)
                    results = await self.fetch_all(cursor)
                    if commit: await connection.commit()

        # SQL-specific errors:
        except aiomysql.MySQLError as exception:
            self.__printer("SQL Exception", exception)
            excp = exception

        # Other errors:
        except Exception as exception:
            self.__printer("Other Exception", exception)
            excp = exception

        # Done here:
        if return_exception: return rows_affected, results, excp
        else: return rows_affected, results

    async def execute_many(
            self,
            query: str,
            data: List[tuple | List],
            commit: bool = True,
            return_exception: bool = False
    ):

        """
        Runs many commands / queries in SQL.
        Consider the following example:
        QUERY: "INSERT INTO pincodeMaster (pincode, city, state) VALUES (%s, %s, %s);"
        DATA: [
            ('110001', 'New Delhi', 'Delhi'),
            ('500001', 'Hyderabad', 'Telangana'),
            ('600001', 'Chennai', 'Tamil Nadu')
        ]
        :param query: The query / command to run.
        :param data: The data to substitute into the query string.
        :param commit: Whether, or not, you would like to commit the execution.
        :param return_exception: Whether, or not, you would like to return the exception from this function.
        :return: Either just the result or the result and the exception.
        """

        # Make sure we have a connection:
        await self.ensure_connection()

        # Start by assuming failure:
        rows_affected = None
        results = None
        excp = None

        try:

            # Get a connection and execute the command:
            async with self.__pool.acquire() as connection:
                async with connection.cursor() as cursor:
                    rows_affected = await cursor.executemany(query, data)
                    results = await self.fetch_all(cursor)
                    if commit: await connection.commit()

        # SQL-specific errors:
        except aiomysql.MySQLError as exception:
            self.__printer("SQL Exception", exception)
            excp = exception

        # Other errors:
        except Exception as exception:
            self.__printer("Other Exception", exception)
            excp = exception

        # Done here:
        if return_exception: return rows_affected, results, excp
        else: return rows_affected, results


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


# --- Nothing Yet


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


if __name__ == "__main__":

    pass
