"""Load manager — request queuing, rate limiting, and throughput control.

Manages concurrency and QPS limits per provider, distributes load across
provider instances, tracks queue depth and latency percentiles, and
exposes SLO metrics. Stdlib-only (ADR-021).
"""

from __future__ import annotations

import queue
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field


@dataclass
class LoadConfig:
    """Configuration for load management."""

    max_qps: float = 50.0
    max_concurrent: int = 10
    queue_capacity: int = 1000
    queue_timeout_s: float = 30.0
    slo_p99_ms: float = 5000.0

    def to_dict(self) -> dict:
        return {
            "max_qps": self.max_qps,
            "max_concurrent": self.max_concurrent,
            "queue_capacity": self.queue_capacity,
            "queue_timeout_s": self.queue_timeout_s,
            "slo_p99_ms": self.slo_p99_ms,
        }


@dataclass
class RequestMetrics:
    """Metrics for a single completed request."""

    provider: str
    latency_ms: float
    queued_ms: float
    success: bool
    timestamp: float = field(default_factory=time.monotonic)


@dataclass
class LoadStats:
    """Runtime statistics for the load manager."""

    total_requests: int = 0
    total_completed: int = 0
    total_rejected: int = 0
    total_timeouts: int = 0
    total_errors: int = 0
    current_queue_depth: int = 0
    current_concurrent: int = 0
    latency_p50_ms: float = 0.0
    latency_p95_ms: float = 0.0
    latency_p99_ms: float = 0.0
    avg_queue_wait_ms: float = 0.0
    slo_violations: int = 0
    requests_per_second: float = 0.0

    def to_dict(self) -> dict:
        return {
            "total_requests": self.total_requests,
            "total_completed": self.total_completed,
            "total_rejected": self.total_rejected,
            "total_timeouts": self.total_timeouts,
            "total_errors": self.total_errors,
            "current_queue_depth": self.current_queue_depth,
            "current_concurrent": self.current_concurrent,
            "latency_p50_ms": round(self.latency_p50_ms, 2),
            "latency_p95_ms": round(self.latency_p95_ms, 2),
            "latency_p99_ms": round(self.latency_p99_ms, 2),
            "avg_queue_wait_ms": round(self.avg_queue_wait_ms, 2),
            "slo_violations": self.slo_violations,
            "requests_per_second": round(self.requests_per_second, 2),
        }


