"""Audit tests for local frame sampling.

A synthetic recording paints the source frame index into the red channel.
The sampler must recover those indices, including when the MP4 index sits
at the end of the file, and the proxy must stay inside the video-token budget.
"""

from __future__ import annotations

import json
import shutil
import subprocess
import tempfile
import unittest
from pathlib import Path

from cystoqa.engine import MM_PROCESSOR_KWARGS
from cystoqa.sampling import (
    SampleManifest,
    frame_geometry,
    linspace_indices,
    plan_sample,
    probe_video,
    sample_video,
    smart_resize,
    tokens_for,
)


def _ffmpeg(*args: str) -> None:
    proc = subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", *args],
        capture_output=True, text=True,
    )
    if proc.returncode != 0:
        raise AssertionError(proc.stderr)


def _proxy_reds(path: Path, width: int, height: int) -> list[float]:
    proc = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", str(path), "-f", "rawvideo", "-pix_fmt", "rgb24", "-"],
        capture_output=True, check=True,
    )
    frame = width * height * 3
    if frame < 1 or len(proc.stdout) % frame:
        raise AssertionError(f"unexpected raw size {len(proc.stdout)} for {width}x{height}")
    means = []
    for start in range(0, len(proc.stdout), frame):
        reds = proc.stdout[start:start + frame:3]
        means.append(sum(reds) / len(reds))
    return means


class LinspaceTests(unittest.TestCase):
    def test_matches_numpy_endpoints(self) -> None:
        # np.linspace(0, 19, 4).round()
        self.assertEqual(linspace_indices(19, 4), [0, 6, 13, 19])
        self.assertEqual(linspace_indices(19, 1), [0])

    def test_token_budget_never_exceeds_cap(self) -> None:
        indices, width, height, estimate = plan_sample(
            nb_frames=50 * 1574,
            source_fps=50.0,
            src_w=1920,
            src_h=1080,
            target_fps=2.0,
            token_budget=200_000,
            max_frames=4096,
        )
        self.assertLessEqual(estimate, 200_000)
        self.assertEqual(estimate, tokens_for(width, height, len(indices)))
        self.assertEqual(len(indices) % 2, 0)
        self.assertGreater(len(indices), 1000)
        self.assertEqual(indices[0], 0)
        self.assertEqual(indices[-1], 50 * 1574 - 1)

    def test_one_frame_window_is_repeated_in_the_plan(self) -> None:
        indices, width, height, estimate = plan_sample(
            nb_frames=1, source_fps=25.0, src_w=64, src_h=64,
            target_fps=2.0, token_budget=1_000,
        )
        self.assertEqual(indices, [0, 0])
        self.assertEqual((width, height), (64, 64))
        self.assertEqual(estimate, tokens_for(64, 64, 2))

    def test_preamble_keeps_span_rate_and_question(self) -> None:
        overview = SampleManifest(
            method="ffmpeg_proxy_video_v1",
            source="recording.mp4",
            source_nb_frames=39348,
            source_fps=25.0,
            duration_sec=1573.92,
            width_src=1920,
            height_src=1080,
            target_fps=2.0,
            start_sec=0.0,
            end_sec=1573.92,
            n_emitted=3146,
            token_budget=200_000,
            width=480,
            height=256,
            estimated_tokens=188760,
            frames=[],
            proxy_fps=2.0,
        )
        text = overview.preamble("Any lesions?")
        self.assertEqual(
            text,
            "Procedure time 0.0–1573.9s, sampled at 2 fps.\n"
            "Frame timestamps are procedure time.\n"
            "\n"
            "Question: Any lesions?",
        )
        clip = SampleManifest(
            method="ffmpeg_proxy_video_v1",
            source="recording.mp4",
            source_nb_frames=39348,
            source_fps=25.0,
            duration_sec=1573.92,
            width_src=1920,
            height_src=1080,
            target_fps=8.0,
            start_sec=100.0,
            end_sec=105.0,
            n_emitted=40,
            token_budget=200_000,
            width=480,
            height=256,
            estimated_tokens=1000,
            frames=[],
            proxy_fps=8.0,
        )
        self.assertEqual(
            clip.preamble(),
            "Procedure time 100.0–105.0s, sampled at 8 fps.\n"
            "Frame timestamps start at 0, which is procedure time 100.0s.",
        )

    def test_geometry_is_patch_aligned(self) -> None:
        width, height = frame_geometry(1920, 1080, 65_536)
        self.assertEqual(width % 32, 0)
        self.assertEqual(height % 32, 0)
        self.assertLessEqual(width * height, 65_536)

    def test_smart_resize_matches_qwen_video_rule(self) -> None:
        # 90 rounds to the nearest multiple of 32, which is 96, not the floor 64.
        height, width = smart_resize(
            8, 90, 90, temporal_factor=2, factor=32,
            min_pixels=4096, max_pixels=10**12,
        )
        self.assertEqual((width, height), (96, 96))
        # Half-integers use Python's round-half-even: 80/32 = 2.5 → 2 → 64.
        height, width = smart_resize(
            8, 80, 80, temporal_factor=2, factor=32,
            min_pixels=4096, max_pixels=10**12,
        )
        self.assertEqual((width, height), (64, 64))
        # Hugging Face video longest_edge is 25,165,824 pixels for the whole clip.
        height, width = smart_resize(
            3146, 1080, 1920, temporal_factor=2, factor=32,
            min_pixels=4096, max_pixels=25_165_824,
        )
        self.assertEqual((width, height), (96, 64))
        self.assertLessEqual(3146 * height * width, 25_165_824)

    def test_plan_uses_smart_resize(self) -> None:
        indices, width, height, estimate = plan_sample(
            nb_frames=80, source_fps=10.0, src_w=90, src_h=90,
            target_fps=1.0, token_budget=200_000, max_frames=4096,
        )
        self.assertEqual(len(indices), 8)
        self.assertEqual((width, height), (96, 96))
        self.assertEqual(estimate, tokens_for(96, 96, 8))
        self.assertLessEqual(estimate, 200_000)


