"""Tests for the RawAdapter."""

from rewardprobe.adapters.raw import RawAdapter, _wrap_callable

TASKS = [
    {"prompt": "What is 2+2?", "answer": "4", "question": "math"},
]


class TestWrapCallable:
    def test_two_param_completion_answer(self):
        def fn(completion, answer):
            return 1.0 if answer in completion else 0.0

        wrapped = _wrap_callable(fn)
        assert wrapped("4", {"answer": "4"}) == 1.0
        assert wrapped("wrong", {"answer": "4"}) == 0.0

    def test_single_param(self):
        def fn(completion):
            return 1.0 if len(completion) > 0 else 0.0

        wrapped = _wrap_callable(fn)
        assert wrapped("hello", {}) == 1.0
        assert wrapped("", {}) == 0.0

    def test_kwargs_passthrough(self):
        def fn(completion, **kwargs):
            answer = kwargs.get("answer", "")
            return 1.0 if completion == answer else 0.0

        wrapped = _wrap_callable(fn)
        assert wrapped("4", {"answer": "4"}) == 1.0

    def test_named_params_from_task(self):
        def fn(completion, answer, question):
            return 1.0 if completion == answer else 0.0

        wrapped = _wrap_callable(fn)
        assert wrapped("4", {"answer": "4", "question": "math"}) == 1.0


class TestRawAdapter:
    def test_single_fn(self):
        def fn(completion, answer):
            return 1.0 if answer in completion else 0.0

        adapter = RawAdapter(fn=fn, tasks=TASKS)
        env = adapter.adapt()
        assert env.name == "fn"
        assert len(env.reward_fns) == 1
        assert len(env.tasks) == 1

    def test_multiple_fns(self):
        def fn1(completion, answer):
            return 1.0

        def fn2(completion, answer):
            return 0.5

        adapter = RawAdapter(fn=[fn1, fn2], tasks=TASKS, weights=[1.0, 0.5])
        env = adapter.adapt()
        assert len(env.reward_fns) == 2
        assert env.reward_fns[0].weight == 1.0
        assert env.reward_fns[1].weight == 0.5

    def test_wrapped_fn_works_with_tasks(self):
        def fn(completion, answer):
            return 1.0 if completion.strip() == answer else 0.0

        adapter = RawAdapter(fn=fn, tasks=TASKS)
        env = adapter.adapt()
        # Test the wrapped function directly
        score = env.reward_fns[0].fn("4", TASKS[0])
        assert score == 1.0
        score = env.reward_fns[0].fn("wrong", TASKS[0])
        assert score == 0.0
