"""
Tests for v0.42.0 features:
  - ChainTrustResult CI fields (ci_lo, ci_hi, ci_width)
  - score_chain(return_ci=True)
  - AgentGuard.retrieval_quality()
  - AgentGuard.autotune_threshold()
  - AgentGuard.recommend_step_budget()
  - AgentGuard.recalibrate_for_model()
  - QppgMonitor.generate_report()
"""
import time
import pytest
import numpy as np


# ── Fixtures ──────────────────────────────────────────────────────────────────

SAMPLE_STEPS = [
    {
        "thought": "I need to find information about the Battle of Hastings.",
        "action_type": "Search",
        "action_arg": "Battle of Hastings date",
        "observation": "The Battle of Hastings was fought on 14 October 1066 between "
                       "King Harold II of England and the Norman-French army of William.",
    },
    {
        "thought": "I have enough information to answer.",
        "action_type": "Finish",
        "action_arg": "1066",
        "observation": "",
    },
]

SAMPLE_QUESTION = "When was the Battle of Hastings fought?"
SAMPLE_ANSWER   = "1066"


@pytest.fixture
def guard():
    from llm_guard import AgentGuard
    return AgentGuard()


@pytest.fixture
def monitor():
    from qppg_service import QppgMonitor
    return QppgMonitor(threshold=0.65)


# ── ChainTrustResult CI fields ─────────────────────────────────────────────────

class TestChainTrustResultCI:
    def test_ci_fields_zero_when_return_ci_false(self, guard):
        # v0.45.0: return_ci=True is the new default; pass False explicitly to test disabled path
        result = guard.score_chain(SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=False)
        assert result.ci_lo    == 0.0
        assert result.ci_hi    == 0.0
        assert result.ci_width == 0.0

    def test_ci_fields_populated_when_return_ci(self, guard):
        result = guard.score_chain(
            SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=True
        )
        assert 0.0 <= result.ci_lo    <= result.risk_score
        assert result.risk_score      <= result.ci_hi <= 1.0
        assert result.ci_width        == round(result.ci_hi - result.ci_lo, 4)
        assert result.ci_width        > 0.0

    def test_ci_width_reasonable(self, guard):
        """CI should be less than 0.6 wide (Wilson with n=30 on [0.2, 0.8] range < 0.36)."""
        result = guard.score_chain(
            SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=True
        )
        assert result.ci_width < 0.6

    def test_ci_bounds_valid_range(self, guard):
        result = guard.score_chain(
            SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=True
        )
        assert result.ci_lo >= 0.0
        assert result.ci_hi <= 1.0
        assert result.ci_lo <= result.ci_hi

    def test_ci_narrows_with_larger_cal_pool(self, guard):
        """Simulate larger calibration pool; CI should narrow."""
        # Patch _n_cal to simulate fitted state
        guard._scorer._n_cal = 200
        result_large = guard.score_chain(
            SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=True
        )
        guard._scorer._n_cal = 0
        result_small = guard.score_chain(
            SAMPLE_QUESTION, SAMPLE_STEPS, SAMPLE_ANSWER, return_ci=True
        )
        # Larger calibration pool → narrower CI
        assert result_large.ci_width <= result_small.ci_width + 0.01


# ── AgentGuard.retrieval_quality() ────────────────────────────────────────────

class TestRetrievalQuality:
    def test_returns_expected_keys(self, guard):
        result = guard.retrieval_quality(SAMPLE_QUESTION, SAMPLE_STEPS)
        assert "mean_sim"      in result
        assert "max_sim"       in result
        assert "min_sim"       in result
        assert "var_sim"       in result
        assert "n_relevant"    in result
        assert "per_step"      in result
        assert "quality_label" in result

    def test_quality_label_valid(self, guard):
        result = guard.retrieval_quality(SAMPLE_QUESTION, SAMPLE_STEPS)
        assert result["quality_label"] in {"GOOD", "FAIR", "POOR"}

    def test_per_step_length_matches_search_steps(self, guard):
        result = guard.retrieval_quality(SAMPLE_QUESTION, SAMPLE_STEPS)
        # SAMPLE_STEPS has 1 Search step
        assert len(result["per_step"]) == 1

    def test_mean_sim_in_range(self, guard):
        result = guard.retrieval_quality(SAMPLE_QUESTION, SAMPLE_STEPS)
        assert 0.0 <= result["mean_sim"] <= 1.0

    def test_empty_steps_returns_poor(self, guard):
        result = guard.retrieval_quality(SAMPLE_QUESTION, [])
        assert result["quality_label"] == "POOR"
        assert result["mean_sim"]       == 0.0
        assert result["per_step"]       == []

    def test_no_observation_returns_poor(self, guard):
        steps = [{"thought": "x", "action_type": "Search",
                  "action_arg": "foo", "observation": ""}]
        result = guard.retrieval_quality(SAMPLE_QUESTION, steps)
        assert result["quality_label"] == "POOR"

    def test_high_relevance_on_matching_observation(self, guard):
        steps = [
            {
                "thought": "Searching for battle date.",
                "action_type": "Search",
                "action_arg": "Battle of Hastings",
                "observation": "The Battle of Hastings was fought in 1066 in England.",
            }
        ]
        result = guard.retrieval_quality("When was the Battle of Hastings?", steps)
        # Should score at least FAIR (mean_sim >= 0.35)
        assert result["mean_sim"] >= 0.30


