"""
tests/test_domain_router.py
Tests for DomainRouter (v0.64.0)
"""
import pytest
from unittest.mock import patch, MagicMock

from llm_guard import (
    QuestionType,
    DomainRoutedResult,
    QuestionTypeClassifier,
    DomainRouter,
)


# ── QuestionTypeClassifier (regex) ────────────────────────────────────────────

class TestQuestionTypeClassifier:
    def setup_method(self):
        self.clf = QuestionTypeClassifier(use_llm_fallback=False)

    def test_fact_verification_true_false(self):
        qtype, src = self.clf.classify("Is it true that the Earth is flat?")
        assert qtype == QuestionType.FACT_VERIFICATION
        assert src == "regex"

    def test_fact_verification_claim(self):
        qtype, _ = self.clf.classify(
            "The claim that Einstein won the Nobel Prize in Physics is correct."
        )
        assert qtype == QuestionType.FACT_VERIFICATION

    def test_multi_hop_entity_both(self):
        qtype, _ = self.clf.classify(
            "Were both Skin Yard and Ostava from the US?"
        )
        assert qtype == QuestionType.MULTI_HOP_ENTITY

    def test_multi_hop_entity_share(self):
        qtype, _ = self.clf.classify(
            "What occupation do Chris Menges and Aram Avakian share?"
        )
        assert qtype == QuestionType.MULTI_HOP_ENTITY

    def test_trivia_recall_in_what_year(self):
        qtype, _ = self.clf.classify(
            "In what year was the king who made the 1925 Birthday Honours born?"
        )
        assert qtype == QuestionType.TRIVIA_RECALL

    def test_open_factual_default(self):
        qtype, _ = self.clf.classify("What is the capital of France?")
        assert qtype == QuestionType.OPEN_FACTUAL

    def test_returns_tuple(self):
        result = self.clf.classify("Where is the Eiffel Tower?")
        assert len(result) == 2
        assert isinstance(result[0], QuestionType)
        assert result[1] in ("regex", "haiku")


# ── DomainRouter routing table ────────────────────────────────────────────────

class TestDomainRouterTable:
    def setup_method(self):
        self.router = DomainRouter(api_key="fake-key")

    def test_routing_table_returns_all_types(self):
        table = self.router.routing_table()
        for qt in QuestionType:
            assert qt.value in table

    def test_routing_table_fields(self):
        table = self.router.routing_table()
        for entry in table.values():
            assert "signal" in entry
            assert "auroc" in entry
            assert "ci" in entry
            assert "confidence" in entry

    def test_signal_assignments(self):
        table = self.router.routing_table()
        assert table["fact_verification"]["signal"] == "faj_r"
        assert table["multi_hop_entity"]["signal"] == "mfe_sem_sim"
        assert table["trivia_recall"]["signal"] == "faj_a"
        assert table["open_factual"]["signal"] == "mfe_sc_pair"
        assert table["multi_hop_bridge"]["signal"] == "sc_old"

    def test_validated_aurocs(self):
        table = self.router.routing_table()
        assert table["fact_verification"]["auroc"] == pytest.approx(0.810, abs=0.001)
        assert table["multi_hop_entity"]["auroc"] == pytest.approx(0.788, abs=0.001)
        assert table["trivia_recall"]["auroc"] == pytest.approx(0.751, abs=0.001)
        assert table["open_factual"]["auroc"] == pytest.approx(0.667, abs=0.001)
        assert table["multi_hop_bridge"]["auroc"] == pytest.approx(0.625, abs=0.001)

    def test_confidence_tiers(self):
        table = self.router.routing_table()
        assert table["fact_verification"]["confidence"] == "HIGH"
        assert table["multi_hop_entity"]["confidence"] == "HIGH"
        assert table["trivia_recall"]["confidence"] == "HIGH"
        assert table["open_factual"]["confidence"] == "LOW"
        assert table["multi_hop_bridge"]["confidence"] == "LOW"


# ── DomainRouter.score() with mocked signals ─────────────────────────────────

FAKE_STEPS = [
    {"action": "Search", "observation": "The Eiffel Tower is in Paris, France."}
]


