# tests/test_exp_proxy_chain_quality.py
"""Unit tests for exp_openai_proxy_chain_quality helpers."""
import sys
from pathlib import Path
import pytest

ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "experiments"))


# ── Label logic tests ────────────────────────────────────────────────────────

def test_label_fuzzy_match_correct():
    # gold is the short dataset answer; predicted is the model's longer response
    # Use a case where gold tokens are a large fraction of predicted tokens → F1 >= 0.5
    from exp_openai_proxy_chain_quality import compute_label
    assert compute_label("Paris France", "Paris France capital") == 1  # F1 = 0.8

def test_label_fuzzy_match_wrong():
    from exp_openai_proxy_chain_quality import compute_label
    assert compute_label("London England", "Paris is the capital of France") == 0

def test_label_numeric_matches_anywhere_in_pred():
    from exp_openai_proxy_chain_quality import compute_label
    assert compute_label("1945", "1945") == 1
    assert compute_label("1945", "around 1945 approximately") == 1  # number appears in pred
    assert compute_label("1945", "the answer is 1946") == 0          # wrong number

def test_label_yesno_first_word_match():
    from exp_openai_proxy_chain_quality import compute_label
    assert compute_label("yes", "yes") == 1
    assert compute_label("yes", "yes it is definitely true") == 1        # first word
    assert compute_label("yes", "no it is not") == 0                     # first word wrong
    assert compute_label("no", "No, they are different") == 1            # case-insensitive
    # Verbose GPT-style answers: gold word anywhere or semantic words
    assert compute_label("yes", "Both are American filmmakers.") == 1    # "both" in YES_WORDS
    assert compute_label("yes", "They share the same nationality.") == 1 # "same" in YES_WORDS
    assert compute_label("no", "They are from different countries.") == 1 # "different" in NO_WORDS

def test_label_short_non_numeric_uses_fuzzy():
    from exp_openai_proxy_chain_quality import compute_label
    # "Rome" is short but not numeric/yes-no → uses fuzzy F1 (not exact match)
    # F1("Rome", "capital Rome") = 2*(1/2*1.0)/(0.5+1.0) = 0.667 >= 0.5 → correct (1)
    assert compute_label("Rome", "capital Rome") == 1
    # Wrong answer: no overlap → F1=0 < 0.5 → returns 0
    assert compute_label("Rome", "Athens is the capital") == 0


# ── RawMessageCapture tests ──────────────────────────────────────────────────

def test_raw_message_capture_accumulates_tool_calls():
    """Messages accumulated across multiple on_llm_end calls."""
    from exp_openai_proxy_chain_quality import RawMessageCapture
    cap = RawMessageCapture()

    # Simulate first LLM call: emits a tool_call
    class FakeGen:
        message = type("M", (), {
            "tool_calls": [{"function": {"name": "wikipedia_search",
                                         "arguments": '{"query": "Paris"}'}}],
            "content": "",
        })()
    class FakeResult:
        generations = [[FakeGen()]]

    cap.on_llm_end(FakeResult(), **{})
    assert len(cap.messages) == 1

    # Simulate second LLM call: plain content (final answer)
    class FakeGen2:
        message = type("M", (), {"tool_calls": [], "content": "Paris is the capital."})()
    class FakeResult2:
        generations = [[FakeGen2()]]

    cap.on_llm_end(FakeResult2(), **{})
    assert len(cap.messages) == 2


def test_raw_message_capture_reset():
    from exp_openai_proxy_chain_quality import RawMessageCapture
    cap = RawMessageCapture()
    cap.messages.append({"content": "test"})
    cap.reset()
    assert cap.messages == []


