from __future__ import annotations

import argparse
import json
from pathlib import Path
import tempfile
import unittest
from unittest import mock

import numpy as np

import evaluate_annotated_floor_events as evaluator
from evaluate_annotated_floor_events import load_annotated_opportunities
from evaluate_floor_state_changes import NUM_CLASSES, NUM_HORIZONS
from materialize_provisional_llm_reviews import LABEL_SOURCE, materialize


def llm_candidate(uid: str, *, unanimous: bool) -> dict[str, object]:
    suggestion = {
        "verdict": "POSITIVE",
        "event_type": "WEARER_ENTRY",
        "subtype": None,
        "exclude_reason": None,
        "selected_opportunity_uid": None,
        "confidence": 0.9,
        "reason": "The wearer takes the floor.",
        "needs_human_review": True,
        "unanimous": unanimous,
    }
    return {
        "uid": uid,
        "session": "Session_4",
        "split": "train",
        "task": "wearer_entry",
        "candidate_source": "WEARER_ONSET",
        "anchor_dense_frame": 40,
        "anchor_dense_seconds": 2.0,
        "temporal_evidence": {
            "behavior_anchor_frame": 40,
            "behavior_anchor_seconds": 2.0,
            "opportunity_options": [],
        },
        "ai_verdict": "POSITIVE",
        "ai_event_type": "WEARER_ENTRY",
        "ai_subtype": None,
        "ai_exclude_reason": None,
        "ai_confidence": 0.9,
        "llm_suggestion": suggestion,
    }


class ProvisionalLLMReviewTests(unittest.TestCase):
    def test_only_unanimous_rows_are_materialized_with_nonhuman_provenance(self) -> None:
        with tempfile.TemporaryDirectory() as temporary:
            root = Path(temporary)
            events = root / "events.jsonl"
            events.write_text(
                "".join(
                    json.dumps(row) + "\n"
                    for row in (
                        llm_candidate("unanimous", unanimous=True),
                        llm_candidate("disagreement", unanimous=False),
                    )
                ),
                encoding="utf-8",
            )
            output = root / "provisional_reviews.jsonl"
            records, summary = materialize([events], output)

            self.assertEqual(len(records), 1)
            self.assertEqual(summary["skipped"], {"nonunanimous": 1})
            self.assertEqual(records[0]["review_source"], LABEL_SOURCE)
            self.assertFalse(records[0]["human_reviewed"])
            self.assertFalse(records[0]["review"]["human_reviewed"])
            self.assertEqual(records[0]["review"]["human_status"], "REVIEWED")
            self.assertTrue(output.with_suffix(".csv").is_file())
            self.assertTrue(output.with_suffix(".summary.json").is_file())

            with self.assertRaisesRegex(ValueError, "allow-provisional-llm-reviews"):
                load_annotated_opportunities([output], ["Session_4"])
            opportunities, diagnostics = load_annotated_opportunities(
                [output], ["Session_4"], allow_provisional_llm=True
            )
            self.assertEqual(len(opportunities), 1)
            self.assertEqual(diagnostics["num_provisional_llm"], 1)

            run_dir = root / "run"
            (run_dir / "pred").mkdir(parents=True)
            (run_dir / "valid").mkdir()
            np.save(
                run_dir / "pred" / "Session_4.npy",
                np.full((20, NUM_CLASSES, NUM_HORIZONS), 0.5, dtype=np.float32),
            )
            np.save(
                run_dir / "valid" / "Session_4.npy",
                np.ones((20, NUM_HORIZONS), dtype=bool),
            )
            evaluation_output = root / "evaluation.json"
            args = argparse.Namespace(
                reviews=[output],
                run=[f"model={run_dir}"],
                sessions=["Session_4"],
                calibration_reviews=None,
                calibration_run=None,
                calibration_sessions=None,
                matched_recalls=[0.8],
                allow_in_sample_thresholds=False,
                allow_provisional_llm_reviews=True,
                output=evaluation_output,
                csv_output=None,
                event_csv_output=None,
                print_event_examples=0,
            )
            with mock.patch.object(evaluator, "parse_args", return_value=args):
                evaluator.main()
            evaluation = json.loads(evaluation_output.read_text(encoding="utf-8"))
            self.assertTrue(evaluation["provisional_llm"])
            self.assertEqual(
                evaluation["protocol"], "provisional-llm-fixed-horizon-floor-events-v1"
            )
            self.assertIn("no human review", evaluation["semantic_scope"])


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