#!/usr/bin/env python3
"""Writer 全评测集预注册顺序与盲化分配测试。"""

from __future__ import annotations

import hashlib
import pathlib
import sys
import unittest

PROJECT_ROOT = pathlib.Path(__file__).resolve().parents[3]
SCRIPT_DIR = PROJECT_ROOT / ".claude" / "skills" / "evaluate-frozen-replay" / "scripts"
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from writer_eval_preregister import (  # noqa: E402
    PreregistrationError,
    build_balanced_preregistration,
    validate_balanced_position_counts,
)


class WriterEvalPreregisterTest(unittest.TestCase):
    """覆盖评测集级平衡、确定性和独立盲化命名空间。"""

    def setUp(self):
        """固定五个样本，复现 Gate A 最小预注册集合。"""

        self.evaluation_set_version = "writer-gate-a-deep-space-v1"
        self.sample_ids = [
            "deep-space-489-battle",
            "deep-space-321-character-dialogue",
            "deep-space-544-turning-point",
            "deep-space-199-information-reveal",
            "deep-space-523-returning-character",
        ]

    def test_five_samples_use_hash_sort_and_balanced_abc_rotations(self):
        """五样本先按规定哈希排序，再整体循环 ABC/BCA/CAB。"""

        registration = build_balanced_preregistration(
            evaluation_set_version=self.evaluation_set_version,
            sample_ids=self.sample_ids,
        )
        expected_order = sorted(
            self.sample_ids,
            key=lambda sample_id: hashlib.sha256(
                f"{self.evaluation_set_version}{sample_id}".encode("utf-8")
            ).hexdigest(),
        )
        self.assertEqual(
            [item["sampleId"] for item in registration["armOrderTable"]],
            expected_order,
        )
        self.assertEqual(
            [item["armOrder"] for item in registration["armOrderTable"]],
            [
                ["A", "B", "C"],
                ["B", "C", "A"],
                ["C", "A", "B"],
                ["A", "B", "C"],
                ["B", "C", "A"],
            ],
        )
        validate_balanced_position_counts(registration["armOrderTable"], order_field="armOrder")
        validate_balanced_position_counts(
            registration["blindAssignmentTable"], order_field="armOrder"
        )
        for table_name in ("armOrderTable", "blindAssignmentTable"):
            counts = registration["positionCounts"][table_name]
            for position in counts:
                self.assertLessEqual(max(position.values()) - min(position.values()), 1)

    def test_blind_assignment_is_deterministic_but_uses_independent_namespace(self):
        """盲化分配可复现，但不能复用运行顺序表的样本排序。"""

        first = build_balanced_preregistration(
            evaluation_set_version=self.evaluation_set_version,
            sample_ids=self.sample_ids,
        )
        second = build_balanced_preregistration(
            evaluation_set_version=self.evaluation_set_version,
            sample_ids=list(reversed(self.sample_ids)),
        )
        self.assertEqual(first, second)
        self.assertNotEqual(
            [item["sampleId"] for item in first["armOrderTable"]],
            [item["sampleId"] for item in first["blindAssignmentTable"]],
        )
        self.assertEqual(first["blindNamespace"], "writer-blind-assignment-v1")

    def test_duplicate_samples_and_unbalanced_tables_fail_closed(self):
        """重复样本或逐样本随机构造出的失衡表都不能通过预注册。"""

        with self.assertRaises(PreregistrationError):
            build_balanced_preregistration(
                evaluation_set_version=self.evaluation_set_version,
                sample_ids=["same", "same"],
            )
        with self.assertRaises(PreregistrationError):
            validate_balanced_position_counts(
                [
                    {"sampleId": "one", "armOrder": ["A", "B", "C"]},
                    {"sampleId": "two", "armOrder": ["A", "B", "C"]},
                    {"sampleId": "three", "armOrder": ["A", "B", "C"]},
                ],
                order_field="armOrder",
            )


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