"""Tests for v0.60.0 — SH1 RepairBandit, SH2 heal_with_retry, SH3 track_step CUSUM."""
import sys
sys.path.insert(0, "/Users/amajumder/Downloads/my research/QPPG")

from llm_guard import RepairBandit


# ── SH1: RepairBandit ─────────────────────────────────────────────────────────

def test_repair_bandit_select_returns_valid_strategy():
    bandit = RepairBandit(strategies=["A", "B", "C"])
    s = bandit.select()
    assert s in ["A", "B", "C"]


def test_repair_bandit_ucb1_learns():
    bandit = RepairBandit(strategies=["A", "B"])
    for _ in range(10):
        bandit.update("A", 1.0)
    for _ in range(10):
        bandit.update("B", 0.0)
    counts = {"A": 0, "B": 0}
    for _ in range(100):
        counts[bandit.select()] += 1
    assert counts["A"] > counts["B"]


def test_repair_bandit_sqlite_persistence(tmp_path):
    db = str(tmp_path / "test.db")
    bandit = RepairBandit(strategies=["A", "B"], db_path=db)
    bandit.update("A", 1.0)
    bandit2 = RepairBandit(strategies=["A", "B"], db_path=db)
    assert bandit2._counts["A"] == 1


# ── SH2: SelfHealer.heal_with_retry ──────────────────────────────────────────

def test_heal_with_retry_exhausts_attempts():
    from qppg_service.self_healer import SelfHealer
    healer = SelfHealer()
    calls = {"n": 0}

    def mock_fn(strategy, chain):
        calls["n"] += 1
        return {"success": False, "new_chain": chain}

    result = healer.heal_with_retry(
        chain={"steps": [], "final_answer": "x"},
        failure_mode="RETRIEVAL_FAILURE",
        execute_fn=mock_fn,
        max_attempts=3,
    )
    assert calls["n"] == 3
    assert result["success"] is False


def test_heal_with_retry_stops_on_success():
    from qppg_service.self_healer import SelfHealer
    healer = SelfHealer()
    calls = {"n": 0}

    def mock_fn(strategy, chain):
        calls["n"] += 1
        return {"success": calls["n"] >= 2, "new_chain": chain}

    result = healer.heal_with_retry(
        chain={"steps": [], "final_answer": "x"},
        failure_mode="RETRIEVAL_FAILURE",
        execute_fn=mock_fn,
        max_attempts=5,
    )
    assert calls["n"] == 2
    assert result["success"] is True


# ── SH3: QppgMonitor.track_step CUSUM ────────────────────────────────────────

def test_track_step_detects_drift():
    from qppg_service.monitor import QppgMonitor
    monitor = QppgMonitor()
    alerts = []
    for i, risk in enumerate([0.1, 0.1, 0.1, 0.9, 0.9]):
        r = monitor.track_step("chain_1", i, risk)
        if r and r.get("drift_detected"):
            alerts.append(i)
    assert len(alerts) > 0


def test_track_step_no_drift_stable():
    from qppg_service.monitor import QppgMonitor
    monitor = QppgMonitor()
    alerts = []
    for i, risk in enumerate([0.5, 0.51, 0.49, 0.50, 0.52]):
        r = monitor.track_step("chain_stable", i, risk)
        if r and r.get("drift_detected"):
            alerts.append(i)
    assert len(alerts) == 0
