from functools import wraps
import time
import random
from typing import Callable, TypeVar, Any
from abc import ABCMeta, abstractmethod
from elasticsearch import Elasticsearch
from logging import Logger, getLogger, INFO, StreamHandler, Formatter


T = TypeVar('T')


class Error(Exception):
    pass


class MaxRetriesExceededError(Error):
    pass


class ESGuard():
    def __init__(
        self,
        es: Elasticsearch=Elasticsearch(),
        os_cpu_percent: int=90,
        os_mem_used_percent: int=90,
        jvm_mem_heap_used_percent: int=90,
        logger: Logger=None
    ) -> None:
        self.logger = logger
        self.es = es
        self.os_cpu_percent = os_cpu_percent
        self.os_mem_used_percent = os_mem_used_percent
        self.jvm_mem_heap_used_percent = jvm_mem_heap_used_percent

    def _warning(self, message: str) -> None:
        if self.logger is None:
            return
        self.logger.warning(message)
    
    def _wait(self) -> bool:
        wait = False
        res = self.es.nodes.stats(metric=['os', 'jvm'])
        for k in res['nodes'].keys():
            cpu = res['nodes'][k]['os']['cpu']['percent']
            if cpu >= self.os_cpu_percent and self.os_cpu_percent > 0:
                self._warning(f'node({k}) OS CPU usage {cpu}% over {self.os_cpu_percent}% ')
                wait = True
                break

            mem = res['nodes'][k]['os']['mem']['used_percent']
            if mem >= self.os_mem_used_percent and self.os_mem_used_percent > 0:
                self._warning(f'node({k}) OS MEM usage {mem}% over {self.os_mem_used_percent}% ')
                wait = True
                break

            jvm_heap = res['nodes'][k]['jvm']['mem']['heap_used_percent']
            if jvm_heap >= self.jvm_mem_heap_used_percent and self.jvm_mem_heap_used_percent > 0:
                self._warning(f'node({k}) JVM heap usage {jvm_heap}% over {self.jvm_mem_heap_used_percent}% ')
                wait = True
                break

        return wait

    def decotator(self) ->  Callable[[], T]:
        def _retry(func) -> Callable[[], T]:
            @wraps(func)
            def wrapper(*args, **kwargs) -> Any:
                x = 0
                backoff_in_seconds = 1
                retries = 3
                
                while self._wait():
                    if x == retries:
                        raise MaxRetriesExceededError(f'max retries exceeded {retries}')
                    sleep = (backoff_in_seconds * 2 ** x + random.uniform(0, 1))
                    time.sleep(sleep)
                    x += 1

                return func(*args, **kwargs)

            return wrapper

        return _retry
