from __future__ import annotations

import json
from pathlib import Path
import tempfile
import unittest

from prepare_ego4d_av_eot_int_review import (
    CandidateConfig, derive_clip_candidates, iter_clips, load_clip_segments,
)
from review_floor_events import CandidateRepository, load_candidates


def segment(person: str, start: float, end: float) -> dict[str, object]:
    return {"person": person, "start_time": start, "end_time": end}


def fixture_clip() -> dict[str, object]:
    return {
        "clip_uid": "clip-001",
        "video_uid": "video-001",
        "clip_start_sec": 0,
        "clip_end_sec": 12.0,
        "valid": True,
        "camera_wearer": {
            "person_id": "wearer",
            "camera_wearer": True,
            "voice_segments": [segment("wearer", 2.0, 3.0), segment("wearer", 6.0, 8.0)],
        },
        "persons": [
            {
                "person_id": "p1", "camera_wearer": False,
                "voice_segments": [segment("p1", 0.5, 1.8), segment("p1", 3.2, 4.0), segment("p1", 6.5, 8.8)],
            },
            {
                "person_id": "p2", "camera_wearer": False,
                "voice_segments": [segment("p2", 4.3, 5.2)],
            },
        ],
        "transcriptions": [
            {"person_id": "p1", "start_time_sec": 0.5, "end_time_sec": 1.8, "transcription": "Can you answer?"},
            {"person_id": "wearer", "start_time_sec": 2.0, "end_time_sec": 3.0, "transcription": "Yes."},
            {"person_id": "p1", "start_time_sec": 6.5, "end_time_sec": 8.8, "transcription": "Wait, one thing."},
        ],
        "social_segments_talking": [
            {"person": "p1", "start_time": 0.5, "end_time": 1.8, "is_at_me": True},
        ],
    }


class Ego4DAVCandidateTests(unittest.TestCase):
    def test_loads_roles_transcript_and_social_evidence(self) -> None:
        rows, wearer_id = load_clip_segments(fixture_clip())
        self.assertEqual(wearer_id, "wearer")
        question = next(row for row in rows if row.start == 0.5)
        self.assertEqual(question.transcription, "Can you answer?")
        self.assertIs(question.addressed_to_wearer, True)
        self.assertEqual(next(row.role for row in rows if row.person_id == "wearer"), "wearer")

    def test_derives_raw_wearer_boundaries_and_controls(self) -> None:
        rows, timeline = derive_clip_candidates(
            fixture_clip(), split="train", config=CandidateConfig()
        )
        sources = [row["candidate_source"] for row in rows]
        self.assertEqual(sources.count("WEARER_ONSET"), 2)
        self.assertEqual(sources.count("WEARER_OFFSET"), 2)
        self.assertIn("OTHER_TO_OTHER", sources)
        onset = next(row for row in rows if row["candidate_source"] == "WEARER_ONSET")
        option = onset["temporal_evidence"]["opportunity_options"][0]
        self.assertEqual(option["participant_id"], "p1")
        self.assertIs(option["addressed_to_wearer"], True)
        offset = next(row for row in rows if row["candidate_source"] == "WEARER_OFFSET" and row["event_time"] == 8.0)
        self.assertTrue(offset["temporal_evidence"]["other_crosses_wearer_offset"])
        self.assertTrue(offset["temporal_evidence"]["new_other_onset_during_wearer"])
        self.assertTrue(all(row["start_seconds"] < row["end_seconds"] for row in timeline))

    def test_duplicate_clip_uid_is_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            payload = {"videos": [{"split": "train", "clips": [fixture_clip()]}]}
            paths = []
            for name in ("av_train.json", "av_val.json"):
                path = root / name
                path.write_text(json.dumps(payload), encoding="utf-8")
                paths.append(path)
            with self.assertRaisesRegex(ValueError, "duplicate clip_uid"):
                list(iter_clips(paths))

    def test_review_repository_uses_ego4d_media_and_embedded_transcript(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            media = root / "v1" / "clips" / "clip-001.mp4"
            media.parent.mkdir(parents=True)
            media.write_bytes(b"not-a-real-video")
            rows, _ = derive_clip_candidates(
                fixture_clip(), split="val", config=CandidateConfig()
            )
            events = root / "ego4d_av_candidates.jsonl"
            events.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8")
            candidates = load_candidates([events])
            repository = CandidateRepository(
                candidates, root, dataset="auto",
                context_before_seconds=3.0, context_after_seconds=3.0,
            )
            self.assertEqual(repository.dataset, "ego4d")
            self.assertEqual(repository.media_path(candidates[0]), media.resolve())
            detail = repository.detail(candidates[0])
            self.assertTrue(detail["transcript_available"])
            self.assertTrue(any(row["transcription"] == "Can you answer?" for row in detail["transcript"]))

    def test_rejects_missing_camera_wearer(self) -> None:
        clip = fixture_clip()
        clip.pop("camera_wearer")
        with self.assertRaisesRegex(ValueError, "exactly one camera wearer"):
            load_clip_segments(clip)

    def test_bounded_official_terminal_endpoint_is_preserved_and_audited(self) -> None:
        clip = fixture_clip()
        clip["clip_end_sec"] = 8.76
        rows, _ = derive_clip_candidates(clip, split="val", config=CandidateConfig())
        offset = next(
            row for row in rows
            if row["candidate_source"] == "WEARER_OFFSET" and row["event_time"] == 8.0
        )
        self.assertAlmostEqual(
            offset["temporal_evidence"]["terminal_voice_overrun_seconds"], 0.04
        )

    def test_large_terminal_endpoint_overrun_is_rejected(self) -> None:
        clip = fixture_clip()
        clip["clip_end_sec"] = 11.95
        clip["camera_wearer"]["voice_segments"].append(segment("wearer", 11.8, 12.2))
        with self.assertRaisesRegex(ValueError, "exceeds the audited limit"):
            derive_clip_candidates(clip, split="val", config=CandidateConfig())


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