# ── AgentGuard.autotune_threshold() ───────────────────────────────────────────

class TestAutotuneThreshold:
    def test_returns_none_when_insufficient_labels(self, guard):
        # Empty buffer → None
        result = guard.autotune_threshold(target_precision=0.85, min_n=30)
        assert result is None

    def test_returns_none_when_below_min_n(self, guard):
        guard._iso_buffer = [(0.8, 1)] * 10  # only 10 samples
        result = guard.autotune_threshold(target_precision=0.85, min_n=30)
        assert result is None

    def test_returns_threshold_with_sufficient_labels(self, guard):
        """With 50 pairs: high scores = wrong, low scores = correct.
        At threshold 0.50, precision should be ~0.5 (50% wrong above 0.5).
        At threshold 0.80, precision should be ~1.0 (all above 0.80 are wrong).
        """
        rng = np.random.default_rng(42)
        # 25 correct chains (risk 0.2–0.5) + 25 wrong chains (risk 0.7–0.95)
        correct = [(float(s), 0) for s in rng.uniform(0.2, 0.5, 25)]
        wrong   = [(float(s), 1) for s in rng.uniform(0.7, 0.95, 25)]
        guard._iso_buffer = correct + wrong

        threshold = guard.autotune_threshold(target_precision=0.85, min_n=30)
        assert threshold is not None
        assert 0.50 <= threshold <= 0.95

        # Verify the returned threshold actually achieves target_precision
        above = [(s, l) for s, l in guard._iso_buffer if s >= threshold]
        if above:
            precision = sum(l for _, l in above) / len(above)
            assert precision >= 0.85

    def test_returns_float(self, guard):
        guard._iso_buffer = [(0.9, 1)] * 50
        result = guard.autotune_threshold(target_precision=0.50, min_n=30)
        assert result is None or isinstance(result, float)


# ── AgentGuard.recommend_step_budget() ────────────────────────────────────────

class TestRecommendStepBudget:
    def test_returns_empirical_default_without_data(self, guard):
        # Empty chain buffer → default 3
        result = guard.recommend_step_budget()
        assert result == 3

    def test_returns_int(self, guard):
        result = guard.recommend_step_budget()
        assert isinstance(result, int)

    def test_returns_default_with_few_correct(self, guard):
        # Only 5 correct chains → below 10 threshold → default 3
        guard._iso_buffer       = [(0.3, 0)] * 5
        guard._iso_chain_buffer = [
            {"steps": [{"action_type": "Search"}] * 2} for _ in range(5)
        ]
        result = guard.recommend_step_budget()
        assert result == 3

    def test_returns_budget_from_step_distribution(self, guard):
        """With 20 correct chains averaging 3 search steps, budget should be ≥3."""
        n = 20
        guard._iso_buffer       = [(0.3, 0)] * n
        guard._iso_chain_buffer = [
            {"steps": [{"action_type": "Search"}] * 3} for _ in range(n)
        ]
        result = guard.recommend_step_budget()
        assert 3 <= result <= 8

    def test_budget_clamped_max(self, guard):
        """Even if chains have many steps, budget is capped at 8."""
        n = 20
        guard._iso_buffer       = [(0.3, 0)] * n
        guard._iso_chain_buffer = [
            {"steps": [{"action_type": "Search"}] * 15} for _ in range(n)
        ]
        result = guard.recommend_step_budget()
        assert result <= 8

    def test_budget_clamped_min(self, guard):
        """Budget is at least 3 even for very short chains."""
        n = 20
        guard._iso_buffer       = [(0.3, 0)] * n
        guard._iso_chain_buffer = [
            {"steps": [{"action_type": "Search"}] * 1} for _ in range(n)
        ]
        result = guard.recommend_step_budget()
        assert result >= 3


# ── AgentGuard.recalibrate_for_model() ────────────────────────────────────────

