"""Tests for descriptive EasyCom AI pre-label aggregation."""

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path

from summarize_easycom_event_prelabels import summarize


class EventPrelabelSummaryTests(unittest.TestCase):
    def test_summary_separates_ai_prelabels_and_two_pass_agreement(self) -> None:
        base = {
            "protocol_version": "p1", "schema_version": 6,
            "candidate_source": "WEARER_ONSET", "task": "wearer_entry",
            "human_status": "UNREVIEWED", "ai_confidence": 0.8,
            "llm_adjudication": {"prompt_version": "v1", "model": "m", "passes": 2},
        }
        rows = [
            {**base, "uid": "u1", "session": "Session_1", "ai_verdict": "POSITIVE",
             "ai_event_type": "WEARER_ENTRY", "ai_subtype": None,
             "llm_adjudication": {**base["llm_adjudication"], "decisions": [
                 {"verdict": "POSITIVE", "event_type": "WEARER_ENTRY"},
                 {"verdict": "POSITIVE", "event_type": "WEARER_ENTRY"},
             ]}},
            {**base, "uid": "u2", "session": "Session_4", "task": "int_yield",
             "candidate_source": "WEARER_OFFSET", "ai_verdict": "EXCLUDE",
             "ai_event_type": None, "ai_subtype": None,
             "ai_exclude_reason": "OUT_OF_SCOPE_NORMAL_COMPLETION",
             "ai_exclude_reason_group": "OUT_OF_SCOPE",
             "llm_adjudication": {**base["llm_adjudication"], "decisions": [
                 {"verdict": "EXCLUDE", "exclude_reason": "OUT_OF_SCOPE_NORMAL_COMPLETION"},
                 {"verdict": "EXCLUDE", "exclude_reason": "INSUFFICIENT_EVIDENCE"},
             ]}},
        ]
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            source = root / "input.jsonl"
            source.write_text("".join(json.dumps(row) + "\n" for row in rows))
            summary = summarize([source], root / "out")
            cohorts = {item["cohort"]: item for item in summary["cohorts"]}
            self.assertEqual(cohorts["Test pooled (Sessions 1-3)"]["num_positive"], 1)
            self.assertEqual(cohorts["Session 4 audit (train split)"]["num_exclude"], 1)
            self.assertEqual(cohorts["All supplied (descriptive only)"]["num_two_pass_semantic_agreement"], 1)
            self.assertEqual(cohorts["All supplied (descriptive only)"]["num_two_pass_full_decision_agreement"], 1)
            self.assertIn("AI pre-label audit only", (root / "out/manuscript_ai_prelabel_table.tex").read_text())


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