"""
RBSP Runtime Throttle — Rate limiting and backpressure.

Provides token-bucket rate limiting, sliding-window counters, and
adaptive backpressure to prevent the engine from being overwhelmed
by concurrent indexing or query operations.
"""

from __future__ import annotations

import threading
import time
from typing import Dict


# ---------------------------------------------------------------------------
# Token Bucket
# ---------------------------------------------------------------------------

class TokenBucket:
    """Classic token-bucket rate limiter.

    Tokens are replenished at *rate* tokens per second up to *capacity*.
    Callers acquire tokens before proceeding; if the bucket is empty they
    either block (with optional timeout) or fail immediately.
    """

    def __init__(self, rate: float, capacity: int) -> None:
        self._rate = float(rate)
        self._capacity = int(capacity)
        self._tokens = float(capacity)
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()

    # -- internal ----------------------------------------------------------

    def _refill(self) -> None:
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
        self._last_refill = now

    # -- public API --------------------------------------------------------

    def try_acquire(self, tokens: int = 1) -> bool:
        """Non-blocking acquire.  Returns True if *tokens* were available."""
        with self._lock:
            self._refill()
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True
            return False

    def acquire(self, tokens: int = 1, timeout: float | None = None) -> bool:
        """Blocking acquire.

        Waits up to *timeout* seconds for tokens to become available.
        Returns True on success, False on timeout.  If *timeout* is None
        the call blocks indefinitely.
        """
        deadline = None if timeout is None else time.monotonic() + timeout

        while True:
            with self._lock:
                self._refill()
                if self._tokens >= tokens:
                    self._tokens -= tokens
                    return True
                # How long until enough tokens are replenished?
                deficit = tokens - self._tokens
                wait = deficit / self._rate if self._rate > 0 else float("inf")

            if deadline is not None:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    return False
                wait = min(wait, remaining)

            time.sleep(min(wait, 0.05))  # cap sleep granularity

    def available(self) -> int:
        """Return the current number of available tokens (floor)."""
        with self._lock:
            self._refill()
            return int(self._tokens)


# ---------------------------------------------------------------------------
# Sliding Window Counter
# ---------------------------------------------------------------------------

class SlidingWindowCounter:
    """Count events in a sliding time window.

    Uses two sub-buckets (current + previous) for a smooth approximation
    of the true sliding window count without storing every timestamp.
    """

    def __init__(
        self, window_seconds: float = 60.0, max_count: int = 1000
    ) -> None:
        self._window = float(window_seconds)
        self._max_count = int(max_count)
        self._lock = threading.Lock()

        # Two half-window buckets
        self._cur_count = 0
        self._prev_count = 0
        self._cur_start = time.monotonic()

    # -- internal ----------------------------------------------------------

    def _rotate(self) -> None:
        now = time.monotonic()
        elapsed = now - self._cur_start
        half = self._window / 2.0
        if elapsed >= half:
            # How many full half-windows have elapsed?
            rotations = int(elapsed / half)
            if rotations >= 2:
                self._prev_count = 0
                self._cur_count = 0
            else:
                self._prev_count = self._cur_count
                self._cur_count = 0
            self._cur_start = now

    # -- public API --------------------------------------------------------

    def record(self) -> bool:
        """Record an event.  Returns False if *max_count* is exceeded."""
        with self._lock:
            self._rotate()
            if self._estimated_count_unlocked() >= self._max_count:
                return False
            self._cur_count += 1
            return True

    def count(self) -> int:
        """Estimated event count in the current sliding window."""
        with self._lock:
            self._rotate()
            return self._estimated_count_unlocked()

    def is_exceeded(self) -> bool:
        """True if the event count has reached or exceeded *max_count*."""
        return self.count() >= self._max_count

    # -- helpers -----------------------------------------------------------

    def _estimated_count_unlocked(self) -> int:
        elapsed = time.monotonic() - self._cur_start
        half = self._window / 2.0
        # Weight the previous bucket by how much of it is still in-window
        weight = max(0.0, 1.0 - elapsed / half) if half > 0 else 0.0
        return self._cur_count + int(self._prev_count * weight)


# ---------------------------------------------------------------------------
# Backpressure Controller
# ---------------------------------------------------------------------------