class TestDomainRouterScore:
    def setup_method(self):
        self.router = DomainRouter(api_key="fake-key")

    def test_score_returns_domain_routed_result(self):
        with patch("llm_guard.domain_router._compute_mfe_sem_sim", return_value=0.3):
            result = self.router.score(
                question="Which entity did both X and Y share?",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.MULTI_HOP_ENTITY,
            )
        assert isinstance(result, DomainRoutedResult)

    def test_score_mfe_sem_sim_zero_cost(self):
        with patch("llm_guard.domain_router._compute_mfe_sem_sim", return_value=0.25):
            result = self.router.score(
                question="test",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.MULTI_HOP_ENTITY,
            )
        assert result.signal_used == "mfe_sem_sim"
        assert result.risk_score == pytest.approx(0.25, abs=0.001)
        assert result.domain_type == QuestionType.MULTI_HOP_ENTITY
        assert result.confidence == "HIGH"
        assert result.classifier_source == "override"

    def test_score_sc_old_uses_bridge(self):
        with patch("llm_guard.domain_router._compute_sc_old", return_value=(0.6, {})):
            result = self.router.score(
                question="test",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.MULTI_HOP_BRIDGE,
            )
        assert result.signal_used == "sc_old"
        assert result.risk_score == pytest.approx(0.6, abs=0.001)

    def test_score_faj_r_fact_verification(self):
        with patch("llm_guard.domain_router._compute_faj",
                   return_value=(0.2, {"judge_score": 4})):
            result = self.router.score(
                question="test",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.FACT_VERIFICATION,
            )
        assert result.signal_used == "faj_r"
        assert result.extras.get("judge_score") == 4

    def test_score_faj_a_trivia(self):
        with patch("llm_guard.domain_router._compute_faj",
                   return_value=(0.4, {"judge_score": 3})):
            result = self.router.score(
                question="test",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.TRIVIA_RECALL,
            )
        assert result.signal_used == "faj_a"

    def test_score_mfe_sc_pair_open_factual(self):
        with patch("llm_guard.domain_router._compute_mfe_sc_pair",
                   return_value=(0.7, {"direct_answers": ["Paris", "London"]})):
            result = self.router.score(
                question="test",
                steps=FAKE_STEPS,
                final_answer="Paris",
                domain_type=QuestionType.OPEN_FACTUAL,
            )
        assert result.signal_used == "mfe_sc_pair"
        assert result.extras.get("direct_answers") == ["Paris", "London"]

    def test_score_domain_type_override(self):
        with patch("llm_guard.domain_router._compute_mfe_sem_sim", return_value=0.1):
            result = self.router.score(
                question="Is it true that X?",  # regex → FACT_VERIFICATION
                steps=FAKE_STEPS,
                final_answer="yes",
                domain_type=QuestionType.MULTI_HOP_ENTITY,  # override
            )
        assert result.domain_type == QuestionType.MULTI_HOP_ENTITY
        assert result.signal_used == "mfe_sem_sim"
        assert result.classifier_source == "override"

    def test_risk_score_bounds(self):
        with patch("llm_guard.domain_router._compute_mfe_sem_sim", return_value=0.0):
            r1 = self.router.score("q", FAKE_STEPS, "a",
                                   domain_type=QuestionType.MULTI_HOP_ENTITY)
        assert 0.0 <= r1.risk_score <= 1.0

        with patch("llm_guard.domain_router._compute_mfe_sem_sim", return_value=1.0):
            r2 = self.router.score("q", FAKE_STEPS, "a",
                                   domain_type=QuestionType.MULTI_HOP_ENTITY)
        assert 0.0 <= r2.risk_score <= 1.0


# ── DomainRoutedResult dataclass ──────────────────────────────────────────────

class TestDomainRoutedResult:
    def test_fields_accessible(self):
        r = DomainRoutedResult(
            risk_score=0.42,
            domain_type=QuestionType.OPEN_FACTUAL,
            signal_used="mfe_sc_pair",
            signal_auroc=0.667,
            confidence="LOW",
            raw_signal_value=0.42,
            classifier_source="regex",
        )
        assert r.risk_score == pytest.approx(0.42)
        assert r.domain_type == QuestionType.OPEN_FACTUAL
        assert r.signal_used == "mfe_sc_pair"
        assert r.signal_auroc == pytest.approx(0.667)
        assert r.confidence == "LOW"
        assert r.extras == {}

    def test_extras_populated(self):
        r = DomainRoutedResult(
            risk_score=0.2,
            domain_type=QuestionType.FACT_VERIFICATION,
            signal_used="faj_r",
            signal_auroc=0.810,
            confidence="HIGH",
            raw_signal_value=0.2,
            classifier_source="haiku",
            extras={"judge_score": 4},
        )
        assert r.extras["judge_score"] == 4


# ── Integration smoke test (no API calls) ────────────────────────────────────

