import logging
import os
import time
from functools import wraps
from typing import Awaitable, Callable, List, Optional

from prometheus_client import Counter, Histogram

logger = logging.getLogger(__name__)


class BasicTimeCostMetrics:
    """
    如果启动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 ApiTimeCostMetrics(BasicTimeCostMetrics):
    """
    接口操作的耗时统计
    指标名如下：
    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 ClientTimeCostMetrics(BasicTimeCostMetrics):
    """
    统计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 = BasicTimeCostMetrics()
    __api_metrics = None
    __client_metrics = None

    @classmethod
    def initialize(
        cls,
        *,
        api_metrics_buckets: Optional[List[int]] = None,
        client_metrics_buckets: Optional[List[int]] = None,
    ):
        if api_metrics_buckets is None:
            cls.__api_metrics = ApiTimeCostMetrics()
        elif len(api_metrics_buckets) == 0:
            cls.__api_metrics = BasicTimeCostMetrics()
        else:
            cls.__api_metrics = ApiTimeCostMetrics(buckets=list(api_metrics_buckets))

        if client_metrics_buckets is None:
            cls.__client_metrics = ClientTimeCostMetrics()
        elif len(client_metrics_buckets) == 0:
            cls.__client_metrics = BasicTimeCostMetrics()
        else:
            cls.__client_metrics = ClientTimeCostMetrics(
                buckets=list(client_metrics_buckets)
            )

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

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

    @classmethod
    def metrics_handler(cls, key: str, metrics: BasicTimeCostMetrics):
        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
                    metrics.add(key, cost)
                    return result
                except Exception as e:
                    metrics.add_error(key)
                    raise e

            return wrapper

        return decorator

    @classmethod
    def client_metrics_handler(cls, client_name: str):
        return cls.metrics_handler(client_name, cls.client_metrics_instance())

    @classmethod
    def api_metrics_handler(cls, api_name: str):
        return cls.metrics_handler(api_name, cls.api_metrics_instance())
