"""
Tests for AdversarialChainDetector — v0.31.0

Validates:
  1. Feature extraction shape and range
  2. fit() + score() basic contract
  3. load_default() pre-trained model works
  4. score_chain() returns AdversarialResult with correct routing
  5. Adversarial chains score higher than correct chains (discrimination)
  6. Serialisation round-trip (export_params / from_params)
  7. Edge cases: empty steps, single step, unfitted raises
  8. Exports from package
"""
import sys
import numpy as np
import pytest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.adversarial.adversarial_detector import (
    AdversarialChainDetector,
    AdversarialParams,
    AdversarialResult,
    ADVERSARIAL_FEATURE_NAMES,
    extract_chain_features,
)


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

def _make_adversarial_steps():
    """Mimic a FARL confident_wrong chain: 1 search + finish, long observation."""
    return [
        {
            "thought": "I need to look up when the Titanic sank and who was the captain.",
            "action_type": "Search",
            "action_arg": "Titanic sank year captain",
            "observation": "The RMS Titanic sank on 15 April 1912 at 02:20. " * 10,
        },
        {
            "thought": "The search results confirm that the Titanic sank in 1912 and "
                       "the captain was Edward John Smith, who went down with the ship.",
            "action_type": "Finish",
            "action_arg": "1912, Captain Edward John Smith",
            "observation": "",
        },
    ]


def _make_correct_steps():
    """Mimic an HP-style correct chain: multiple searches, varied observations."""
    return [
        {
            "thought": "I need to find the height of 712 Fifth Avenue.",
            "action_type": "Search",
            "action_arg": "712 Fifth Avenue height",
            "observation": "712 Fifth Avenue is a skyscraper in Midtown Manhattan. Height: 755 ft.",
        },
        {
            "thought": "Now I need the height of Manhattan Life Insurance Building.",
            "action_type": "Search",
            "action_arg": "Manhattan Life Insurance Building height feet",
            "observation": "The Manhattan Life Insurance Building is 348 feet tall.",
        },
        {
            "thought": "712 Fifth Avenue (755 ft) is taller than Manhattan Life (348 ft).",
            "action_type": "Finish",
            "action_arg": "712 Fifth Avenue",
            "observation": "",
        },
    ]


ADVERSARIAL_Q = "What year did the Titanic sink and who was the captain?"
CORRECT_Q = "Which tower is taller, 712 Fifth Avenue or Manhattan Life Insurance Building?"


# ── Feature extraction ────────────────────────────────────────────────────────

class TestFeatureExtraction:

    def test_shape(self):
        feat = extract_chain_features(_make_correct_steps(), CORRECT_Q)
        assert feat.shape == (16,)
        assert len(ADVERSARIAL_FEATURE_NAMES) == 16

    def test_range(self):
        for steps, q in [
            (_make_adversarial_steps(), ADVERSARIAL_Q),
            (_make_correct_steps(), CORRECT_Q),
            ([], "empty question?"),
        ]:
            feat = extract_chain_features(steps, q)
            assert np.all(feat >= 0.0), f"negative feature: {feat}"
            assert np.all(feat <= 3.0), f"feature > 3: {feat}"

    def test_empty_steps_no_crash(self):
        feat = extract_chain_features([], "anything?")
        assert feat.shape == (16,)
        assert np.all(feat == 0.0)

    def test_adversarial_has_fewer_searches(self):
        """frac_search should be lower for adversarial (1 search / 2 steps = 0.5)."""
        fa = extract_chain_features(_make_adversarial_steps(), ADVERSARIAL_Q)
        fc = extract_chain_features(_make_correct_steps(), CORRECT_Q)
        # Adversarial: 1/2=0.5, Correct: 2/3=0.67
        frac_idx = ADVERSARIAL_FEATURE_NAMES.index("frac_search")
        assert fa[frac_idx] < fc[frac_idx]

    def test_adversarial_longer_obs(self):
        """avg_obs_len should be higher for adversarial chain."""
        fa = extract_chain_features(_make_adversarial_steps(), ADVERSARIAL_Q)
        fc = extract_chain_features(_make_correct_steps(), CORRECT_Q)
        obs_idx = ADVERSARIAL_FEATURE_NAMES.index("avg_obs_len")
        assert fa[obs_idx] > fc[obs_idx]


# ── Fit & score ───────────────────────────────────────────────────────────────

