# tests/test_guard_state.py
import json, os, tempfile
import numpy as np
import pytest
from llm_guard.monitoring.guard_state import dump_state, load_state

def _make_fake_state():
    """Minimal state dict matching what AgentGuard.save_state() emits."""
    try:
        from sklearn.isotonic import IsotonicRegression
        ir = IsotonicRegression(out_of_bounds="clip")
        ir.fit([0.1, 0.3, 0.6, 0.9], [0, 0, 1, 1])
        iso_obj = ir
    except ImportError:
        iso_obj = None

    try:
        from llm_guard.scoring.agent_guard import PtrueWeightBandit
        bandit = PtrueWeightBandit()
        bandit.update(0.75, 1.0)
        bandit.update(0.90, 0.5)
        bandit_dict = {
            "counts": bandit._counts.tolist(),
            "rewards": bandit._rewards.tolist(),
            "total": bandit._total,
        }
    except Exception:
        bandit_dict = None

    return {
        "alert_threshold": 0.68,
        "iso_buffer": [(0.2, 0), (0.8, 1), (0.55, 1)],
        "iso_calibrator": iso_obj,
        "ptrue_bandit": bandit_dict,
        "version": "0.28.0",
    }


def test_roundtrip(tmp_path):
    state = _make_fake_state()
    path = str(tmp_path / "guard.llmg")
    dump_state(state, path)
    assert os.path.exists(path)
    loaded = load_state(path)
    assert loaded["alert_threshold"] == pytest.approx(0.68)
    assert loaded["iso_buffer"] == state["iso_buffer"]
    assert loaded["version"] == "0.28.0"

def test_iso_calibrator_roundtrip(tmp_path):
    state = _make_fake_state()
    if state["iso_calibrator"] is None:
        pytest.skip("sklearn not available")
    path = str(tmp_path / "guard.llmg")
    dump_state(state, path)
    loaded = load_state(path)
    # Loaded iso should make same predictions
    orig_pred = state["iso_calibrator"].predict([0.4, 0.7])
    loaded_pred = loaded["iso_calibrator"].predict([0.4, 0.7])
    np.testing.assert_allclose(orig_pred, loaded_pred, atol=1e-6)

def test_bandit_roundtrip(tmp_path):
    state = _make_fake_state()
    if state["ptrue_bandit"] is None:
        pytest.skip("PtrueWeightBandit not available")
    path = str(tmp_path / "guard.llmg")
    dump_state(state, path)
    loaded = load_state(path)
    assert loaded["ptrue_bandit"]["total"] == state["ptrue_bandit"]["total"]

def test_missing_file_returns_none():
    result = load_state("/nonexistent/path/guard.llmg")
    assert result is None

def test_corrupt_file_returns_none(tmp_path):
    path = str(tmp_path / "corrupt.llmg")
    with open(path, "w") as f:
        f.write("not valid json at all {{{{")
    result = load_state(path)
    assert result is None


import numpy as np
from llm_guard.scoring.agent_guard import AgentGuard


class TestAgentGuardPersistence:
    def test_save_and_load_threshold(self, tmp_path):
        guard = AgentGuard(alert_threshold=0.72)
        path = str(tmp_path / "state.llmg")
        guard.save_state(path)

        guard2 = AgentGuard()  # default threshold 0.70
        assert guard2._alert_threshold == pytest.approx(0.70)
        guard2.load_state(path)
        assert guard2._alert_threshold == pytest.approx(0.72)

    def test_save_and_load_iso_buffer(self, tmp_path):
        guard = AgentGuard()
        guard._iso_buffer = [(0.2, 0), (0.7, 1), (0.85, 1)]
        path = str(tmp_path / "state.llmg")
        guard.save_state(path)

        guard2 = AgentGuard()
        guard2.load_state(path)
        assert guard2._iso_buffer == [(0.2, 0), (0.7, 1), (0.85, 1)]

    def test_load_state_missing_file_is_noop(self, tmp_path):
        guard = AgentGuard(alert_threshold=0.65)
        guard.load_state(str(tmp_path / "nonexistent.llmg"))
        assert guard._alert_threshold == pytest.approx(0.65)  # unchanged

    def test_save_load_bandit_stats(self, tmp_path):
        guard = AgentGuard()
        guard.enable_ptrue_bandit()
        guard._ptrue_bandit.update(0.75, 1.0)
        guard._ptrue_bandit.update(0.90, 0.5)
        path = str(tmp_path / "state.llmg")
        guard.save_state(path)

        guard2 = AgentGuard()
        guard2.enable_ptrue_bandit()
        guard2.load_state(path)
        # After loading, the bandit should know it pulled 2 arms
        assert guard2._ptrue_bandit._total == 2

    def test_roundtrip_iso_calibrator(self, tmp_path):
        try:
            from sklearn.isotonic import IsotonicRegression
        except ImportError:
            pytest.skip("sklearn not available")
        guard = AgentGuard()
        guard.calibrate_isotonic([0.1, 0.3, 0.6, 0.9], [0, 0, 1, 1])
        path = str(tmp_path / "state.llmg")
        guard.save_state(path)

        guard2 = AgentGuard()
        guard2.load_state(path)
        assert guard2._iso_calibrator is not None
        pred_orig = guard._iso_calibrator.predict([0.5])
        pred_loaded = guard2._iso_calibrator.predict([0.5])
        np.testing.assert_allclose(pred_orig, pred_loaded, atol=1e-6)
