"""Adaptive per-service token-bucket sampler for otelctx."""

from __future__ import annotations

import threading
import time


class _TokenBucket:
    """Thread-safe token bucket for one service."""

    def __init__(self, rate: float) -> None:
        self._rate = rate  # tokens per second (== fraction  ≤ 1.0 → fraction of events)
        self._tokens: float = 1.0
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()

    def consume(self) -> bool:
        with self._lock:
            now = time.monotonic()
            elapsed = now - self._last_refill
            self._tokens = min(1.0, self._tokens + elapsed * self._rate)
            self._last_refill = now
            if self._tokens >= 1.0:
                self._tokens -= 1.0
                return True
            return False


class AdaptiveSampler:
    """Adaptive head-sampler with per-service token buckets.

    Sampling decisions:
    *  Always ingest if ``status_code`` indicates an error and
       ``always_errors`` is True.
    *  Always ingest if ``duration_ms`` exceeds ``latency_threshold_ms``.
    *  Otherwise, use a per-service token bucket with the configured
       ``rate`` (1.0 = 100 %, 0.1 = 10 %, etc.).

    Thread-safe.
    """

    def __init__(
        self,
        rate: float = 1.0,
        always_errors: bool = True,
        latency_threshold_ms: float = 1000.0,
    ) -> None:
        if not 0.0 <= rate <= 1.0:
            raise ValueError(f"rate must be in [0, 1], got {rate!r}")
        self._rate = rate
        self._always_errors = always_errors
        self._latency_threshold_ms = latency_threshold_ms

        self._lock = threading.Lock()
        self._buckets: dict[str, _TokenBucket] = {}

        # Counters
        self._total_seen: int = 0
        self._total_dropped: int = 0

    # ------------------------------------------------------------------
    # Sampling decision
    # ------------------------------------------------------------------

    def should_ingest(
        self,
        service: str,
        status_code: str = "UNSET",
        duration_ms: float = 0.0,
    ) -> bool:
        """Return True if the event should be ingested."""
        with self._lock:
            self._total_seen += 1

        # Always-on conditions bypass the bucket
        if self._always_errors and self._is_error(status_code):
            return True
        if duration_ms > self._latency_threshold_ms:
            return True

        # Rate == 1.0 → accept everything
        if self._rate >= 1.0:
            return True

        bucket = self._get_bucket(service)
        accepted = bucket.consume()
        if not accepted:
            with self._lock:
                self._total_dropped += 1
        return accepted

    # ------------------------------------------------------------------
    # Configuration helpers
    # ------------------------------------------------------------------

    def set_rate(self, rate: float) -> None:
        if not 0.0 <= rate <= 1.0:
            raise ValueError(f"rate must be in [0, 1], got {rate!r}")
        self._rate = rate
        with self._lock:
            # Rebuild all buckets with the new rate
            self._buckets = {}

    # ------------------------------------------------------------------
    # Observability
    # ------------------------------------------------------------------

    def stats(self) -> dict:
        with self._lock:
            return {
                "total_seen": self._total_seen,
                "total_dropped": self._total_dropped,
                "total_accepted": self._total_seen - self._total_dropped,
                "rate": self._rate,
                "always_errors": self._always_errors,
                "latency_threshold_ms": self._latency_threshold_ms,
                "active_services": len(self._buckets),
            }

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _get_bucket(self, service: str) -> _TokenBucket:
        with self._lock:
            if service not in self._buckets:
                self._buckets[service] = _TokenBucket(rate=self._rate)
            return self._buckets[service]

    @staticmethod
    def _is_error(status_code: str) -> bool:
        sc = status_code.upper()
        return sc in {"ERROR", "STATUS_CODE_ERROR", "2"} or sc.startswith("5")
