"""Tests for ActiveSamplingStrategy and GuardPhaseController integration."""
import sys
import pytest
import numpy as np
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.active_sampling import ActiveSamplingStrategy


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

def make_pool(n=10):
    return [
        {
            "question": f"What is the capital of country_{i}?",
            "steps": [{"thought": f"search {i}", "action_type": "Search",
                       "action_arg": f"capital country_{i}",
                       "observation": f"The capital is city_{i}."}],
            "final_answer": f"city_{i}",
        }
        for i in range(n)
    ]


# ── Basic interface ───────────────────────────────────────────────────────────

def test_select_returns_k_chains():
    strategy = ActiveSamplingStrategy(strategy="diversity", embedder="tfidf")
    pool = make_pool(10)
    selected = strategy.select(pool, k=3)
    assert len(selected) == 3


def test_select_returns_subset_of_pool():
    strategy = ActiveSamplingStrategy(strategy="diversity", embedder="tfidf")
    pool = make_pool(10)
    selected = strategy.select(pool, k=3)
    for chain in selected:
        assert chain in pool


def test_select_no_duplicates():
    strategy = ActiveSamplingStrategy(strategy="diversity", embedder="tfidf")
    pool = make_pool(10)
    selected = strategy.select(pool, k=5)
    indices = [pool.index(c) for c in selected]
    assert len(set(indices)) == len(indices)


def test_select_k_larger_than_pool_returns_all():
    strategy = ActiveSamplingStrategy(strategy="diversity", embedder="tfidf")
    pool = make_pool(3)
    selected = strategy.select(pool, k=10)
    assert len(selected) == 3


# ── Diversity strategy ────────────────────────────────────────────────────────

def test_diversity_selects_different_chains():
    strategy = ActiveSamplingStrategy(strategy="diversity", embedder="tfidf")
    pool = make_pool(10)
    selected = strategy.select(pool, k=3)
    assert len(selected) == 3


# ── Uncertainty strategy ──────────────────────────────────────────────────────

def test_uncertainty_requires_scores():
    strategy = ActiveSamplingStrategy(strategy="uncertainty", embedder="tfidf")
    pool = make_pool(5)
    with pytest.raises((ValueError, TypeError)):
        strategy.select(pool, k=2, scores=None)


def test_uncertainty_selects_boundary_chains():
    strategy = ActiveSamplingStrategy(strategy="uncertainty", embedder="tfidf")
    pool = make_pool(5)
    scores = [0.1, 0.45, 0.9, 0.55, 0.0]
    selected = strategy.select(pool, k=2, scores=scores)
    assert len(selected) == 2
    selected_questions = [c["question"] for c in selected]
    assert "What is the capital of country_1?" in selected_questions
    assert "What is the capital of country_3?" in selected_questions


# ── Hybrid strategy ───────────────────────────────────────────────────────────

def test_hybrid_without_scores_falls_back_to_diversity():
    strategy = ActiveSamplingStrategy(strategy="hybrid", embedder="tfidf")
    pool = make_pool(10)
    selected = strategy.select(pool, k=4, scores=None)
    assert len(selected) == 4


def test_hybrid_with_scores_alternates():
    strategy = ActiveSamplingStrategy(strategy="hybrid", embedder="tfidf")
    pool = make_pool(10)
    scores = [0.1 * i for i in range(10)]
    selected = strategy.select(pool, k=4, scores=scores)
    assert len(selected) == 4


# ── GuardPhaseController integration ─────────────────────────────────────────

def test_phase_controller_accepts_active_sampling_param():
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController(active_sampling=True, active_sampling_k=3)
    assert ctrl is not None


def test_observe_unlabeled_appends_to_buffer():
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController(active_sampling=True)
    chain = make_pool(1)[0]
    ctrl.observe_unlabeled(chain, score=0.6)
    assert len(ctrl._unlabeled_buffer) == 1
    assert len(ctrl._unlabeled_scores) == 1


def test_chains_needing_labels_returns_list():
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController(active_sampling=True, active_sampling_k=2)
    for chain in make_pool(5):
        ctrl.observe_unlabeled(chain, score=0.5)
    result = ctrl.chains_needing_labels()
    assert isinstance(result, list)
    assert len(result) <= 2


def test_chains_needing_labels_k_override():
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController(active_sampling=True, active_sampling_k=2)
    for chain in make_pool(8):
        ctrl.observe_unlabeled(chain, score=0.5)
    result = ctrl.chains_needing_labels(k=4)
    assert len(result) <= 4


def test_observe_unchanged_signature():
    """Existing observe(label: int) must still work without modification."""
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController()
    ctrl.observe(1)
    ctrl.observe(0)


def test_chains_needing_labels_without_active_sampling_raises():
    from llm_guard.agent_guard import GuardPhaseController
    ctrl = GuardPhaseController(active_sampling=False)
    with pytest.raises((RuntimeError, AttributeError)):
        ctrl.chains_needing_labels()