class RateLimiter:
    """Token bucket rate limiter for QPS control.

    Allows max_qps requests per second with burst capacity.
    """

    def __init__(self, max_qps: float = 50.0) -> None:
        self._max_qps = max(0.1, max_qps)
        self._tokens = max_qps
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()

    def acquire(self, timeout_s: float = 1.0) -> bool:
        """Try to acquire a token. Returns True if allowed."""
        deadline = time.monotonic() + timeout_s
        while time.monotonic() < deadline:
            with self._lock:
                self._refill()
                if self._tokens >= 1.0:
                    self._tokens -= 1.0
                    return True
            time.sleep(0.01)
        return False

    def _refill(self) -> None:
        """Refill tokens based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self._last_refill
        self._tokens = min(
            self._max_qps,
            self._tokens + elapsed * self._max_qps,
        )
        self._last_refill = now

    @property
    def available_tokens(self) -> float:
        with self._lock:
            self._refill()
            return self._tokens


class ConcurrencyLimiter:
    """Semaphore-based concurrency limiter."""

    def __init__(self, max_concurrent: int = 10) -> None:
        self._semaphore = threading.Semaphore(max_concurrent)
        self._max = max_concurrent
        self._active = 0
        self._lock = threading.Lock()

    def acquire(self, timeout_s: float = 30.0) -> bool:
        """Acquire a concurrency slot."""
        acquired = self._semaphore.acquire(timeout=timeout_s)
        if acquired:
            with self._lock:
                self._active += 1
        return acquired

    def release(self) -> None:
        """Release a concurrency slot."""
        with self._lock:
            self._active = max(0, self._active - 1)
        self._semaphore.release()

    @property
    def active(self) -> int:
        with self._lock:
            return self._active

    @property
    def available(self) -> int:
        with self._lock:
            return self._max - self._active


class LoadManager:
    """Central load management with queuing, rate limiting, and SLO tracking.

    Manages request flow to providers with:
    - Token bucket rate limiting (QPS)
    - Semaphore concurrency limiting
    - Request queuing with backpressure
    - Latency percentile tracking (p50/p95/p99)
    - SLO violation monitoring

    Usage::

        mgr = LoadManager(LoadConfig(max_qps=100, max_concurrent=10))

        # Before calling provider
        slot = mgr.acquire()
        if slot:
            try:
                response = provider.generate(prompt)
                mgr.record_success("vllm", latency_ms)
            except Exception:
                mgr.record_error("vllm")
            finally:
                mgr.release()
        else:
            # Queue full or timeout
            handle_overload()
    """

    def __init__(self, config: LoadConfig | None = None) -> None:
        self._config = config or LoadConfig()
        self._rate_limiter = RateLimiter(self._config.max_qps)
        self._concurrency = ConcurrencyLimiter(self._config.max_concurrent)
        self._queue: queue.Queue = queue.Queue(maxsize=self._config.queue_capacity)
        self._metrics: list[RequestMetrics] = []
        self._lock = threading.Lock()
        self._total_requests = 0
        self._total_rejected = 0
        self._total_timeouts = 0
        self._total_errors = 0
        self._started_at = time.monotonic()
        self._max_metrics = 10000

    def acquire(self, timeout_s: float | None = None) -> bool:
        """Acquire both rate limit token and concurrency slot.

        Returns True if the request is allowed to proceed.
        """
        t = timeout_s if timeout_s is not None else self._config.queue_timeout_s

        with self._lock:
            self._total_requests += 1

        if not self._rate_limiter.acquire(timeout_s=min(t, 1.0)):
            with self._lock:
                self._total_rejected += 1
            return False

        if not self._concurrency.acquire(timeout_s=t):
            with self._lock:
                self._total_timeouts += 1
            return False

        return True

    def release(self) -> None:
        """Release concurrency slot after request completes."""
        self._concurrency.release()

    def record_success(
        self,
        provider: str,
        latency_ms: float,
        queued_ms: float = 0.0,
    ) -> None:
        """Record a successful request."""
        metric = RequestMetrics(
            provider=provider,
            latency_ms=latency_ms,
            queued_ms=queued_ms,
            success=True,
        )
        with self._lock:
            self._metrics.append(metric)
            if len(self._metrics) > self._max_metrics:
                self._metrics = self._metrics[-self._max_metrics :]

    def record_error(
        self,
        provider: str,
        latency_ms: float = 0.0,
    ) -> None:
        """Record a failed request."""
        metric = RequestMetrics(
            provider=provider,
            latency_ms=latency_ms,
            queued_ms=0.0,
            success=False,
        )
        with self._lock:
            self._total_errors += 1
            self._metrics.append(metric)
            if len(self._metrics) > self._max_metrics:
                self._metrics = self._metrics[-self._max_metrics :]

    def _percentile(self, values: list[float], pct: float) -> float:
        """Compute percentile from sorted values."""
        if not values:
            return 0.0
        sorted_v = sorted(values)
        idx = int(len(sorted_v) * pct / 100)
        idx = min(idx, len(sorted_v) - 1)
        return sorted_v[idx]

    @property
    def stats(self) -> LoadStats:
        """Compute current load statistics."""
        with self._lock:
            metrics = list(self._metrics)
            total_req = self._total_requests
            total_rej = self._total_rejected
            total_to = self._total_timeouts
            total_err = self._total_errors

        latencies = [m.latency_ms for m in metrics if m.success]
        queue_waits = [m.queued_ms for m in metrics]
        completed = len(metrics)

        elapsed = time.monotonic() - self._started_at
        rps = completed / elapsed if elapsed > 0 else 0.0

        slo_violations = sum(1 for m in metrics if m.latency_ms > self._config.slo_p99_ms)

        return LoadStats(
            total_requests=total_req,
            total_completed=completed,
            total_rejected=total_rej,
            total_timeouts=total_to,
            total_errors=total_err,
            current_queue_depth=self._queue.qsize(),
            current_concurrent=self._concurrency.active,
            latency_p50_ms=self._percentile(latencies, 50),
            latency_p95_ms=self._percentile(latencies, 95),
            latency_p99_ms=self._percentile(latencies, 99),
            avg_queue_wait_ms=(sum(queue_waits) / len(queue_waits) if queue_waits else 0.0),
            slo_violations=slo_violations,
            requests_per_second=rps,
        )

    def provider_stats(self) -> dict[str, dict]:
        """Per-provider breakdown of metrics."""
        with self._lock:
            metrics = list(self._metrics)

        by_provider: dict[str, list[RequestMetrics]] = defaultdict(list)
        for m in metrics:
            by_provider[m.provider].append(m)

        result = {}
        for name, provmetrics in sorted(by_provider.items()):
            latencies = [m.latency_ms for m in provmetrics if m.success]
            errors = sum(1 for m in provmetrics if not m.success)
            result[name] = {
                "total": len(provmetrics),
                "successes": len(provmetrics) - errors,
                "errors": errors,
                "avg_latency_ms": (round(sum(latencies) / len(latencies), 2) if latencies else 0.0),
                "p95_latency_ms": round(self._percentile(latencies, 95), 2),
            }
        return result

    def summary(self) -> dict:
        """Complete summary for monitoring."""
        return {
            "config": self._config.to_dict(),
            "stats": self.stats.to_dict(),
            "providers": self.provider_stats(),
            "rate_limiter_tokens": round(self._rate_limiter.available_tokens, 1),
            "concurrency_available": self._concurrency.available,
        }

    def reset(self) -> None:
        """Reset all metrics and counters."""
        with self._lock:
            self._metrics.clear()
            self._total_requests = 0
            self._total_rejected = 0
            self._total_timeouts = 0
            self._total_errors = 0
            self._started_at = time.monotonic()


__all__ = [
    "ConcurrencyLimiter",
    "LoadConfig",
    "LoadManager",
    "LoadStats",
    "RateLimiter",
    "RequestMetrics",
]
