"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Thursday, 11th April, 2024

    OBJECTIVE:

        To provide an easy way to work with '.json' files in an async manner.

    REFERENCES:

        1) https://www.w3schools.com/python/python_json.asp

    DOWNLOADS:

        N/A

"""

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


# To use redis:
import redis.asyncio as redis

# Other utils:
from utils_v2.string import json

# For debugging:
from icecream import IceCreamDebugger

# To make a decorator:
from functools import wraps

# For hashing and shortening the hash:
import hashlib
import base64


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


# --- Nothing Yet


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


# --- Nothing Yet


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


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   WRAPPERS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


def cache_it(cache = None, expiry = 120):

    """
    This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there.
    If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the
    cache instead of going through the whole function again.
    :param cache: The instance of 'AsyncRedisCache'.
    :param expiry: The time in seconds after which the cached data must be cleared.
    :return: The decorator that automatically caches your data.
    """

    def decorator(func):

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

            # We first use the name of the function and the inputs given to it to generate a key for Redis with the
            # simple hashing and shortening by way of base64 strings:
            inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs)
            sha256_hash = hashlib.sha256()
            sha256_hash.update(inputs_given.encode("utf-8"))
            hashed_key = sha256_hash.digest()
            base64_key = base64.b64encode(hashed_key).decode("utf-8")

            # Now we check if we have the value in cache:
            try: response = await cache.get(base64_key, raise_exception = True)

            # If the key doesn't exist, we pass through the function and store the results.
            except:
                response = await func(*args, **kwargs)
                await cache.set(key = base64_key, value = response, expiry = expiry)

            # Return the response from the wrapped function.
            return response

        return wrapper

    return decorator


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


def cache_class_methods(attr_name, expiry = 120):

    """
    This decorator factory takes an instance of the async caching class 'AsyncRedisCache' and holds your data there.
    If a subsequent call is made to the same decorated function with the same inputs, the result is fetched from the
    cache instead of going through the whole function again.
    :param attr_name: The name of the variable that has an instance of "AsyncRedisCache".
    :param expiry: The time in seconds after which the cached data must be cleared.
    :return: The decorator that automatically caches your data.
    """

    def decorator(func):

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

            # Get the cache object first:
            cache_obj = getattr(self, attr_name)

            # We first use the name of the function and the inputs given to it to generate a key for Redis with the
            # simple hashing and shortening by way of base64 strings:
            inputs_given = func.__name__ + str([_ for _ in args]) + str(kwargs)
            sha256_hash = hashlib.sha256()
            sha256_hash.update(inputs_given.encode("utf-8"))
            hashed_key = sha256_hash.digest()
            base64_key = base64.b64encode(hashed_key).decode("utf-8")

            # Now we check if we have the value in cache:
            try: response = await cache_obj.get(base64_key, raise_exception = True)

            # If the key doesn't exist, we pass through the function and store the results.
            except:
                response = await func(self, *args, **kwargs)
                await cache_obj.set(key = base64_key, value = response, expiry = expiry)

            # Return the response from the wrapped function.
            return response

        return wrapper

    return decorator


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


class AsyncRedisCache:

    # Handle datatypes when setting values to cache:
    __set_converters = {
        "set": lambda x: list(x),
        "tuple": lambda x: list(x),
        "complex": lambda x: {"r": x.real, "i": x.imag}
    }

    # Handle datatypes when getting from cache:
    __get_converters = {
        "set": lambda x: set(x),
        "tuple": lambda x: tuple(x),
        "complex": lambda x: complex(x["r"], x["i"])
    }

    def __init__(
            self,
            connection_string,
            ping_counter = 1_000,
            debug = False,
            debug_prefix = "R-Cache | "
    ):

        """
        Implements a simple cache in Redis which holds and returns all native datatypes like ints, floats, bools,
        strings, dicts, lists, sets, and tuples :)
        :param connection_string: The connection URL for connecting to Redis.
        :param ping_counter: The number of requests to Redis after which you want to ping to ensure connection.
        :param debug: Whether, or not, you want to show debugging messages from the start.
        :param debug_prefix: The prefix text to show with the debugging messages.
        """

        # Note down the configuration:
        self.__client = None
        self.__ping_counter = ping_counter
        self.__requests_since_last_ping = 0
        self.__connection_string = connection_string

        # For debugging:
        self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self.__printer.disable()

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

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

    @staticmethod
    def make_key(*args, **kwargs):

        """
        Generates a key by hashing the args and kwargs sent to it. Can be useful to generate a predictable key. If the
        inputs stay the same, the output stays the same.
        :param args: Any number of args that you would like to use to generate the key.
        :param kwargs: Any number of kwargs that you would like to use to generate the key.
        :return: A string that can be used as a key to store values on Redis.
        """

        # Take everything into one plain text string:
        plain_text = "".join(str(a) for a in args)
        plain_text += json.to_string(kwargs, no_space = True)

        # Hash the plain text value, and create a key from it:
        sha256_hash = hashlib.sha256()
        sha256_hash.update(plain_text.encode("utf-8"))
        hashed_key = sha256_hash.digest()
        base64_key = base64.b64encode(hashed_key).decode("utf-8")

        # Done here:
        return base64_key

    def add_converter(
            self,
            type_name,
            set_converter_func,
            get_converter_func,
    ):

        """
        To add custom data converters to handle new values.
        RULE: Each of the converter functions must take in exactly on argument.
        :param type_name: The name of the type of the object to handle. HINT: type(obj).__name__.
        :param set_converter_func: A converter function that converts from the custom datatype to a datatype that Redis
            can work with. Try converting to an bool, int, float, str, list or dict.
        :param get_converter_func: The inverse of the set converter. This function will be used to convert from a
            datatype that Redis can work with to the custom datatype.
        :return: None.
        """

        self.__set_converters[type_name] = lambda x: set_converter_func(x)
        self.__get_converters[type_name] = lambda x: get_converter_func(x)

    async def connect(self):

        """
        To make an asynchronous connection request to the Redis server to establish a connection.
        :return: True if connected, else False.
        """

        try:
            self.__client = redis.from_url(
                self.__connection_string
            )
            return True

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

    async def disconnect(self):

        """
        Close the connection to the Redis server.
        :return: True if successful, else False.
        """

        if self.__client is not None:
            try: await self.__client.close()
            except Exception as exception:
                self.__printer(exception)
                return False
            return True
        return True

    async def ensure_connection(self):

        """
        To ensure that we are connected. We keep pinging the Redis server every once in a while even when connected.
        :return: True if connected, else False.
        """

        # If we are not connected, we try to establish a connection:
        if self.__client is None: return await self.connect()

        # Else we check if we are connected. If not, we try to connect.
        # But we check only once in a while. In the meantime, we assume that we are connected.
        elif self.__requests_since_last_ping > self.__ping_counter:
            try:
                await self.__client.ping()
                self.__requests_since_last_ping = 0
                return True
            except Exception as exception:
                self.__printer(exception)
                return await self.connect()
        else:
            self.__requests_since_last_ping += 1
            return True

    async def set(self, key, value, expiry: float = 120, raise_exception = False):

        """
        Saves some value to the cache.
        :param key: The key with which the data will be stored and retrieved.
        :param value: The value to store.
        :param expiry: The time in seconds after which the data will expire. Must be a positive number.
        :param raise_exception: If you want to raise an exception if the process fails.
        :return: True if cached, else False.
        """

        await self.ensure_connection()

        try:

            # We handle conversion for unsupported Pythonic datatypes:
            original_type = type(value).__name__
            converter = self.__set_converters.get(original_type)
            if converter is not None: value = converter(value)

            # Here we actually try to store the data:
            response = await self.__client.setex(
                key,
                int(expiry),
                json.to_string({
                    "value": value,
                    "type": original_type
                }, indent = 0, no_space = False)
            )
            return response

        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return False

    async def get(self, key, raise_exception = False, on_fail = None):

        """
        Retrieve the cached value.
        :param key: The key with which the data was saved.
        :param raise_exception: If you want to raise an exception if the process fails.
        :param on_fail: What to return if the process fails due to an exception.
        :return: The retrieved data or null if not found.
        """

        await self.ensure_connection()

        try:
            data = await self.__client.get(key)
            data = data.decode("utf-8")
            data = json.from_string(data)
            converter = self.__get_converters.get(data["type"])
            if converter is not None: return converter(data["value"])
            return data["value"]

        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return on_fail

    async def delete(self, key, raise_exception = False):

        """
        Prematurely delete the value from the cache before it expires.
        :param key: The key with which the data was saved.
        :param raise_exception: If you want to raise an exception if the process fails.
        :return: True if deleted, else False.
        """

        await self.ensure_connection()

        try:
            response = await self.__client.delete(key)
            return True if response else False

        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return False

    async def count(self, key, value: int = 1, expiry: float = None, raise_exception = False):

        """
        To use simple counters. If the counter (identified by the 'key') exists, it will be incremented, else the
        counter will be created and the value will be incremented from 0.
        :param key: The name of the counter.
        :param value: The amount to increment the value by. Send negative values to count backwards.
        :param expiry: The time (in seconds) in which the counter expires. Starts from the time the counter is created.
            This value has to be an integer. If a float is passed, the value will be rounded off.
        :param raise_exception: If you want to raise an exception if the process fails.
        :return: The latest value of the counter. Will be null if something went wrong and the exception was suppressed.
        """

        await self.ensure_connection()

        try:

            already_existed = await self.__client.exists(key)
            new_value = await self.__client.incrby(key, value)
            if expiry and not already_existed: await self.__client.expire(key, int(expiry))
            return new_value

        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return None


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


if __name__ == "__main__":

    import asyncio

    def set_complex(x):
        return {"r": x.real, "i": x.imag}

    def get_complex(x):
        return complex(x["r"], x["i"])

    async def main():

        my_cache = AsyncRedisCache(
            connection_string = r"redis://:dc4da94197c843ab6a730113c2b801d9@redis.ditscentre.in/0",
            ping_counter = 100,
            debug = False
        )

        my_cache.add_converter(
            type_name = type(2j).__name__,
            set_converter_func = set_complex,
            get_converter_func = get_complex
        )

        value = await my_cache.get(key = "6b61afd0-b066-4611-9791-411a30d34624")
        print("GET:", value)
        print("TYP:", type(value), end = "\n\n")

        # success = await my_cache.set(
        #     key = "name",
        #     value = {"first": "John", "last": "Doe"},
        #     expiry = 10
        # )
        # print("SET:", success, end = "\n\n")
        #
        # value = await my_cache.get(key = "name")
        # print("GET:", value)
        # print("TYP:", type(value), end = "\n\n")
        #
        # success = await my_cache.delete(key = "name")
        # print("DEL:", success, end = "\n\n")
        #
        # value = await my_cache.get(key = "cnt")
        # print("GET:", value)
        # print("TYP:", type(value), end = "\n\n")
        #
        # await my_cache.delete(key = "cnt")
        # for _ in range(50):
        #     await asyncio.sleep(1.0)
        #     counter = await my_cache.count(key = "cnt", value = 1, expiry = 10)
        #     print("COUNTER:", counter)

    asyncio.run(main())
