"""
SelectiveVerificationRouter — cost-proportional guard routing.

Routes chains based on a cheap behavioral pre-score. HIGH-confidence chains
(risk < high_threshold) are returned immediately without API calls. MEDIUM/LOW
chains are delegated to the full AgentGuard pipeline.

Standalone: no API key required for the router itself. The provided ``guard``
instance may or may not have a key — only MEDIUM/LOW chains will trigger API calls.
Purpose:    Returns HIGH-confidence chains immediately without API calls; delegates MEDIUM/LOW chains to the full AgentGuard pipeline. Reduces judge API costs proportional to the fraction of clearly-safe chains.
Exports:    SelectiveVerificationRouter
Added:      v0.1.0
Status:     Stable
"""
from __future__ import annotations

import time
from typing import Dict

from llm_guard.scoring.agent_guard import AgentGuard, ChainTrustResult  # no circular dep


class SelectiveVerificationRouter:
    """
    Routes guard calls based on a cheap behavioral pre-score.

    HIGH confidence (behavioral risk < high_threshold, non-empty steps):
        Returns behavioral result immediately. No judge or P(True) API calls made.
    MEDIUM / LOW confidence, or empty steps:
        Delegates to the provided AgentGuard instance (full pipeline).

    Exposes ``score_chain()`` as an alias for ``route()``, making this class
    a drop-in replacement for AgentGuard as a GuardPipeline stage.

    Parameters
    ----------
    guard : AgentGuard
        Configured AgentGuard. Use use_judge=True or score_with_ptrue() for
        maximum accuracy on escalated chains.
    high_threshold : float, default 0.50
        Behavioral risk below which a chain is considered HIGH-confidence.
        Matches the HIGH/MEDIUM boundary used by AgentGuard._risk_to_tier().
    """

    def __init__(self, guard: AgentGuard, high_threshold: float = 0.50) -> None:
        self._guard = guard
        self._high_threshold = high_threshold
        self._stats: Dict[str, int] = {"bypassed": 0, "full_pipeline": 0}

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

    def route(
        self,
        question: str,
        steps: list,
        final_answer: str,
        finished: bool = True,
    ) -> ChainTrustResult:
        """
        Score a chain with cost-proportional routing.

        Empty ``steps`` always escalate to the full pipeline regardless of score
        (no behavioral signal available for unanswered queries).

        Returns a standard ChainTrustResult. Use routing_stats() to observe
        the bypass rate over time.
        """
        t0 = time.perf_counter()

        # Empty steps: no behavioral signal — always escalate
        if not steps:
            result = self._guard.score_chain(question, steps, final_answer, finished=finished)
            self._stats["full_pipeline"] += 1
            return result

        # Behavioral pre-score: reuses the same fitted LabelFreeScorer instance
        # as the full score_chain() path, so any GMM/obs-pool calibration is reflected.
        lf = self._guard._scorer.score(question, steps, final_answer, finished)

        if lf.risk_score < self._high_threshold:
            latency_ms = (time.perf_counter() - t0) * 1000
            result = self._build_result(lf, steps, latency_ms)
            self._stats["bypassed"] += 1
            return result

        # Full pipeline
        result = self._guard.score_chain(question, steps, final_answer, finished=finished)
        self._stats["full_pipeline"] += 1
        return result

    def score_chain(
        self,
        question: str,
        steps: list,
        final_answer: str,
        finished: bool = True,
    ) -> ChainTrustResult:
        """
        Alias for route(). Implements the score_chain() interface required by
        GuardPipeline._dispatch(), making this a drop-in stage replacement.
        """
        return self.route(question, steps, final_answer, finished=finished)

    def routing_stats(self) -> dict:
        """
        Returns {
            "bypassed":      int,
            "full_pipeline": int,
            "bypass_rate":   float,  # 0.0 if no calls made yet
        }.

        Use this in production to measure actual cost savings.
        Note: expected ~40-50% bypass rate on HotpotQA/TriviaQA distributions;
        actual rates vary by domain.
        """
        total = self._stats["bypassed"] + self._stats["full_pipeline"]
        return {
            "bypassed": self._stats["bypassed"],
            "full_pipeline": self._stats["full_pipeline"],
            "bypass_rate": self._stats["bypassed"] / total if total > 0 else 0.0,
        }

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _build_result(self, lf, steps: list, latency_ms: float) -> ChainTrustResult:
        """Construct a ChainTrustResult from a LabelFreeResult (bypassed path)."""
        risk = lf.risk_score
        if risk < 0.50:
            tier = "HIGH"
        elif risk < 0.70:
            tier = "MEDIUM"
        else:
            tier = "LOW"

        components = dict(lf.components)
        components["routing"] = 0.0  # 0.0 = bypassed; 1.0 = full_pipeline
        return ChainTrustResult(
            risk_score=risk,
            confidence_tier=tier,
            needs_alert=risk >= self._guard._alert_threshold,
            failure_mode=lf.failure_mode,
            judge_label=None,
            step_count=len(steps),
            behavioral_score=risk,
            behavioral_components=components,
            latency_ms=latency_ms,
        )
