from __future__ import annotations

import functools
import time
from typing import Any

from ._promrs import (
    CONTENT_TYPE_LATEST,
    REGISTRY,
    CollectorRegistry,
    Counter as _RustCounter,
    Gauge as _RustGauge,
    Histogram as _RustHistogram,
    _CounterChild,
    _GaugeChild,
    _HistogramChild,
    generate_latest,
)


class _Timer:
    """Context manager and decorator for timing code blocks."""

    def __init__(self, metric: Any, callback_name: str) -> None:
        self._metric = metric
        self._callback_name = callback_name
        self._start: float = 0.0

    def __enter__(self) -> _Timer:
        self._start = time.perf_counter()
        return self

    def __exit__(self, _typ: Any, _value: Any, _traceback: Any) -> bool:
        duration = time.perf_counter() - self._start
        getattr(self._metric, self._callback_name)(duration)
        return False

    def __call__(self, f: Any) -> Any:
        @functools.wraps(f)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            with self:
                return f(*args, **kwargs)

        return wrapped


class _ExceptionCounter:
    """Context manager and decorator for counting exceptions."""

    def __init__(self, counter: Any, exception: type[BaseException] = BaseException) -> None:
        self._counter = counter
        self._exception = exception

    def __enter__(self) -> _ExceptionCounter:
        return self

    def __exit__(self, typ: Any, _value: Any, _traceback: Any) -> bool:
        if typ is not None and issubclass(typ, self._exception):
            self._counter.inc()
        return False

    def __call__(self, f: Any) -> Any:
        @functools.wraps(f)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            with self:
                return f(*args, **kwargs)

        return wrapped


class _InprogressTracker:
    """Context manager and decorator for tracking in-progress operations."""

    def __init__(self, gauge: Any) -> None:
        self._gauge = gauge

    def __enter__(self) -> _InprogressTracker:
        self._gauge.inc()
        return self

    def __exit__(self, _typ: Any, _value: Any, _traceback: Any) -> bool:
        self._gauge.dec()
        return False

    def __call__(self, f: Any) -> Any:
        @functools.wraps(f)
        def wrapped(*args: Any, **kwargs: Any) -> Any:
            with self:
                return f(*args, **kwargs)

        return wrapped


class _CounterChildWrapper:
    __slots__ = ("_child",)

    def __init__(self, child: _CounterChild) -> None:
        self._child = child

    def inc(self, amount: float = 1.0) -> None:
        self._child.inc(amount)

    def count_exceptions(self, exception: type[BaseException] = BaseException) -> _ExceptionCounter:
        return _ExceptionCounter(self, exception)


class _GaugeChildWrapper:
    __slots__ = ("_child",)

    def __init__(self, child: _GaugeChild) -> None:
        self._child = child

    def set(self, value: float) -> None:
        self._child.set(value)

    def inc(self, amount: float = 1.0) -> None:
        self._child.inc(amount)

    def dec(self, amount: float = 1.0) -> None:
        self._child.dec(amount)

    def set_to_current_time(self) -> None:
        self._child.set_to_current_time()

    def time(self) -> _Timer:
        return _Timer(self, "set")

    def track_inprogress(self) -> _InprogressTracker:
        return _InprogressTracker(self)


class _HistogramChildWrapper:
    __slots__ = ("_child",)

    def __init__(self, child: _HistogramChild) -> None:
        self._child = child

    def observe(self, value: float) -> None:
        self._child.observe(value)

    def time(self) -> _Timer:
        return _Timer(self, "observe")


class Counter(_RustCounter):
    def labels(self, *args: str, **kwargs: str) -> _CounterChildWrapper:
        return _CounterChildWrapper(super().labels(*args, **kwargs))

    def count_exceptions(self, exception: type[BaseException] = BaseException) -> _ExceptionCounter:
        return _ExceptionCounter(self, exception)


class Gauge(_RustGauge):
    def labels(self, *args: str, **kwargs: str) -> _GaugeChildWrapper:
        return _GaugeChildWrapper(super().labels(*args, **kwargs))

    def time(self) -> _Timer:
        return _Timer(self, "set")

    def track_inprogress(self) -> _InprogressTracker:
        return _InprogressTracker(self)


class Histogram(_RustHistogram):
    def labels(self, *args: str, **kwargs: str) -> _HistogramChildWrapper:
        return _HistogramChildWrapper(super().labels(*args, **kwargs))

    def time(self) -> _Timer:
        return _Timer(self, "observe")


__all__ = ["CONTENT_TYPE_LATEST", "REGISTRY", "CollectorRegistry", "Counter", "Gauge", "Histogram", "generate_latest"]
