from functools import wraps
from typing import Optional, List, Callable, Awaitable
from prometheus_client import Histogram, Counter
import os
from loguru import logger
import time

from huibiao_framework.config.config import MetricsConfig


class RequestOperationMetrics:
    """
    如果启动prometheus则使用此类
    """

    def __init__(self, buckets: Optional[List[float]] = None):
        default_buckets = [2, 5, 10, 30, 60, 120, 180, 240, 300]
        self.buckets = buckets if buckets is not None else default_buckets
        self.pid = os.getpid()

    def add(self, key, cost): ...

    def add_error(self, key): ...


class PrometheusApiMetrics(RequestOperationMetrics):
    """
    算法接口中算法操作的耗时统计
    指标名如下：
    operation_duration_seconds_count
    operation_duration_seconds_bucket
    operation_errors_total
    """

    def __init__(self, buckets: Optional[List[float]] = None):
        super().__init__(buckets=buckets)

        self.ops_duration = Histogram(
            "operation_duration_seconds",
            "Api operation duration in seconds",
            ["operation", "pid"],
            buckets=self.buckets,
        )
        self.ops_error_cnt = Counter(
            "operation_errors_total",
            "The total number of errors in operations.",
            ["operation", "pid"],
        )

    def add(self, key, cost):
        self.ops_duration.labels(operation=key, pid=self.pid).observe(cost)

    def add_error(self, key):
        self.ops_error_cnt.labels(operation=key, pid=self.pid).inc(1)


class PrometheusClientMetrics(RequestOperationMetrics):
    """
    统计client请求的耗时
    """

    def __init__(self, buckets: Optional[List[float]] = None):
        super().__init__(buckets=buckets)

        self.duration = Histogram(
            "client_duration_seconds",
            "Client response duration in seconds",
            ["client", "pid"],
            buckets=self.buckets,
        )
        self.error_cnt = Counter(
            "client_errors_total",
            "The total number of errors in client response.",
            ["client", "pid"],
        )

    def add(self, key, cost):
        self.duration.labels(client=key, pid=self.pid).observe(cost)

    def add_error(self, key):
        self.error_cnt.labels(client=key, pid=self.pid).inc(1)


class PrometheusTimeCostMetricSetting:
    __default_metrics = RequestOperationMetrics()
    __api_metrics = None
    __client_metrics = None

    @classmethod
    def initialize(
        cls,
        *,
        api_metrics_buckets: Optional[List[float]] = None,
        client_metrics_buckets: Optional[List[float]] = None,
    ):
        if (
            cls.__api_metrics is None
            and not MetricsConfig.IGNORE_REQUEST_OPERATION_METRICS
        ):
            cls.__api_metrics = PrometheusApiMetrics(buckets=api_metrics_buckets)
            logger.info(f"apiRequestOperationBuckets => {api_metrics_buckets}")

        if (
            cls.__client_metrics is None
            and not MetricsConfig.IGNORE_CLIENT_OPERATION_METRICS
        ):
            cls.__client_metrics = PrometheusClientMetrics(
                buckets=client_metrics_buckets
            )
            logger.info(f"clientRequestOperationBuckets => {client_metrics_buckets}")

    @classmethod
    def api_metrics_instance(cls) -> RequestOperationMetrics:
        return cls.__api_metrics or cls.__default_metrics

    @classmethod
    def client_metrics_instance(cls) -> RequestOperationMetrics:
        return cls.__client_metrics or cls.__default_metrics

    @classmethod
    def client_metrics_handler(cls, client_name: str):
        def decorator(func: Callable[..., Awaitable]) -> Callable[..., Awaitable]:
            @wraps(func)
            async def wrapper(*args, **kwargs):
                try:
                    start = time.perf_counter()
                    result = await func(*args, **kwargs)
                    cost = time.perf_counter() - start
                    cls.client_metrics_instance().add(client_name, cost)
                    return result
                except Exception as e:
                    cls.client_metrics_instance().add_error(client_name)
                    raise e

            return wrapper

        return decorator
