"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Wednesday, 15th May, 2024

    OBJECTIVE:

        To provide a rate-limiting mechanism using the simple token bucket algorithm.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


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


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

# For date and time keeping:
import datetime
import time

# For asynchronous activities:
import asyncio

# For mathematical operations:
import math


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


# --- Nothing Yet


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


# --- Nothing Yet


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


class TokenBucket:

    """
    MODES: "add" - Add 'rate_limit' no. of tokens to the bucket every interval.
           "reset" - Reset the token count to 'rate_limit' every interval
    """

    MODE_RESET = 0
    MODE_ADD = 1

    ONE_YEAR = 3_15_36_000.0
    ONE_MONTH = 26_78_400.0
    ONE_WEEK = 6_04_800.0
    ONE_DAY = 86_400.0
    ONE_HOUR = 3_600.0
    ONE_MINUTE = 60.0
    ONE_SECOND = 1.0

    def __init__(self, rate_limit, seconds = 1.0, mode = "reset", sleep = 0.1, in_sequence = False):

        """
        Initialize the rate controller.
        :param rate_limit: The no. of operations allowed per unit of time.
        :param seconds: The time period in seconds in which the tokens get reset or added.
        :param mode: To select what happens when the period is over. 'reset' mode means that the remaining tokens from
            the previous period are discarded and the counter is set back to the rate limit, and 'add' mode means that
            new tokens are added on top of the exiting ones that were unused in the previous period.
        :param sleep: The delay to add before checking back to see if tokens are available. Keep it longer for longer
            periods (which is adjusted by the 'seconds' parameter).
        :param in_sequence: Whether, or not, you want to maintain the sequence in which the requests were made
            (experimental).
        """

        if rate_limit is not None:
            rate_limit = int(rate_limit)
            self.__semaphore = asyncio.Semaphore(1 if in_sequence else rate_limit)
        mode = mode.lower()
        if mode not in [self.MODE_RESET, self.MODE_ADD]: mode = self.MODE_RESET
        self.__mode = mode
        self.__rate_limit = rate_limit
        self.__token_count = rate_limit
        self.__last_token_generation_time = time.time()
        self.__sleep = sleep
        self.__seconds = seconds

    async def has_turn(self):

        """
        Just to check if a token is available.
        Should be used only if you don't want to wait for turn in case it is unavailable.
        :return: True if the token is available, False if not.
        """

        # If the user doesn't want any rate-limits, we return immediately.
        # Otherwise, we check if the user has any tokens available:
        if self.__rate_limit is None: return True
        if self.__token_count > 0: return True
        else: return False

    async def get_turn(self, timeout = None):

        """
        To wait till either you get your turn or the wait gets timed-out.
        :param timeout: The time (in seconds) to wait to get a turn before exiting with failure.
        :return: True if a turn was received, else False if timed-out.
        """

        # If the user doesn't want any rate-limits, we return immediately:
        if self.__rate_limit is None: return

        # Start by making variables:
        got_turn = False
        wait_start = time.time()

        # Otherwise, we wait for the semaphore:
        async with self.__semaphore:

            # Wait till you get your turn or the attempt gets timed-out:
            while True:

                # If the timeout has been exceeded:
                if (
                        timeout is not None and
                        time.time() - wait_start >= timeout
                ): break

                # If it is time to generate new tokens:
                time_delta = time.time() - self.__last_token_generation_time
                if time_delta >= self.__seconds:
                    if self.__mode == self.MODE_RESET: self.__token_count = self.__rate_limit
                    else: self.__token_count += int(self.__rate_limit * math.floor(time_delta / self.__seconds))
                    self.__last_token_generation_time = time.time()

                # If a token is available (or not):
                if self.__token_count > 0:
                    self.__token_count -= 1
                    got_turn = True
                    break
                else: await asyncio.sleep(self.__sleep)

        # Done here:
        return got_turn


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


if __name__ == "__main__":

    import random

    async def single_test(rate_controller, count):
        has_turn = await rate_controller.has_turn()
        got_turn = await rate_controller.get_turn()
        print(f"TURN: {count:.<5} {'Y' if has_turn else '-'} / {'Y' if got_turn else '-'} ({datetime.datetime.now()})")
        if got_turn: await asyncio.sleep(random.random())

    async def multi_test(max_count):
        rate_controller = TokenBucket(
            rate_limit = 10,
            seconds = 1,
            in_sequence = False,
            sleep = 0.025,
            mode = "reset",
            timeout = 3
        )
        tasks = [single_test(rate_controller, count + 1) for count in range(max_count)]
        await asyncio.gather(*tasks)


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