"""
Tests for llm-guard-kit v0.29.0 new components:
  - ConformalRiskBudget
  - QuestionComplexityScorer
  - DebateVerifier
  - GuardPhaseController / PhaseStatus
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import pytest
import numpy as np
from unittest.mock import patch, MagicMock


# ── ConformalRiskBudget ───────────────────────────────────────────────────────

from llm_guard.calibration.conformal_risk_budget import ConformalRiskBudget, BudgetState

class TestConformalRiskBudget:

    def test_bonferroni_threshold(self):
        b = ConformalRiskBudget(0.15, "bonferroni")
        t = b.allocate(5)
        assert len(t) == 5
        assert abs(t[0] - 0.15/5) < 1e-9

    def test_product_threshold(self):
        b = ConformalRiskBudget(0.20, "product")
        t = b.allocate(5)
        expected = 1 - (1-0.20)**(1/5)
        assert abs(t[0] - expected) < 1e-9

    def test_product_no_abort_on_safe_steps(self):
        b = ConformalRiskBudget(0.20, "product")
        b.allocate(5)
        aborts = [b.observe(i, 0.03) for i in range(5)]
        assert not any(aborts)

    def test_abort_triggered_at_correct_step(self):
        b = ConformalRiskBudget(0.20, "product")
        b.allocate(5)
        # cumsum: 0.04, 0.10, 0.14, 0.23 → abort at step 3
        risks = [0.04, 0.06, 0.04, 0.09, 0.01]
        aborts = [b.observe(i, r) for i, r in enumerate(risks)]
        assert aborts == [False, False, False, True, True]
        assert b.summary()["aborted_at"] == 3

    def test_adaptive_abort_on_spike(self):
        b = ConformalRiskBudget(0.20, "adaptive")
        b.allocate(5)
        abort0 = b.observe(0, 0.05)
        abort1 = b.observe(1, 0.16)   # cumulative 0.21 > 0.20
        assert not abort0
        assert abort1

    def test_reset(self):
        b = ConformalRiskBudget(0.20, "product")
        b.allocate(5)
        b.observe(0, 0.25)
        b.reset()
        assert b.remaining_budget == pytest.approx(0.20)
        assert b._history == []

    def test_invalid_budget_raises(self):
        with pytest.raises(ValueError):
            ConformalRiskBudget(1.5)
        with pytest.raises(ValueError):
            ConformalRiskBudget(0.0)

    def test_allocate_before_observe_required(self):
        b = ConformalRiskBudget(0.20)
        with pytest.raises(RuntimeError):
            b.observe(0, 0.1)

    def test_summary_structure(self):
        b = ConformalRiskBudget(0.10, "bonferroni")
        b.allocate(3)
        b.observe(0, 0.03)
        s = b.summary()
        assert "total_budget" in s
        assert "cumulative_used" in s
        assert "remaining" in s
        assert "aborted_at" in s

    def test_single_step(self):
        b = ConformalRiskBudget(0.30, "bonferroni")
        b.allocate(1)
        assert not b.observe(0, 0.25)
        assert b.observe(1, 0.25)   # now over 0.30


# ── QuestionComplexityScorer ──────────────────────────────────────────────────

from llm_guard.adversarial.novelty_detector import QuestionComplexityScorer, _complexity_features, ComplexityResult

HP_QUESTIONS = [
    {"question": "Which film came first, the French comedy film or the American science fiction film directed by James Cameron?"},
    {"question": "Are both Cangzhou and Qinhuangdao located in the same province?"},
    {"question": "Which magazine was started first, Arthur's Magazine or First for Women?"},
    {"question": "What was the nationality of the person who invented the telephone?"},
    {"question": "In which country was the actress who played the Countess in The Grand Budapest Hotel born?"},
    {"question": "Were Scott Derrickson and Ed Wood of the same nationality?"},
    {"question": "What is the capital of the country where the person who invented the telephone was born?"},
    {"question": "Was the film starring the actor who appeared in Inception directed by Christopher Nolan?"},
    {"question": "Who was the director of the 2010 film that starred the actor born in Los Angeles?"},
    {"question": "What year did the person who founded the company that made the iPhone die?"},
]

TV_QUESTIONS = [
    {"question": "What year did World War II end?"},
    {"question": "Who invented the telephone?"},
    {"question": "What is the capital of France?"},
    {"question": "How many sides does a hexagon have?"},
    {"question": "Who wrote Romeo and Juliet?"},
]

class TestQuestionComplexityScorer:

    def test_fit_and_score(self):
        scorer = QuestionComplexityScorer()
        scorer.fit(HP_QUESTIONS * 5)   # pad to get decent cal set
        result = scorer.score(HP_QUESTIONS[0])
        assert isinstance(result, ComplexityResult)
        assert 0.0 <= result.complexity_score <= 1.0
        assert result.routing in ("trusted", "uncertain", "force_judge")
        assert result.n_cal == len(HP_QUESTIONS) * 5

    def test_hp_scores_higher_than_tv(self):
        scorer = QuestionComplexityScorer()
        scorer.fit(HP_QUESTIONS * 10)
        hp_scores = [scorer.score(q).complexity_score for q in HP_QUESTIONS]
        tv_scores = [scorer.score(q).complexity_score for q in TV_QUESTIONS]
        assert np.mean(hp_scores) > np.mean(tv_scores)

    def test_complexity_features_shape(self):
        vec = _complexity_features("Who was the director of the film starring the actor born in LA?")
        assert vec.shape == (7,)
        assert all(0.0 <= v <= 3.0 for v in vec)

    def test_complexity_features_simple_question(self):
        vec = _complexity_features("What year?")
        # short question → low token_count_norm and char_len_norm
        assert vec[0] < 0.3   # token_count_norm
        assert vec[6] < 0.2   # char_len_norm

    def test_complexity_features_complex_question(self):
        vec = _complexity_features(
            "Which city was the birthplace of the person who founded the company "
            "that made the device used by the actor who won the Oscar in 2019?"
        )
        assert vec[0] > 0.8   # long question

    def test_fit_requires_nonempty(self):
        scorer = QuestionComplexityScorer()
        with pytest.raises(Exception):
            scorer.score(HP_QUESTIONS[0])  # not fitted

    def test_score_batch(self):
        scorer = QuestionComplexityScorer()
        scorer.fit(HP_QUESTIONS * 5)
        results = scorer.score_batch(HP_QUESTIONS)
        assert len(results) == len(HP_QUESTIONS)

    def test_fit_with_ood(self):
        scorer = QuestionComplexityScorer()
        scorer.fit_with_ood(HP_QUESTIONS * 5, TV_QUESTIONS * 5)
        # After supervised fit, HP should score higher than TV
        hp_s = [scorer.score(q).complexity_score for q in HP_QUESTIONS]
        tv_s = [scorer.score(q).complexity_score for q in TV_QUESTIONS]
        assert np.mean(hp_s) > np.mean(tv_s)

    def test_exported_from_package(self):
        from llm_guard import QuestionComplexityScorer as QCS, ComplexityResult as CR
        assert QCS is not None
        assert CR is not None


# ── DebateVerifier ────────────────────────────────────────────────────────────

from llm_guard.healing.debate_verifier import DebateVerifier, DebateResult

class TestDebateVerifier:

    def test_single_agent_no_debate(self):
        dv = DebateVerifier()
        r = dv.verify("Q?", ["answer A"])
        assert r.verdict == 0
        assert not r.debate_triggered
        assert r.routing == "trusted"

    def test_agreeing_agents_no_debate(self):
        dv = DebateVerifier(disagree_threshold=0.40)
        # Both agents give the same answer
        r = dv.verify("Who invented the telephone?", ["Alexander Graham Bell", "Alexander Graham Bell"])
        assert not r.debate_triggered
        assert r.routing == "trusted"
        assert r.confidence == pytest.approx(0.95)

    def test_disagreeing_agents_triggers_debate(self):
        """Mock API call to avoid hitting real API in tests."""
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text='{"verdict": 0, "confidence": 0.85, "reason": "A has evidence"}')]

        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.return_value = mock_response
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None)
            r = dv.verify(
                "Who invented the telephone?",
                ["Alexander Graham Bell", "Elisha Gray", "Antonio Meucci"],
            )

        assert r.debate_triggered
        assert r.verdict == 0
        assert r.confidence == pytest.approx(0.85)
        assert r.risk == pytest.approx(0.15)
        assert r.routing == "trusted"
        assert not r.override_agent0

    def test_override_routing(self):
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text='{"verdict": 1, "confidence": 0.80, "reason": "B is correct"}')]

        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.return_value = mock_response
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None)
            r = dv.verify(
                "Capital of France?",
                ["London", "Paris", "Berlin"],
            )

        assert r.debate_triggered
        assert r.verdict == 1
        assert r.override_agent0
        assert r.routing == "override_agent0"

    def test_uncertain_routing(self):
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text='{"verdict": 0, "confidence": 0.40, "reason": "unclear"}')]

        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.return_value = mock_response
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None, uncertain_threshold=0.50)
            r = dv.verify("Q?", ["A", "B completely different answer here"])

        assert r.debate_triggered
        assert r.routing == "uncertain"

    def test_tie_verdict(self):
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text='{"verdict": "tie", "confidence": 0.3, "reason": "cannot decide"}')]

        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.return_value = mock_response
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None)
            r = dv.verify("Q?", ["Answer X", "totally different response Y"])

        assert r.verdict == -1
        assert r.debate_triggered

    def test_parse_response_valid_json(self):
        idx, conf, reason = DebateVerifier._parse_response(
            '{"verdict": 2, "confidence": 0.75, "reason": "best evidence"}', n_agents=3
        )
        assert idx == 2
        assert conf == pytest.approx(0.75)
        assert reason == "best evidence"

    def test_parse_response_tie(self):
        idx, conf, _ = DebateVerifier._parse_response(
            '{"verdict": "tie", "confidence": 0.4, "reason": ""}', n_agents=2
        )
        assert idx == -1

    def test_parse_response_out_of_range(self):
        # verdict=5 with only 3 agents → should return -1
        idx, conf, _ = DebateVerifier._parse_response(
            '{"verdict": 5, "confidence": 0.8, "reason": ""}', n_agents=3
        )
        assert idx == -1

    def test_parse_response_markdown_fenced(self):
        text = '```json\n{"verdict": 1, "confidence": 0.9, "reason": "test"}\n```'
        idx, conf, reason = DebateVerifier._parse_response(text, n_agents=2)
        assert idx == 1
        assert conf == pytest.approx(0.9)

    def test_api_failure_returns_safe_default(self):
        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.side_effect = Exception("network error")
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None)
            r = dv.verify("Q?", ["A", "B completely different"])

        # Soft fail — should not raise
        assert r.debate_triggered
        assert r.confidence == pytest.approx(0.5)

    def test_exported_from_package(self):
        from llm_guard import DebateVerifier as DV, DebateResult as DR
        assert DV is not None
        assert DR is not None

    def test_verify_batch(self):
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text='{"verdict": 0, "confidence": 0.9, "reason": "ok"}')]
        with patch("anthropic.Anthropic") as MockAnthropic:
            MockAnthropic.return_value.messages.create.return_value = mock_response
            dv = DebateVerifier(api_key="sk-fake", cache_dir=None)
            items = [
                {"question": "Q1?", "answers": ["A1", "B1 different"]},
                {"question": "Q2?", "answers": ["A2", "B2 different"]},
            ]
            results = dv.verify_batch(items)
        assert len(results) == 2
        assert all(isinstance(r, DebateResult) for r in results)


# ── GuardPhaseController ──────────────────────────────────────────────────────

from llm_guard.scoring.agent_guard import GuardPhaseController, PhaseStatus

class TestGuardPhaseController:

    def test_bootstrap_phase_initially(self):
        ctrl = GuardPhaseController(min_wrong=10, min_total=50)
        assert ctrl.phase == "bootstrap"
        assert ctrl.should_use_judge

    def test_transitions_to_calibrating(self):
        ctrl = GuardPhaseController(min_wrong=10, min_total=50)
        # 5 wrong + 15 correct = 20 total, 5 wrong  → calibrating
        for _ in range(15):
            ctrl.observe(0)
        for _ in range(5):
            ctrl.observe(1)
        assert ctrl.phase == "calibrating"
        assert ctrl.should_use_judge

    def test_transitions_to_self_sustaining(self):
        ctrl = GuardPhaseController(min_wrong=10, min_total=50)
        for _ in range(40):
            ctrl.observe(0)
        for _ in range(10):
            ctrl.observe(1)
        assert ctrl.phase == "self_sustaining"
        assert not ctrl.should_use_judge

    def test_is_starved(self):
        ctrl = GuardPhaseController()
        # 97 correct, 3 wrong → 3% < 5% default → starved
        for _ in range(97):
            ctrl.observe(0)
        for _ in range(3):
            ctrl.observe(1)
        assert ctrl.is_starved(starvation_ratio=0.05)   # 3% < 5% → True
        assert not ctrl.is_starved(starvation_ratio=0.02)  # 3% > 2% → not starved

    def test_status_dataclass(self):
        ctrl = GuardPhaseController()
        s = ctrl.status
        assert isinstance(s, PhaseStatus)
        assert s.phase in ("bootstrap", "calibrating", "self_sustaining")
        assert s.n_total == 0
        assert isinstance(s.should_use_judge, bool)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
