import psutil
import asyncio
import time
import util

class DynamicThrottle(object):

    def __init__(self, **kwargs):
        self.target_cpu_usage = kwargs.get("target_cpu_usage", 50)
        self.sleep_interval = 0.0

        self.sleep_increment = kwargs.get("sleep_increment", 0.01)
        self.sleep_decrement = kwargs.get("sleep_decrement", 0.01)

        self.sleep_max = kwargs.get("sleep_max", 0.1)
        self.sleep_min = kwargs.get("sleep_min", 0.01)

        self.psutil_interval = kwargs.get("psutil_interval", 0.1)

        self.log = kwargs.get("log", util.get_logger(self.__class__.__name__))

        self.consecutive_increments = 0
        self.consecutive_decrements = 0

        self.last_was_increment = False

        if kwargs.get("async"):
            self.dynamic_throttle = self.dynamic_throttle_async
        else:
            self.dynamic_throttle = self.dynamic_throttle


    async def dynamic_throttle_async(self):
        """
        Dynamically sleeps before a request if CPU usage is above our target.
        """
        current_cpu_usage = psutil.cpu_percent(interval=self.psutil_interval)

        if current_cpu_usage > self.target_cpu_usage:
            self.sleep_interval += self.sleep_increment
            if self.sleep_interval > self.sleep_max:
                self.sleep_interval = self.sleep_max
            self.log.info(
                f"CPU {current_cpu_usage}% > {self.target_cpu_usage}%, "
                f"=> sleep {self.sleep_interval:.3f}s"
            )
        elif current_cpu_usage < self.target_cpu_usage and self.sleep_interval > self.sleep_min:
            self.sleep_interval -= self.sleep_decrement
            self.log.info(
                f"CPU {current_cpu_usage}% < {self.target_cpu_usage}%, "
                f"=> sleep {self.sleep_interval:.3f}s"
            )

        if self.sleep_interval > 0:
            await asyncio.sleep(self.sleep_interval)

    def dynamic_throttle(self):
        """
        Dynamically sleeps before a request if CPU usage is above our target.
        """
        current_cpu_usage = psutil.cpu_percent(interval=self.psutil_interval)

        if current_cpu_usage > self.target_cpu_usage:
            self.sleep_interval += self.sleep_increment
            if self.sleep_interval > self.sleep_max:
                self.sleep_interval = self.sleep_max
            self.log.info(
                f"CPU {current_cpu_usage}% > {self.target_cpu_usage}%, "
                f"=> sleep {self.sleep_interval:.3f}s"
            )
        elif current_cpu_usage < self.target_cpu_usage and self.sleep_interval > self.sleep_min:
            self.sleep_interval -= self.sleep_decrement
            self.log.info(
                f"CPU {current_cpu_usage}% < {self.target_cpu_usage}%, "
                f"=> sleep {self.sleep_interval:.3f}s"
            )

        if self.sleep_interval > 0:
            time.sleep(self.sleep_interval)