"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Saturday, 26th Apr., 2025

    OBJECTIVE:

        To provide an easy way to cache data for fast access. This version has the change that it can handle custom
        serializers by way of dependency injection.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""

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


# To use redis:
import redis.asyncio as redis
from redis.asyncio.sentinel import Sentinel

# Other utils:
from utils_v2.string import json
from utils_v2.serialization.pickle_serializer import PickleSerializer

# For debugging:
from icecream import IceCreamDebugger

# To make a decorator:
from functools import wraps

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

# To work with datatypes:
from typing import List, Any


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           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:

    def __init__(
            self,
            connection_string: str | dict = None,
            serializer = None,
            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 or Sentinel JSON for connecting to Redis.
        :param serializer: The serializer to use.
        :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.
        """

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

        # Note down the configuration:
        self.__client = None
        self.__serializer = serializer or PickleSerializer()
        self.__ping_counter = ping_counter
        self.__requests_since_last_ping = 0

        # Figure out the connection mechanism:
        if isinstance(connection_string, dict):
            self.__sentinel_json = connection_string
            self.__connection_string = None
            self.__printer("Received Sentinel JSON.")
        elif isinstance(connection_string, str):
            try:
                self.__sentinel_json = json.from_string(connection_string)
                self.__connection_string = None
                self.__printer("Seems like a Sentinel JSON String.")
            except:
                self.__sentinel_json = None
                self.__connection_string = connection_string
                self.__printer("Seems like a regular Connection String.")


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

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

    @staticmethod
    def make_key(*args, **kwargs) -> str:

        """
        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

    async def connect(self) -> bool:

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

        try:

            # First priority given to the cluster connection if a sentinel JSON is provided.
            # This provides robustness against the failure of one single node.
            if self.__sentinel_json is not None:
                sentinel_conn = Sentinel(
                    sentinels = [(s["host"], s["port"]) for s in self.__sentinel_json["servers"]],
                    socket_timeout = 2.5
                )
                self.__client = sentinel_conn.master_for(
                    service_name = self.__sentinel_json["serviceName"],
                    db = self.__sentinel_json["db"],
                    password = self.__sentinel_json["password"]
                )
                return True

            # If a cluster setup isn't ready,
            # we connect to a single node through a simple connection string:
            elif self.__connection_string is not None:
                self.__client = redis.from_url(
                    self.__connection_string
                )
                return True

            # When none of the available connection mechanisms are given:
            else: raise ValueError("Either a Sentinel JSON or a Connection String is needed.")

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

    async def disconnect(self) -> bool:

        """
        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) -> bool:

        """
        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 list_keys(
            self,
            match: str = "*",
            raise_exception: bool = False
    ) -> List[str] | None:

        """
        To get a list of all the keys that match the pattern.
        :param match: To match a glob-style pattern. THIS IS NOT FULL-FLEDGED REGEX.
        :param raise_exception: If you want to raise an exception if the process fails.
        :return: The list of keys if successful, else None.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        # Start with an empty list:
        keys_list = []

        try:

            # Start at the beginning,
            # and break out of the loop if the pointer returns to zero:
            pointer = 0
            while True:
                pointer, keys = await self.__client.scan(pointer, match = match)
                keys_list = keys_list + [k.decode("utf-8") for k in keys]
                if pointer == 0: break

            # Done here:
            return list(set(keys_list)) if keys_list else None

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return None

    async def ttl(
            self,
            key: str | bytes,
            raise_exception: bool = False
    ) -> int | float | None:

        """
        To get the no. of seconds till the expiry of some key.
        :param key: The key to check the expiry of.
        :param raise_exception:
        :return: -1 if the key is persistent (i.e., no expiry time set), -2 if the key does not exist, or the time left
            in seconds if the key exists and is not persistent. If something goes wrong, you will get a null value.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        # Start by assuming failure:
        ttl = None

        try:

            # Get the TTL of the key in ms,
            # convert it to seconds and return the value:
            ttl_ms = await self.__client.pttl(key)
            if ttl_ms >= 0: ttl = ttl_ms / 1_000
            return ttl

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return None

    async def set(
            self,
            key: str | bytes,
            value: Any,
            expiry: float = None,
            raise_exception: bool = False
    ) -> bool:

        """
        Saves some value to the cache. If an expiry is specified, the data will be deleted after that many seconds.
        :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.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        try:

            # Here we actually try to set the data:
            data = self.__serializer.serialize(value)
            if expiry: response = await self.__client.setex(key, int(expiry), data)
            else: response = await self.__client.set(key, data)
            return response

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return False

    async def get(
            self,
            key: str | bytes,
            raise_exception: bool = False,
            on_fail: Any = None
    ) -> Any:

        """
        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.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        try:

            # Here we try to fetch the data:
            data = await self.__client.get(key)
            data = self.__serializer.deserialize(data)
            return data

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return on_fail

    async def delete(
            self,
            key: str | bytes,
            raise_exception: bool = False
    ) -> bool:

        """
        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.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        try:

            # Try to manually delete the key before expiry:
            response = await self.__client.delete(key)
            return True if response else False

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return False

    async def count(
            self,
            key: str | bytes,
            value: int = 1,
            expiry: float = None,
            raise_exception: bool = False
    ) -> int | None:

        """
        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.
        """

        # Standard connectivity check:
        await self.ensure_connection()

        try:

            # check if the key already exists,
            # regardless of that, increment the counter:
            already_existed = await self.__client.exists(key)
            new_value = await self.__client.incrby(key, value)

            # If the key didn't already exist, specify the expiry:
            if expiry and not already_existed: await self.__client.expire(key, int(expiry))

            # Done here:
            return new_value

        # If an exception occurs in th eprocess:
        except Exception as exception:
            self.__printer(exception)
            if raise_exception: raise
            else: return None


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


if __name__ == "__main__":

    import asyncio
    from utils_v2.string import json

    async def main():

        my_cache = AsyncRedisCache(
            connection_string = input("Your Connection String or Sentinel JSON: "),
            ping_counter = 100,
            debug = True
        )
        print("Connecting.")
        connected = await my_cache.connect()
        print("Success:", connected)

        # Proceed only if we connected to the database successfully:
        if connected:

            existing_keys = await my_cache.list_keys()
            print(f"KEYS ({len(existing_keys)}):", json.to_string(existing_keys, default = str))

            # key = str(input("Paste a key to check its TTL: "))
            # key_ttl = await my_cache.ttl(key = key)
            # print(f"KEY '{key}' has {key_ttl} seconds of TTL.")
            #
            # value = await my_cache.get(key = str(input("Paste a key to get: ")))
            # print("GET:", json.to_string(value, default = str) if isinstance(value, (dict, list)) else 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")

            # print("Starting count test.")
            # await my_cache.delete(key = "cnt")
            # for _ in range(5):
            #     print("Step:", _)
            #     await asyncio.sleep(1.0)
            #     counter = await my_cache.count(key = "cnt", value = 1, expiry = 10)
            #     print("COUNTER:", counter)

    asyncio.run(main())
