"""tests/test_scoring_mini_judge.py — Unit tests for llm_guard.scoring.mini_judge."""
import pytest
from unittest.mock import MagicMock, patch
import numpy as np

pytestmark = pytest.mark.unit

SAMPLE_CHAIN = {
    "question": "Who founded Apple?",
    "steps": [
        {"thought": "I need to search", "action_type": "Search",
         "action_args": "Apple founder", "observation": "Steve Jobs founded Apple in 1976."},
    ],
    "final_answer": "Steve Jobs",
}

SAMPLE_STEPS = [
    {"thought": "I need to search", "action_type": "Search",
     "action_args": "Apple founder", "observation": "Steve Jobs founded Apple in 1976."},
]


class TestImport:
    def test_import(self):
        from llm_guard.scoring.mini_judge import MiniJudge
        assert MiniJudge is not None

    def test_import_extract_features(self):
        from llm_guard.scoring.mini_judge import _extract_features
        assert callable(_extract_features)


class TestInstantiation:
    def test_default_construction(self, tmp_path):
        """MiniJudge without a pkl file should instantiate without error."""
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "nonexistent.pkl"))
        assert judge is not None
        assert not judge.is_fitted

    def test_is_fitted_false_when_no_pkl(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "no.pkl"))
        assert judge.is_fitted is False

    def test_auroc_none_when_unfitted(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "no.pkl"))
        assert judge.auroc is None


class TestFeatureExtraction:
    def test_extract_features_shape(self):
        from llm_guard.scoring.mini_judge import _extract_features
        feats = _extract_features(SAMPLE_CHAIN)
        assert feats.shape == (14,)

    def test_extract_features_empty_steps(self):
        from llm_guard.scoring.mini_judge import _extract_features
        chain = {"question": "test?", "steps": [], "final_answer": "ans"}
        feats = _extract_features(chain)
        assert feats.shape == (14,)

    def test_extract_features_returns_floats(self):
        from llm_guard.scoring.mini_judge import _extract_features
        feats = _extract_features(SAMPLE_CHAIN)
        assert feats.dtype == np.float64

    def test_extract_features_range(self):
        """Most features should be in a reasonable range."""
        from llm_guard.scoring.mini_judge import _extract_features
        feats = _extract_features(SAMPLE_CHAIN)
        # First 11 features are behavioral signals 0–1
        assert all(np.isfinite(feats))


class TestFitAndScore:
    def _make_chains(self, n=20):
        chains = []
        for i in range(n):
            chains.append({
                "question": "Who is X?",
                "steps": [{"thought": "think", "action_type": "Search",
                            "action_args": "X", "observation": "X is a person."}],
                "final_answer": "X is a person.",
                "correct": (i % 2 == 0),
            })
        return chains

    def test_fit_creates_pipeline(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "model.pkl"))
        chains = self._make_chains(20)
        judge.fit(chains, n_splits=2)
        assert judge.is_fitted

    def test_score_after_fit_returns_float(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "model.pkl"))
        chains = self._make_chains(20)
        judge.fit(chains, n_splits=2)
        risk = judge.score("Who?", SAMPLE_STEPS, "Someone")
        assert 0.0 <= risk <= 1.0

    def test_score_chain_alias(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "model.pkl"))
        chains = self._make_chains(20)
        judge.fit(chains, n_splits=2)
        r1 = judge.score("Q", SAMPLE_STEPS, "A")
        r2 = judge.score_chain("Q", SAMPLE_STEPS, "A")
        assert r1 == r2

    def test_score_raises_when_unfitted(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "no.pkl"))
        with pytest.raises(RuntimeError, match="not fitted"):
            judge.score("Q", [], "A")

    def test_save_load_roundtrip(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        model_path = str(tmp_path / "judge.pkl")
        judge = MiniJudge(model_path=model_path)
        chains = self._make_chains(20)
        judge.fit(chains, n_splits=2)
        judge.save(model_path)
        loaded = MiniJudge.load(model_path)
        assert loaded.is_fitted

    def test_save_raises_unfitted(self, tmp_path):
        from llm_guard.scoring.mini_judge import MiniJudge
        judge = MiniJudge(model_path=str(tmp_path / "no.pkl"))
        with pytest.raises(RuntimeError):
            judge.save(str(tmp_path / "out.pkl"))