class BackpressureController:
    """Adaptive backpressure based on observed latency.

    Maintains a concurrency semaphore that auto-adjusts: when latencies
    exceed the target the controller *reduces* max_concurrent; when
    latencies are well below the target it *increases* max_concurrent
    (up to the configured ceiling).
    """

    _MIN_CONCURRENT = 1
    _ADJUST_INTERVAL = 1.0  # seconds between adjustments
    _INCREASE_FACTOR = 1.25
    _DECREASE_FACTOR = 0.75

    def __init__(
        self, target_latency_ms: float = 100.0, max_concurrent: int = 10
    ) -> None:
        self._target_latency = float(target_latency_ms)
        self._ceiling = int(max_concurrent)
        self._max_concurrent = int(max_concurrent)
        self._in_flight = 0
        self._lock = threading.Lock()

        # Latency tracking
        self._recent_latencies: list[float] = []
        self._last_adjust = time.monotonic()

    # -- public API --------------------------------------------------------

    def acquire(self) -> bool:
        """Try to acquire a concurrency slot.  Returns False if full."""
        with self._lock:
            if self._in_flight >= self._max_concurrent:
                return False
            self._in_flight += 1
            return True

    def release(self, latency_ms: float) -> None:
        """Release a slot and report the observed latency."""
        with self._lock:
            self._in_flight = max(0, self._in_flight - 1)
            self._recent_latencies.append(latency_ms)
            self._maybe_adjust()

    def current_concurrency(self) -> int:
        """Return the number of in-flight operations."""
        with self._lock:
            return self._in_flight

    def should_shed_load(self) -> bool:
        """True if the system is overloaded and callers should back off."""
        with self._lock:
            if self._in_flight >= self._max_concurrent:
                return True
            if len(self._recent_latencies) >= 3:
                avg = sum(self._recent_latencies[-3:]) / 3.0
                return avg > self._target_latency * 2.0
            return False

    # -- internal ----------------------------------------------------------

    def _maybe_adjust(self) -> None:
        """Adjust max_concurrent based on recent latencies (lock held)."""
        now = time.monotonic()
        if now - self._last_adjust < self._ADJUST_INTERVAL:
            return
        if len(self._recent_latencies) < 3:
            return

        avg = sum(self._recent_latencies) / len(self._recent_latencies)
        self._recent_latencies.clear()
        self._last_adjust = now

        if avg > self._target_latency * 1.5:
            # Latency too high — reduce concurrency
            self._max_concurrent = max(
                self._MIN_CONCURRENT,
                int(self._max_concurrent * self._DECREASE_FACTOR),
            )
        elif avg < self._target_latency * 0.5:
            # Latency well below target — allow more concurrency
            self._max_concurrent = min(
                self._ceiling,
                max(self._max_concurrent + 1, int(self._max_concurrent * self._INCREASE_FACTOR)),
            )


# ---------------------------------------------------------------------------
# Throttle (combined facade)
# ---------------------------------------------------------------------------

class Throttle:
    """Combines rate limiting and backpressure for query and index paths.

    Typical usage::

        throttle = Throttle()
        if throttle.allow_query():
            do_query()
        if throttle.allow_index():
            do_index()
    """

    def __init__(
        self,
        queries_per_second: float = 100,
        max_concurrent_queries: int = 10,
        max_concurrent_indexing: int = 2,
    ) -> None:
        self._query_bucket = TokenBucket(
            rate=queries_per_second, capacity=int(queries_per_second)
        )
        self._query_bp = BackpressureController(
            target_latency_ms=100.0, max_concurrent=max_concurrent_queries
        )
        self._index_bp = BackpressureController(
            target_latency_ms=500.0, max_concurrent=max_concurrent_indexing
        )
        self._query_window = SlidingWindowCounter(
            window_seconds=60.0, max_count=int(queries_per_second * 60)
        )
        self._index_window = SlidingWindowCounter(
            window_seconds=60.0, max_count=max_concurrent_indexing * 60
        )

    # -- public API --------------------------------------------------------

    def allow_query(self) -> bool:
        """Return True if a new query is permitted right now."""
        if not self._query_bucket.try_acquire():
            return False
        if not self._query_bp.acquire():
            return False
        self._query_window.record()
        return True

    def allow_index(self) -> bool:
        """Return True if a new indexing operation is permitted right now."""
        if not self._index_bp.acquire():
            return False
        self._index_window.record()
        return True

    def release_query(self, latency_ms: float) -> None:
        """Release a query slot and report latency for adaptive tuning."""
        self._query_bp.release(latency_ms)

    def release_index(self, latency_ms: float) -> None:
        """Release an indexing slot and report latency for adaptive tuning."""
        self._index_bp.release(latency_ms)

    def stats(self) -> Dict[str, object]:
        """Return a snapshot of throttle statistics."""
        return {
            "query_tokens_available": self._query_bucket.available(),
            "query_concurrency": self._query_bp.current_concurrency(),
            "query_window_count": self._query_window.count(),
            "query_overloaded": self._query_bp.should_shed_load(),
            "index_concurrency": self._index_bp.current_concurrency(),
            "index_window_count": self._index_window.count(),
            "index_overloaded": self._index_bp.should_shed_load(),
        }