class TestFitAndScore:

    def _make_dataset(self, n=20):
        adversarial = [(_make_adversarial_steps(), ADVERSARIAL_Q)] * n
        correct     = [(_make_correct_steps(), CORRECT_Q)] * n
        return adversarial, correct

    def test_fit_returns_self(self):
        det = AdversarialChainDetector(experimental=True)
        adv, cor = self._make_dataset(10)
        result = det.fit(adv, cor)
        assert result is det

    def test_is_fitted_after_fit(self):
        det = AdversarialChainDetector(experimental=True)
        assert not det.is_fitted
        adv, cor = self._make_dataset(10)
        det.fit(adv, cor)
        assert det.is_fitted

    def test_score_range(self):
        det = AdversarialChainDetector(experimental=True)
        adv, cor = self._make_dataset(15)
        det.fit(adv, cor)
        s = det.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert 0.0 <= s <= 1.0

    def test_unfitted_raises(self):
        det = AdversarialChainDetector(experimental=True)
        with pytest.raises(RuntimeError):
            det.score(_make_correct_steps(), CORRECT_Q)

    def test_adversarial_scores_higher(self):
        """After fit on representative data, adversarial chains should score > correct."""
        det = AdversarialChainDetector(experimental=True)
        adv, cor = self._make_dataset(20)
        det.fit(adv, cor)
        s_adv = det.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        s_cor = det.score(_make_correct_steps(), CORRECT_Q)
        assert s_adv > s_cor

    def test_predict_bool(self):
        det = AdversarialChainDetector(threshold=0.50, experimental=True)
        adv, cor = self._make_dataset(20)
        det.fit(adv, cor)
        result = det.predict(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert isinstance(result, bool)

    def test_score_chain_returns_result(self):
        det = AdversarialChainDetector(experimental=True)
        adv, cor = self._make_dataset(15)
        det.fit(adv, cor)
        r = det.score_chain(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert isinstance(r, AdversarialResult)
        assert 0.0 <= r.adv_score <= 1.0
        assert r.routing in ("safe", "adversarial_suspected")
        assert isinstance(r.is_adversarial, bool)
        assert len(r.top_features) == 3

    def test_routing_adversarial(self):
        det = AdversarialChainDetector(threshold=0.01, experimental=True)  # very low threshold
        adv, cor = self._make_dataset(20)
        det.fit(adv, cor)
        r = det.score_chain(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert r.routing == "adversarial_suspected"
        assert r.is_adversarial is True

    def test_routing_safe(self):
        det = AdversarialChainDetector(threshold=0.99, experimental=True)  # very high threshold
        adv, cor = self._make_dataset(20)
        det.fit(adv, cor)
        r = det.score_chain(_make_correct_steps(), CORRECT_Q)
        assert r.routing == "safe"
        assert r.is_adversarial is False

    def test_score_batch(self):
        det = AdversarialChainDetector(experimental=True)
        adv, cor = self._make_dataset(15)
        det.fit(adv, cor)
        chains = [
            (_make_adversarial_steps(), ADVERSARIAL_Q),
            (_make_correct_steps(), CORRECT_Q),
        ]
        results = det.score_batch(chains)
        assert len(results) == 2
        assert all(isinstance(r, AdversarialResult) for r in results)


# ── load_default ──────────────────────────────────────────────────────────────

class TestLoadDefault:

    def test_load_default_is_fitted(self):
        det = AdversarialChainDetector.load_default(experimental=True)
        assert det.is_fitted

    def test_load_default_score_range(self):
        det = AdversarialChainDetector.load_default(experimental=True)
        s = det.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert 0.0 <= s <= 1.0

    def test_load_default_discrimination(self):
        """Pre-trained model should score adversarial > correct."""
        det = AdversarialChainDetector.load_default(experimental=True)
        s_adv = det.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        s_cor = det.score(_make_correct_steps(), CORRECT_Q)
        assert s_adv > s_cor, f"pre-trained: adv={s_adv:.3f}, cor={s_cor:.3f}"

    def test_load_default_auroc_field(self):
        det = AdversarialChainDetector.load_default(experimental=True)
        params = det.export_params()
        assert params.auroc == pytest.approx(0.9960)


# ── Serialisation ─────────────────────────────────────────────────────────────

class TestSerialisation:

    def _fitted_det(self, n=15):
        det = AdversarialChainDetector(experimental=True)
        adv = [(_make_adversarial_steps(), ADVERSARIAL_Q)] * n
        cor = [(_make_correct_steps(), CORRECT_Q)] * n
        det.fit(adv, cor)
        return det

    def test_export_params_type(self):
        params = self._fitted_det().export_params()
        assert isinstance(params, AdversarialParams)
        assert len(params.coef) == 16
        assert len(params.scale_mean) == 16
        assert len(params.scale_std) == 16

    def test_round_trip_dict(self):
        det = self._fitted_det()
        params = det.export_params()
        d = params.to_dict()
        restored = AdversarialParams.from_dict(d)
        det2 = AdversarialChainDetector.from_params(restored)
        s1 = det.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        s2 = det2.score(_make_adversarial_steps(), ADVERSARIAL_Q)
        assert s1 == pytest.approx(s2, abs=1e-6)

    def test_round_trip_json(self):
        det = self._fitted_det()
        params = det.export_params()
        j = params.to_json()
        restored = AdversarialParams.from_json(j)
        det2 = AdversarialChainDetector.from_params(restored)
        s1 = det.score(_make_correct_steps(), CORRECT_Q)
        s2 = det2.score(_make_correct_steps(), CORRECT_Q)
        assert s1 == pytest.approx(s2, abs=1e-6)

    def test_params_no_raw_data(self):
        params = self._fitted_det().export_params()
        d = params.to_dict()
        assert set(d.keys()) == {"coef", "intercept", "scale_mean", "scale_std",
                                  "threshold", "auroc"}


# ── Package export ────────────────────────────────────────────────────────────

class TestPackageExport:

    def test_exported_from_package(self):
        from llm_guard import AdversarialChainDetector as ACD
        from llm_guard import AdversarialResult as AR
        assert ACD is not None
        assert AR is not None


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