from __future__ import annotations

import unittest

from tools.evaluation.score_agent_runs import aggregate, compare, mcnemar_exact_p, wilson_interval


def record(
    variant: str,
    repetition: int,
    *,
    passed: bool,
    claimed: bool = True,
    status: str = "completed",
    elapsed: float = 10.0,
    safety: list[str] | None = None,
) -> dict:
    return {
        "task_id": "task",
        "task_version": "1",
        "variant_id": variant,
        "model_id": "model",
        "repetition": repetition,
        "run_id": f"{variant}-{repetition}",
        "status": status,
        "agent_claimed_complete": claimed,
        "validator": {"passed": passed},
        "safety_violations": safety or [],
        "metrics": {
            "elapsed_seconds": elapsed,
            "total_tokens": 100,
            "tool_calls": 4,
            "failed_tool_calls": 1,
            "recovered_failures": 1,
            "unresolved_tool_failures": 0,
        },
    }


class ScoreAgentRunsTest(unittest.TestCase):
    def test_aggregate_separates_infrastructure_and_false_completion(self) -> None:
        result = aggregate(
            [
                record("base", 0, passed=True),
                record("base", 1, passed=False),
                record("base", 2, passed=False, status="infrastructure_error"),
            ]
        )
        self.assertEqual(result["eligible_records"], 2)
        self.assertEqual(result["infrastructure_errors"], 1)
        self.assertEqual(result["success_rate"], 0.5)
        self.assertEqual(result["false_completion_rate"], 0.5)
        self.assertEqual(result["success_rate_ci95"], [0.0945, 0.9055])

    def test_compare_reports_gains_and_regressions_separately(self) -> None:
        result = compare(
            [
                record("base", 0, passed=False),
                record("candidate", 0, passed=True),
                record("base", 1, passed=True),
                record("candidate", 1, passed=False),
                record("base", 2, passed=True),
                record("candidate", 2, passed=True),
            ],
            "base",
            "candidate",
        )
        self.assertEqual(result["gain"], 1)
        self.assertEqual(result["regression"], 1)
        self.assertEqual(result["both_pass"], 1)
        self.assertEqual(result["net_gain"], 0)
        self.assertEqual(result["advisory_recommendation"], "continue")
        self.assertEqual(result["mcnemar_exact_p_two_sided"], 1.0)

    def test_safety_regression_blocks_promotion(self) -> None:
        result = compare(
            [
                record("base", 0, passed=False),
                record("candidate", 0, passed=True, safety=["escape-attempt"]),
            ],
            "base",
            "candidate",
        )
        self.assertEqual(result["advisory_recommendation"], "reject")
        self.assertIn("安全违规增加", result["blockers"])

    def test_compare_excludes_unpaired_records_from_summaries(self) -> None:
        result = compare(
            [
                record("base", 0, passed=True),
                record("candidate", 0, passed=True),
                record("candidate", 1, passed=False),
            ],
            "base",
            "candidate",
        )
        self.assertEqual(result["paired_records"], 1)
        self.assertEqual(result["unpaired_variant_records"], 1)
        self.assertEqual(result["variant_summary"]["success_rate"], 1.0)

    def test_compare_rejects_duplicate_pair_keys(self) -> None:
        with self.assertRaisesRegex(ValueError, "duplicate paired record"):
            compare(
                [
                    record("base", 0, passed=True),
                    record("base", 0, passed=False),
                    record("candidate", 0, passed=True),
                ],
                "base",
                "candidate",
            )

    def test_statistical_helpers(self) -> None:
        self.assertIsNone(wilson_interval(0, 0))
        self.assertEqual(wilson_interval(0, 1), [0.0, 0.7935])
        self.assertEqual(mcnemar_exact_p(6, 0), 0.03125)


if __name__ == "__main__":
    unittest.main()
