"""Tests for the Probe orchestrator and Report output."""

import json
import tempfile
from pathlib import Path

from rewardprobe import Probe
from rewardprobe.report import print_report, report_to_dict, report_to_json


def _exact_match(completion, answer):
    return 1.0 if completion.strip() == str(answer).strip() else 0.0


def _always_one(completion, answer):
    return 1.0


TASKS = [
    {"prompt": "What is 2+2?", "answer": "4"},
    {"prompt": "What is 3+3?", "answer": "6"},
    {"prompt": "What is 1+1?", "answer": "2"},
]


class TestProbe:
    def test_test_fn_returns_report(self):
        probe = Probe()
        report = probe.test_fn(_exact_match, TASKS)
        assert report.environment == "_exact_match"
        assert report.rewardprobe_version == "0.1.0"
        assert isinstance(report.passed, bool)
        assert len(report.families) == 6

    def test_buggy_fn_finds_issues(self):
        probe = Probe()
        report = probe.test_fn(_always_one, TASKS)
        assert not report.passed
        assert report.critical_count > 0

    def test_families_filter(self):
        probe = Probe(families=["basic_sanity"])
        report = probe.test_fn(_exact_match, TASKS)
        assert len(report.families) == 1
        assert report.families[0].name == "basic_sanity"

    def test_num_tasks_respected(self):
        probe = Probe(num_tasks=2)
        report = probe.test_fn(_exact_match, TASKS)
        # Should still run (just with fewer tasks)
        assert len(report.families) > 0

    def test_seed_reproducibility(self):
        probe1 = Probe(seed=42)
        probe2 = Probe(seed=42)
        r1 = probe1.test_fn(_exact_match, TASKS)
        r2 = probe2.test_fn(_exact_match, TASKS)
        assert r1.total_passed == r2.total_passed
        assert r1.critical_count == r2.critical_count


class TestReport:
    def test_to_dict(self):
        probe = Probe()
        report = probe.test_fn(_exact_match, TASKS)
        d = report_to_dict(report)
        assert "environment" in d
        assert "families" in d
        assert "summary" in d
        assert isinstance(d["passed"], bool)

    def test_to_json(self):
        probe = Probe()
        report = probe.test_fn(_exact_match, TASKS)
        with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
            report_to_json(report, f.name)
            data = json.loads(Path(f.name).read_text())
            assert data["environment"] == "_exact_match"
            assert "families" in data

    def test_print_report_does_not_crash(self, capsys):
        probe = Probe()
        report = probe.test_fn(_always_one, TASKS)
        # Just verify it doesn't crash
        print_report(report)