CLINICAL = Path(
    "<LAB_ROOT>/inference/agent/uploads/"
    "M_03022020093955_00000hn20025294o_1_001_001-1.MP4"
)
SHORT = Path(
    "<LAB_ROOT>/inference/agent/uploads/HN26041968N.mp4"
)


class SyntheticVideoTests(unittest.TestCase):
    def test_recovers_frame_indices_with_moov_at_end(self) -> None:
        if shutil.which("ffmpeg") is None:
            self.skipTest("ffmpeg is not installed")
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            video = root / "indexed.mp4"
            # 10 fps, 2.0 s, 20 frames. Red channel = source frame index.
            # No +faststart, so the moov atom is written at the end.
            _ffmpeg(
                "-f", "lavfi",
                "-i", "color=c=black:s=64x64:r=10:d=2,format=rgb24,geq=r='N':g=0:b=0",
                "-c:v", "libx264", "-pix_fmt", "yuv420p",
                str(video),
            )
            raw = video.read_bytes()
            self.assertGreaterEqual(raw.rfind(b"moov"), len(raw) // 2)

            meta = probe_video(video)
            self.assertEqual(meta["nb_frames"], 20)
            self.assertAlmostEqual(meta["source_fps"], 10.0, places=3)

            manifest = sample_video(video, 2.0, root / "out", token_budget=10_000)
            self.assertEqual([frame.source_index for frame in manifest.frames], [0, 6, 13, 19])
            self.assertEqual(manifest.method, "ffmpeg_proxy_video_v1")
            self.assertLessEqual(manifest.estimated_tokens, 10_000)
            self.assertEqual(manifest.width % 32, 0)
            self.assertEqual(manifest.height % 32, 0)
            self.assertTrue(Path(manifest.manifest_path).is_file())
            self.assertTrue(Path(manifest.proxy_path).is_file())
            saved = json.loads(Path(manifest.manifest_path).read_text())
            self.assertEqual(saved["n_emitted"], 4)
            reds = _proxy_reds(Path(manifest.proxy_path), manifest.width, manifest.height)
            self.assertEqual(len(reds), 4)
            for frame, red in zip(manifest.frames, reds):
                self.assertAlmostEqual(red, frame.source_index, delta=8.0)

    def test_message_is_one_passthrough_video(self) -> None:
        if shutil.which("ffmpeg") is None:
            self.skipTest("ffmpeg is not installed")
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            video = root / "clip.mp4"
            _ffmpeg(
                "-f", "lavfi", "-i", "color=c=white:s=64x64:r=10:d=1",
                "-c:v", "libx264", "-pix_fmt", "yuv420p", str(video),
            )
            manifest = sample_video(video, 2.0, root / "out", token_budget=10_000)
            from cystoqa.engine import Hooks
            part = Hooks().video_part(manifest, "http://media/model_video.mp4")
            messages = [{"role": "user", "content": [
                {"type": "text", "text": manifest.preamble("q")}, part,
            ]}]
            types = [item["type"] for item in messages[0]["content"]]
            self.assertEqual(types, ["text", "video_url"])
            self.assertNotIn("image_url", types)
            self.assertNotIn("model_read", part)
            self.assertEqual(MM_PROCESSOR_KWARGS["do_sample_frames"], False)
            self.assertEqual(MM_PROCESSOR_KWARGS["do_resize"], False)
            self.assertEqual(MM_PROCESSOR_KWARGS["cap_pixels_per_frame"], False)
            self.assertNotIn("fps", MM_PROCESSOR_KWARGS)
            self.assertNotIn("size", MM_PROCESSOR_KWARGS)

    def test_chat_readout_stays_out_of_the_prompt(self) -> None:
        from types import SimpleNamespace

        from cystoqa.web import WebHooks, _client_history, _for_llm

        manifest = SimpleNamespace(
            proxy_path="<TMP>/model_video.mp4", proxy_fps=2.0, width=480, height=256,
        )
        part = WebHooks(lambda _evt: None).video_part(manifest, "http://media/model_video.mp4")
        messages = [{"role": "user", "content": [part]}]
        shown = _client_history(messages)[0]["content"][0]
        sent = _for_llm(messages)[0]["content"][0]
        self.assertEqual(shown["model_read"], {"fps": 2.0, "width": 480, "height": 256})
        self.assertNotIn("_local_path", shown)
        self.assertEqual(set(sent), {"type", "video_url"})
        self.assertEqual(sent["video_url"]["url"], "http://media/model_video.mp4")

    def test_clinical_probe_follows_decoded_frames(self) -> None:
        """The container lists each field as a sample: 78696 at 50 fps.

        The decoder emits one picture per two samples, 39348 frames at 25 fps.
        Planning against the container count asks select() for indices that
        do not exist, and the proxy comes back with exactly half the frames.
        """
        if not CLINICAL.is_file() or shutil.which("ffmpeg") is None:
            self.skipTest("clinical recording is not on this machine")
        meta = probe_video(CLINICAL)
        self.assertEqual(meta["nb_frames"], 39348)
        self.assertAlmostEqual(meta["source_fps"], 25.0, places=3)
        indices, width, height, estimate = plan_sample(
            nb_frames=meta["nb_frames"],
            source_fps=meta["source_fps"],
            src_w=meta["width"],
            src_h=meta["height"],
            target_fps=2.0,
            token_budget=200_000,
            max_frames=4096,
        )
        self.assertEqual(len(indices), 3146)
        self.assertEqual(indices[0], 0)
        self.assertEqual(indices[-1], meta["nb_frames"] - 1)
        self.assertLessEqual(estimate, 200_000)
        self.assertEqual((width, height), (96, 64))

    def test_header_one_past_the_decoder_is_not_selected(self) -> None:
        """Header nb_frames is 10129; the decoder emits 10128.

        Planning through the extra sample asks select() for an index that
        does not exist, and the proxy comes back one frame short.
        """
        if not SHORT.is_file() or shutil.which("ffmpeg") is None:
            self.skipTest("short recording is not on this machine")
        meta = probe_video(SHORT)
        self.assertEqual(meta["nb_frames"], 10128)
        self.assertAlmostEqual(meta["source_fps"], 50.0, places=3)
        indices, width, height, estimate = plan_sample(
            nb_frames=meta["nb_frames"],
            source_fps=meta["source_fps"],
            src_w=meta["width"],
            src_h=meta["height"],
            target_fps=2.0,
            token_budget=200_000,
            max_frames=4096,
        )
        self.assertEqual(len(indices), 404)
        self.assertEqual(indices[-1], 10127)
        self.assertEqual((width, height), (320, 160))
        self.assertLessEqual(estimate, 200_000)

    def test_clinical_file_head_is_not_black(self) -> None:
        """The 26-minute recording that OpenCV returned as black frames."""
        if not CLINICAL.is_file():
            self.skipTest("clinical recording is not on this machine")
        with tempfile.TemporaryDirectory() as tmp:
            manifest = sample_video(
                CLINICAL, 2.0, Path(tmp), start_sec=0.0, end_sec=2.0, token_budget=20_000,
            )
            self.assertGreaterEqual(manifest.n_emitted, 2)
            bright = [frame.mean_luma for frame in manifest.frames if frame.mean_luma is not None]
            self.assertTrue(bright)
            self.assertGreater(max(bright), 30.0)


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