"""Tests for SelectiveVerificationRouter."""
import time
from unittest.mock import MagicMock, call

import pytest


def _make_lf(risk: float = 0.30):
    """Mock LabelFreeResult with the attributes SelectiveVerificationRouter reads."""
    lf = MagicMock()
    lf.risk_score = risk
    lf.failure_mode = None
    lf.components = {"sc2": 0.3, "sc3": 0.4}
    return lf


def _make_chain_result(risk: float = 0.30):
    """Mock ChainTrustResult returned by AgentGuard.score_chain()."""
    from llm_guard.agent_guard import ChainTrustResult
    tier = "HIGH" if risk < 0.50 else ("MEDIUM" if risk < 0.70 else "LOW")
    return ChainTrustResult(
        risk_score=risk,
        confidence_tier=tier,
        needs_alert=risk >= 0.70,
        failure_mode=None,
        judge_label=None,
        step_count=3,
        behavioral_score=risk,
        behavioral_components={"sc2": risk},
        latency_ms=1.0,
    )


def _make_guard(scorer_risk: float = 0.30, chain_risk: float = 0.30):
    """Mock AgentGuard with _scorer and score_chain."""
    guard = MagicMock()
    guard._scorer = MagicMock()
    guard._scorer.score.return_value = _make_lf(risk=scorer_risk)
    guard.score_chain.return_value = _make_chain_result(risk=chain_risk)
    guard._alert_threshold = 0.70
    return guard


STEPS = [{"thought": "t", "action_type": "search", "action_arg": "x", "observation": "o"}]


class TestSelectiveVerificationRouterRouting:
    def test_high_confidence_bypasses_guard(self):
        guard = _make_guard(scorer_risk=0.30)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        result = router.route("Q?", STEPS, "A")
        guard.score_chain.assert_not_called()
        assert result.confidence_tier == "HIGH"
        assert result.risk_score == pytest.approx(0.30)

    def test_medium_confidence_uses_full_pipeline(self):
        guard = _make_guard(scorer_risk=0.55, chain_risk=0.55)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        result = router.route("Q?", STEPS, "A")
        guard.score_chain.assert_called_once()
        assert result.confidence_tier == "MEDIUM"

    def test_low_confidence_uses_full_pipeline(self):
        guard = _make_guard(scorer_risk=0.75, chain_risk=0.75)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        result = router.route("Q?", STEPS, "A")
        guard.score_chain.assert_called_once()

    def test_empty_steps_always_full_pipeline(self):
        guard = _make_guard(scorer_risk=0.10)  # would be bypassed normally
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        result = router.route("Q?", [], "A")
        guard.score_chain.assert_called_once()
        # _scorer.score should NOT be called for empty steps
        guard._scorer.score.assert_not_called()

    def test_score_chain_alias_identical_to_route(self):
        guard = _make_guard(scorer_risk=0.30)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        r1 = router.route("Q?", STEPS, "A")
        # Reset mocks
        guard._scorer.score.return_value = _make_lf(risk=0.30)
        guard.score_chain.return_value = _make_chain_result(risk=0.30)
        r2 = router.score_chain("Q?", STEPS, "A")
        assert r1.risk_score == r2.risk_score
        assert r1.confidence_tier == r2.confidence_tier

    def test_scorer_reuses_guard_scorer_not_new_instance(self):
        guard = _make_guard(scorer_risk=0.30)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        router.route("Q?", STEPS, "A")
        guard._scorer.score.assert_called_once()

    def test_needs_alert_reflects_alert_threshold(self):
        guard = _make_guard(scorer_risk=0.45)
        guard._scorer.score.return_value = _make_lf(risk=0.45)
        guard._alert_threshold = 0.40
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard, high_threshold=0.50)
        result = router.route("Q?", STEPS, "A")
        assert result.needs_alert is True  # 0.45 >= 0.40
        guard.score_chain.assert_not_called()


class TestSelectiveVerificationRouterStats:
    def test_routing_stats_empty_at_start(self):
        guard = _make_guard()
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        stats = router.routing_stats()
        assert stats["bypassed"] == 0
        assert stats["full_pipeline"] == 0
        assert stats["bypass_rate"] == 0.0

    def test_routing_stats_tracks_counts(self):
        guard = _make_guard(scorer_risk=0.30, chain_risk=0.30)
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        # 2 HIGH (bypassed), 1 MEDIUM (full)
        router.route("Q?", STEPS, "A")
        router.route("Q?", STEPS, "A")
        guard._scorer.score.return_value = _make_lf(risk=0.60)
        guard.score_chain.return_value = _make_chain_result(risk=0.60)
        router.route("Q?", STEPS, "A")
        stats = router.routing_stats()
        assert stats["bypassed"] == 2
        assert stats["full_pipeline"] == 1
        assert stats["bypass_rate"] == pytest.approx(2 / 3)


class TestGuardPipelineCompatibility:
    def test_score_chain_method_exists(self):
        guard = _make_guard()
        from llm_guard.selective_router import SelectiveVerificationRouter
        router = SelectiveVerificationRouter(guard=guard)
        assert callable(getattr(router, "score_chain", None))

    def test_works_as_guard_pipeline_stage(self):
        guard = _make_guard(scorer_risk=0.30)
        from llm_guard.selective_router import SelectiveVerificationRouter
        from llm_guard.guard_pipeline import GuardPipeline
        router = SelectiveVerificationRouter(guard=guard)
        pipeline = GuardPipeline().add_stage("guard", router)
        result = pipeline.run("Q?", STEPS, "A")
        assert result.final_risk == pytest.approx(0.30, abs=0.05)
        assert result.action in ("proceed", "rewrite", "escalate", "abort")
