"""A/B rollout controller for ModernBERT adapter weight variants.

``ABRollout`` implements a deterministic traffic split between the incumbent
(control) adapter and a newly trained (treatment) adapter.  The split is
determined by hashing the request identifier so the same request always goes
to the same variant — enabling consistent comparison while gradually rolling
out new weights.

Thread-safe: all metric recording uses a ``threading.Lock``.
"""

from __future__ import annotations

import hashlib
import logging
import threading
from typing import Any

logger = logging.getLogger(__name__)


class ABRollout:
    """Deterministic A/B traffic splitter for router weight variants.

    Implements a hash-modulo split: ``hash(request_hash) % 100 < split_pct * 100``
    routes to the treatment variant (new weights); all other traffic stays on
    the control (current production) weights.

    Args:
        split_pct: Fraction of traffic to route to the treatment variant.
            Must be in the range [0.0, 1.0].  Defaults to 0.10 (10%).

    Usage::

        rollout = ABRollout(split_pct=0.10)
        use_new = rollout.should_use_new_weights(request_id)
        variant = "treatment" if use_new else "control"
        # ... run inference ...
        rollout.record_result(variant, latency_ms=32.5, success=True)
        metrics = rollout.get_metrics()
    """

    _CONTROL: str = "control"
    _TREATMENT: str = "treatment"

    def __init__(self, split_pct: float = 0.10) -> None:
        if not 0.0 <= split_pct <= 1.0:
            raise ValueError(f"split_pct must be in [0.0, 1.0], got {split_pct}")
        self._split_pct = split_pct
        self._lock = threading.Lock()
        # Per-variant accumulators
        self._counts: dict[str, int] = {self._CONTROL: 0, self._TREATMENT: 0}
        self._successes: dict[str, int] = {self._CONTROL: 0, self._TREATMENT: 0}
        self._latencies: dict[str, list[float]] = {self._CONTROL: [], self._TREATMENT: []}

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def should_use_new_weights(self, request_hash: str) -> bool:
        """Return True when this request should use the treatment (new) weights.

        The decision is deterministic: the same ``request_hash`` always returns
        the same result.  This ensures consistent routing across retries and
        simplifies A/B analysis.

        Args:
            request_hash: An opaque identifier for the request (e.g. a UUID or
                query hash).  Only its hash modulo 100 is used.

        Returns:
            True when the request falls in the treatment bucket.
        """
        # Hash the request_hash to an integer bucket 0-99
        digest = hashlib.md5(request_hash.encode(), usedforsecurity=False).hexdigest()
        bucket = int(digest[:8], 16) % 100
        threshold = int(self._split_pct * 100)
        return bucket < threshold

    def record_result(self, variant: str, latency_ms: float, success: bool) -> None:
        """Record a routing outcome for the given variant.

        Args:
            variant: Either ``"control"`` or ``"treatment"``.
            latency_ms: End-to-end inference latency in milliseconds.
            success: Whether the request succeeded (model returned a valid response).
        """
        if variant not in (self._CONTROL, self._TREATMENT):
            logger.warning("Unknown A/B variant '%s'; ignoring.", variant)
            return

        with self._lock:
            self._counts[variant] += 1
            if success:
                self._successes[variant] += 1
            self._latencies[variant].append(latency_ms)

    def get_metrics(self) -> dict[str, Any]:
        """Return per-variant statistics.

        Returns a dict of the form::

            {
                "control": {
                    "requests": N,
                    "success_rate": 0.98,
                    "p50_latency_ms": 12.4,
                },
                "treatment": {
                    "requests": M,
                    "success_rate": 0.97,
                    "p50_latency_ms": 11.8,
                },
            }

        Latency percentiles are computed over all recorded observations.
        Returns 0.0 for latency when no observations have been recorded.
        """
        with self._lock:
            result: dict[str, Any] = {}
            for variant in (self._CONTROL, self._TREATMENT):
                n = self._counts[variant]
                successes = self._successes[variant]
                latencies = sorted(self._latencies[variant])

                success_rate = successes / n if n > 0 else 0.0
                p50 = _percentile(latencies, 0.50)

                result[variant] = {
                    "requests": n,
                    "success_rate": round(success_rate, 4),
                    "p50_latency_ms": round(p50, 2),
                }
            return result

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def split_pct(self) -> float:
        """The current treatment traffic fraction."""
        return self._split_pct


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _percentile(sorted_values: list[float], p: float) -> float:
    """Return the ``p``-th percentile of a pre-sorted list.

    Uses the nearest-rank method: ``ceil(n * p) - 1`` (0-indexed).

    Args:
        sorted_values: Sorted list of floats.
        p: Percentile in [0, 1].

    Returns:
        The percentile value, or 0.0 if the list is empty.
    """
    if not sorted_values:
        return 0.0
    import math

    n = len(sorted_values)
    # Nearest-rank: ordinal rank = ceil(n * p), convert to 0-indexed
    idx = max(0, math.ceil(n * p) - 1)
    return sorted_values[min(idx, n - 1)]