class TestRecalibrateForModel:
    def _make_seed_chains(self, n=5):
        chains = []
        for i in range(n):
            label = 1 if i % 2 == 0 else 0  # alternating correct/wrong
            chains.append({
                "question":     SAMPLE_QUESTION,
                "steps":        SAMPLE_STEPS,
                "final_answer": SAMPLE_ANSWER,
                "label":        label,
                "risk_score":   0.3 if label == 1 else 0.8,
            })
        return chains

    def test_returns_dict_with_expected_keys(self, guard):
        seed_chains = self._make_seed_chains(5)
        result = guard.recalibrate_for_model("gpt-4o", seed_chains)
        assert "model_id"                    in result
        assert "n_seed"                      in result
        assert "platt_T"                     in result
        assert "platt_b"                     in result
        assert "sample_raw_scores"           in result
        assert "sample_recalibrated_scores"  in result

    def test_model_id_preserved(self, guard):
        seed_chains = self._make_seed_chains(4)
        result = guard.recalibrate_for_model("llama-3.3-70b", seed_chains)
        assert result["model_id"] == "llama-3.3-70b"

    def test_n_seed_correct(self, guard):
        seed_chains = self._make_seed_chains(5)
        result = guard.recalibrate_for_model("gpt-4o", seed_chains)
        assert result["n_seed"] == 5

    def test_raises_with_fewer_than_2_chains(self, guard):
        with pytest.raises(ValueError, match="at least 2"):
            guard.recalibrate_for_model("gpt-4o", [self._make_seed_chains(5)[0]])

    def test_platt_t_positive(self, guard):
        seed_chains = self._make_seed_chains(6)
        result = guard.recalibrate_for_model("claude-haiku", seed_chains)
        assert result["platt_T"] > 0.0

    def test_recalibrated_scores_in_range(self, guard):
        seed_chains = self._make_seed_chains(6)
        result = guard.recalibrate_for_model("gpt-4o", seed_chains)
        for s in result["sample_recalibrated_scores"]:
            assert 0.0 <= s <= 1.0

    def test_seed_labels_override(self, guard):
        """seed_labels kwarg should override chain['label']."""
        seed_chains = self._make_seed_chains(4)
        labels = [1, 0, 1, 0]
        result = guard.recalibrate_for_model("gpt-4o", seed_chains, seed_labels=labels)
        assert result["n_seed"] == 4


# ── QppgMonitor.generate_report() ─────────────────────────────────────────────

class TestGenerateReport:
    def _seed_monitor(self, monitor, n=20):
        """Add n fake history entries over the last 10 days."""
        import random
        rng = random.Random(42)
        now = time.time()
        for i in range(n):
            ts = now - rng.uniform(0, 7 * 86400)
            monitor._history.append({
                "timestamp":       ts,
                "question":        f"Question {i}",
                "final_answer":    "answer",
                "step_count":      rng.randint(1, 5),
                "finished":        True,
                "risk_score":      rng.uniform(0.1, 0.9),
                "behavioral_score": rng.uniform(0.1, 0.9),
                "needs_review":    rng.random() > 0.7,
                "domain":          "test",
            })

    def test_returns_dict_with_expected_keys(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        for key in ["period_start", "period_end", "days", "total_chains",
                    "alert_count", "alert_rate", "avg_risk", "p95_risk",
                    "per_day", "trend", "calibration_state"]:
            assert key in report, f"Missing key: {key}"

    def test_empty_history_returns_zero(self, monitor):
        report = monitor.generate_report(days=7)
        assert report["total_chains"] == 0
        assert report["alert_count"]  == 0
        assert report["per_day"]      == []

    def test_days_parameter(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        assert report["days"] == 7

    def test_per_day_has_days_entries(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        assert len(report["per_day"]) == 7

    def test_trend_is_valid(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        assert report["trend"] in {"rising", "falling", "stable"}

    def test_calibration_state_behavioral_only(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        assert report["calibration_state"] == "behavioral_only"

    def test_avg_risk_in_range(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        if report["avg_risk"] is not None:
            assert 0.0 <= report["avg_risk"] <= 1.0

    def test_p95_risk_gte_avg_risk(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        if report["avg_risk"] is not None and report["p95_risk"] is not None:
            assert report["p95_risk"] >= report["avg_risk"] - 0.01

    def test_alert_rate_consistent(self, monitor):
        self._seed_monitor(monitor)
        report = monitor.generate_report(days=7)
        if report["total_chains"] > 0:
            expected_rate = report["alert_count"] / report["total_chains"]
            assert abs(report["alert_rate"] - expected_rate) < 0.001

    def test_rising_trend_detected(self, monitor):
        """Insert history where risk rises over time."""
        now = time.time()
        for i in range(20):
            ts = now - (19 - i) * 3600  # oldest first
            risk = 0.2 + i * 0.03        # risk rises from 0.2 to 0.77
            monitor._history.append({
                "timestamp": ts,
                "question": f"Q{i}", "final_answer": "A",
                "step_count": 2, "finished": True,
                "risk_score": risk, "behavioral_score": risk,
                "needs_review": risk > 0.65, "domain": "test",
            })
        report = monitor.generate_report(days=1)
        # With a strong rising trend, should detect "rising" or at least not fall
        assert report["trend"] in {"rising", "stable"}
