from prometheus_client import Counter, Histogram  
from threading import Lock

def singleton(cls):
    instance = None
    lock = Lock()

    def get_instance(*args, **kwargs):
        nonlocal instance
        with lock:
            if instance is None:
                instance = cls(*args, **kwargs)
        return instance

    return get_instance


@singleton
class Metrics:
    def __init__(self, *args, prometheus_collector, **kwargs):
        super().__init__(*args, **kwargs)
        self.prometheus_collector = prometheus_collector

        self.prepare_count: Counter = Counter(
            name="prepare_total",
            documentation="Total messages prepared by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.prepare_duration: Histogram = Histogram(
            name="prepare_duration",
            documentation="Preparation duration by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.process_count: Counter = Counter(
            name="process",
            documentation="Total messages processed by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.process_duration: Histogram = Histogram(
            name="process_duration",
            documentation="Processing duration by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.tool_count: Counter = Counter(
            name="tool_total",
            documentation="Total messages tooled by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.tool_duration: Histogram = Histogram(
            name="tool_duration",
            documentation="Tool duration by agent type",
            buckets=(0.1, 0.5, 1.0, 2.0, float("inf")),
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.generate_count: Counter = Counter(
            name="generate_total",
            documentation="Total messages generated by agent type",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.generate_duration: Histogram = Histogram(
            name="generate_duration",
            documentation="Generate duration by agent type",
            buckets=(0.1, 0.5, 1.0, 2.0, float("inf")),
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.tokens_prompt: Counter = Counter(
            name="tokens_prompt",
            documentation="Total tokens passed as prompt to LLM",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )

        self.tokens_eval: Counter = Counter(
            name="tokens_eval",
            documentation="Total tokens returned by LLM",
            labelnames=("agent",),
            registry=self.prometheus_collector,
        )