def test_raw_message_capture_build_proxy_chain():
    """build_proxy_chain() assembles steps from accumulated messages."""
    from exp_openai_proxy_chain_quality import RawMessageCapture

    cap = RawMessageCapture()

    class Msg1:
        tool_calls = [{"function": {"name": "search", "arguments": '{"query": "Rome"}'}}]
        content = ""
    class Msg2:
        tool_calls = []
        content = "Rome is the capital of Italy."

    class Gen:
        def __init__(self, msg): self.message = msg
    class Res:
        def __init__(self, msg): self.generations = [[Gen(msg)]]

    cap.on_llm_end(Res(Msg1()), **{})
    cap.on_llm_end(Res(Msg2()), **{})

    chain = cap.build_proxy_chain("What is the capital of Italy?", "Rome is the capital of Italy.")
    assert chain["question"] == "What is the capital of Italy?"
    assert any(s["action_type"] == "search" for s in chain["steps"])
    assert any(s["action_type"] == "Finish" for s in chain["steps"])
    assert all(s["observation"] == "" for s in chain["steps"])


# ── Gate logic tests ─────────────────────────────────────────────────────────

def test_gate_pass():
    from exp_openai_proxy_chain_quality import evaluate_gate
    result = evaluate_gate(
        proxy_auroc=0.62, proxy_auroc_ci_lo=0.51,
        proxy_auprc=0.45, class_prior=0.40,
        gold_auroc=0.70,
    )
    assert result == "PASS"

def test_gate_fail_auroc_too_low():
    from exp_openai_proxy_chain_quality import evaluate_gate
    result = evaluate_gate(
        proxy_auroc=0.58, proxy_auroc_ci_lo=0.51,
        proxy_auprc=0.45, class_prior=0.40,
        gold_auroc=0.70,
    )
    assert result == "FAIL"

def test_gate_fail_ci_lo_too_low():
    from exp_openai_proxy_chain_quality import evaluate_gate
    result = evaluate_gate(
        proxy_auroc=0.65, proxy_auroc_ci_lo=0.49,
        proxy_auprc=0.45, class_prior=0.40,
        gold_auroc=0.70,
    )
    assert result == "FAIL"

def test_gate_fail_auprc_below_prior():
    from exp_openai_proxy_chain_quality import evaluate_gate
    result = evaluate_gate(
        proxy_auroc=0.65, proxy_auroc_ci_lo=0.52,
        proxy_auprc=0.30, class_prior=0.40,
        gold_auroc=0.70,
    )
    assert result == "FAIL"

def test_gate_fail_delta_too_large():
    from exp_openai_proxy_chain_quality import evaluate_gate
    result = evaluate_gate(
        proxy_auroc=0.62, proxy_auroc_ci_lo=0.51,
        proxy_auprc=0.45, class_prior=0.40,
        gold_auroc=0.80,  # delta = 0.18 > 0.15
    )
    assert result == "FAIL"


# ── sc_old_score ablation test ───────────────────────────────────────────────

def test_sc_old_ablated_zeroes_obs_signals():
    from exp_openai_proxy_chain_quality import sc_old_score, sc_old_score_ablated
    chain = {
        "question": "What is X?",
        "steps": [
            {"action_type": "Search", "action_arg": "X", "observation": ""},
            {"action_type": "Finish", "action_arg": "X is Y", "observation": ""},
        ],
        "final_answer": "X is Y",
    }
    full_score    = sc_old_score(chain)
    ablated_score = sc_old_score_ablated(chain)
    # Both scores must be valid
    assert 0.0 <= full_score    <= 1.0
    assert 0.0 <= ablated_score <= 1.0
    # With all empty observations:
    #   full:    0.35*loop_rate + 0.25*1.0 + 0.15*ans_gap + 0.10*1.0 + 0.15*1.0
    #   ablated: 0.35*loop_rate + 0.25*0.0 + 0.15*ans_gap + 0.10*0.0 + 0.15*0.0
    # ablated must be strictly less (obs signals add 0.50 to full when obs are empty)
    assert ablated_score < full_score, (
        f"ablated ({ablated_score:.4f}) should be < full ({full_score:.4f}) "
        f"when observations are empty (obs signals = +0.50 in full)"
    )
    # Verify decomposition: full - ablated ≈ 0.25*1.0 + 0.10*1.0 + 0.15*1.0 = 0.50
    assert abs((full_score - ablated_score) - 0.50) < 0.01, (
        f"obs-signal contribution should be ~0.50, got {full_score - ablated_score:.4f}"
    )
