"""Shared test fixtures for rewardprobe tests."""

import pytest

from rewardprobe.types import NormalizedEnvironment, NormalizedRewardFn

# ---------------------------------------------------------------------------
# Sample tasks
# ---------------------------------------------------------------------------

MATH_TASKS = [
    {"prompt": "What is 2+2?", "answer": "4"},
    {"prompt": "What is 3+3?", "answer": "6"},
    {"prompt": "What is 0+0?", "answer": "0"},
    {"prompt": "What is 5*5?", "answer": "25"},
    {"prompt": "What is 1+1?", "answer": "2"},
    {"prompt": "What is 7+3?", "answer": "10"},
    {"prompt": "What is 9-4?", "answer": "5"},
    {"prompt": "What is 8*2?", "answer": "16"},
    {"prompt": "What is 12/3?", "answer": "4"},
    {"prompt": "What is 15-7?", "answer": "8"},
]


# ---------------------------------------------------------------------------
# Reward functions — intentionally buggy
# ---------------------------------------------------------------------------


def substring_reward(completion: str, task: dict) -> float:
    """Buggy: uses substring match, vulnerable to trivial bypasses."""
    answer = task.get("answer", "")
    return 1.0 if answer in completion else 0.0


def always_one(completion: str, task: dict) -> float:
    """Buggy: always returns 1.0 regardless of input."""
    return 1.0


def crashes_on_empty(completion: str, task: dict) -> float:
    """Buggy: crashes when completion is empty."""
    return 1.0 if completion[0] == task["answer"][0] else 0.0


def out_of_range(completion: str, task: dict) -> float:
    """Buggy: returns values outside [0, 1]."""
    answer = task.get("answer", "")
    if completion.strip() == answer:
        return 2.0
    return -0.5


def nondeterministic(completion: str, task: dict) -> float:
    """Buggy: uses random, non-deterministic."""
    import random

    return random.random()


# ---------------------------------------------------------------------------
# Reward functions — correct
# ---------------------------------------------------------------------------


def exact_match(completion: str, task: dict) -> float:
    """Correct: exact match with normalization."""
    c = completion.strip().lower()
    a = str(task.get("answer", "")).strip().lower()
    if not c or not a:
        return 0.0
    return 1.0 if c == a else 0.0


def format_reward(completion: str, task: dict) -> float:
    """Secondary fn: checks if answer is in a boxed format."""
    return 1.0 if "\\boxed{" in completion else 0.0


# ---------------------------------------------------------------------------
# Environment fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def math_tasks():
    return MATH_TASKS


@pytest.fixture
def buggy_env():
    """Environment with a substring-match reward function."""
    return NormalizedEnvironment(
        name="buggy_test",
        reward_fns=[
            NormalizedRewardFn(name="substring", fn=substring_reward, weight=1.0),
        ],
        tasks=MATH_TASKS,
    )


@pytest.fixture
def clean_env():
    """Environment with a correct exact-match reward function."""
    return NormalizedEnvironment(
        name="clean_test",
        reward_fns=[
            NormalizedRewardFn(name="exact_match", fn=exact_match, weight=1.0),
        ],
        tasks=MATH_TASKS,
    )


@pytest.fixture
def multi_fn_env():
    """Environment with multiple reward functions for composition tests."""
    return NormalizedEnvironment(
        name="multi_fn_test",
        reward_fns=[
            NormalizedRewardFn(name="exact_match", fn=exact_match, weight=1.0),
            NormalizedRewardFn(name="format_reward", fn=format_reward, weight=0.2),
        ],
        tasks=MATH_TASKS,
    )


@pytest.fixture
def always_one_env():
    """Environment with an always-1.0 reward function."""
    return NormalizedEnvironment(
        name="always_one_test",
        reward_fns=[
            NormalizedRewardFn(name="always_one", fn=always_one, weight=1.0),
        ],
        tasks=MATH_TASKS,
    )


@pytest.fixture
def crashing_env():
    """Environment with a reward function that crashes on empty input."""
    return NormalizedEnvironment(
        name="crashing_test",
        reward_fns=[
            NormalizedRewardFn(name="crashes_on_empty", fn=crashes_on_empty, weight=1.0),
        ],
        tasks=MATH_TASKS,
    )


@pytest.fixture
def out_of_range_env():
    """Environment with a reward function that returns out-of-range values."""
    return NormalizedEnvironment(
        name="out_of_range_test",
        reward_fns=[
            NormalizedRewardFn(name="out_of_range", fn=out_of_range, weight=1.0),
        ],
        tasks=MATH_TASKS,
    )
