"""tests/test_healing_repair_bandit.py — Unit tests for RepairBandit."""
import pytest

pytestmark = pytest.mark.unit

STRATEGIES = ["ADDITIONAL_SEARCH", "RESOLVE_CONFLICT", "ANCHOR_TO_EVIDENCE"]


class TestImport:
    def test_import(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        assert RepairBandit is not None


class TestInstantiation:
    def test_basic_construction(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        assert b.strategies == STRATEGIES

    def test_custom_exploration(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES, exploration_c=2.0)
        assert b.exploration_c == 2.0

    def test_initial_counts_zero(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        for s in STRATEGIES:
            assert b._counts[s] == 0
            assert b._rewards[s] == 0.0


class TestSelect:
    def test_first_calls_round_robin(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        # Round-robin: must call update() between selects to advance past count=0 arm
        first_three = []
        for _ in range(3):
            s = b.select()
            first_three.append(s)
            b.update(s, 0.5)
        assert set(first_three) == set(STRATEGIES)

    def test_select_returns_valid_strategy(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        for _ in range(10):
            s = b.select()
            assert s in STRATEGIES

    def test_single_strategy(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=["ONLY_STRATEGY"])
        s = b.select()
        assert s == "ONLY_STRATEGY"


class TestUpdate:
    def test_update_increments_count(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        b.update(STRATEGIES[0], reward=1.0)
        assert b._counts[STRATEGIES[0]] == 1
        assert b._rewards[STRATEGIES[0]] == 1.0

    def test_update_accumulates(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES)
        b.update(STRATEGIES[0], reward=0.5)
        b.update(STRATEGIES[0], reward=0.8)
        assert b._counts[STRATEGIES[0]] == 2
        assert b._rewards[STRATEGIES[0]] == pytest.approx(1.3, abs=1e-6)

    def test_ucb_prefers_higher_reward(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=["A", "B"], exploration_c=0.0)
        # Prime both arms
        b.update("A", reward=0.9)
        b.update("B", reward=0.1)
        # With no exploration, should pick A
        selected = b.select()
        assert selected == "A"


class TestPersistence:
    def test_in_memory_no_file(self):
        from llm_guard.healing.repair_bandit import RepairBandit
        b = RepairBandit(strategies=STRATEGIES, db_path=None)
        b.update(STRATEGIES[0], reward=1.0)
        assert b._counts[STRATEGIES[0]] == 1

    def test_sqlite_persistence(self, tmp_path):
        from llm_guard.healing.repair_bandit import RepairBandit
        db = str(tmp_path / "bandit.db")
        b1 = RepairBandit(strategies=STRATEGIES, db_path=db)
        b1.update(STRATEGIES[0], reward=1.0)
        # Re-load from same db
        b2 = RepairBandit(strategies=STRATEGIES, db_path=db)
        assert b2._counts[STRATEGIES[0]] == 1
        assert b2._rewards[STRATEGIES[0]] == pytest.approx(1.0, abs=1e-6)