class TestDomainRouterIntegration:
    def test_import_from_llm_guard(self):
        from llm_guard import DomainRouter, QuestionType, DomainRoutedResult
        assert DomainRouter is not None
        assert QuestionType.FACT_VERIFICATION.value == "fact_verification"

    def test_classifier_regex_only_no_api(self):
        clf = QuestionTypeClassifier(use_llm_fallback=False)
        qtype, src = clf.classify("What is the boiling point of water?")
        assert isinstance(qtype, QuestionType)
        assert src == "regex"

    def test_router_instantiation_no_key(self):
        # Should not raise even without API key
        router = DomainRouter(api_key="")
        assert router is not None

    def test_routing_table_mean_auroc(self):
        router = DomainRouter()
        table = router.routing_table()
        aurocs = [v["auroc"] for v in table.values()]
        mean = sum(aurocs) / len(aurocs)
        assert mean == pytest.approx(0.728, abs=0.005)


# ── AUROC / AUPRC validation against saved signal matrices ───────────────────
# Uses the actual val-split signal matrices from exp_E1 — no mocking.

import json
import sys
import numpy as np
from pathlib import Path

_ROOT = Path(__file__).resolve().parent.parent
_MATRIX_DIR = _ROOT / "results" / "exp_E1_cold_start_sim"

# Oracle signal per domain (from E2 routing table)
_ORACLE_SIGNAL = {
    "fever":         "faj_r",
    "2wikimultihop": "mfe_sem_sim",
    "triviaqa":      "faj_a",
    "nq":            "mfe_sc_pair",
    "hotpotqa":      "sc_old",
}

# Validated AUROCs from E2 (tolerance ±0.02 for floating-point variance)
_EXPECTED_AUROC = {
    "fever":         0.810,
    "2wikimultihop": 0.788,
    "triviaqa":      0.751,
    "nq":            0.667,
    "hotpotqa":      0.625,
}

# AUPRC lower-bound: AUPRC > class prior (fraction of positives)
_CLASS_PRIOR = {
    "fever":         20 / 38,
    "2wikimultihop": 22 / 40,
    "triviaqa":      20 / 140,
    "nq":            23 / 40,
    "hotpotqa":      24 / 43,
}


def _load_matrix_labels(domain: str):
    """Load (scores, labels) from saved E1 signal matrix + chain labels."""
    mat_path = _MATRIX_DIR / f"signal_matrix_{domain}_val.json"
    if not mat_path.exists():
        pytest.skip(f"Signal matrix not found: {mat_path}")
    matrix = json.loads(mat_path.read_text())

    sys.path.insert(0, str(_ROOT))
    from experiments.shared.data_loader import load_split
    chains = load_split(domain, "val")

    signal = _ORACLE_SIGNAL[domain]
    scores, labels = [], []
    for chain in chains:
        row = matrix.get(chain.chain_id, {})
        val = row.get(signal, 0.5)
        scores.append(float(val))
        labels.append(1 - chain.label)  # 1=incorrect

    return np.array(scores), np.array(labels)


@pytest.mark.parametrize("domain", list(_ORACLE_SIGNAL.keys()))
class TestOracleSignalMetrics:

    def test_auroc_matches_validated(self, domain):
        from sklearn.metrics import roc_auc_score
        scores, labels = _load_matrix_labels(domain)
        auc = roc_auc_score(labels, scores)
        auc = max(auc, 1.0 - auc)  # handle inverted
        expected = _EXPECTED_AUROC[domain]
        assert auc == pytest.approx(expected, abs=0.02), (
            f"{domain}: AUROC {auc:.3f} deviates from validated {expected:.3f}"
        )

    def test_auroc_above_random(self, domain):
        from sklearn.metrics import roc_auc_score
        scores, labels = _load_matrix_labels(domain)
        auc = roc_auc_score(labels, scores)
        auc = max(auc, 1.0 - auc)
        assert auc >= 0.60, f"{domain}: AUROC {auc:.3f} < 0.60 (below useful threshold)"

    def test_auprc_above_class_prior(self, domain):
        from sklearn.metrics import average_precision_score
        scores, labels = _load_matrix_labels(domain)
        # Use risk-direction scores (positive class = incorrect)
        from sklearn.metrics import roc_auc_score
        auc = roc_auc_score(labels, scores)
        if auc < 0.5:
            scores = 1.0 - scores  # flip to correct direction
        auprc = average_precision_score(labels, scores)
        prior = _CLASS_PRIOR[domain]
        assert auprc > prior, (
            f"{domain}: AUPRC {auprc:.3f} not above class prior {prior:.3f}"
        )

    def test_no_constant_scores(self, domain):
        """Signal must have variance — constant = no information."""
        scores, _ = _load_matrix_labels(domain)
        assert np.std(scores) > 1e-4, f"{domain}: signal has near-zero variance"

    def test_scores_in_unit_interval(self, domain):
        scores, _ = _load_matrix_labels(domain)
        assert scores.min() >= 0.0 and scores.max() <= 1.0, (
            f"{domain}: scores outside [0, 1]: min={scores.min():.3f} max={scores.max():.3f}"
        )
