tests/test_video.py

from __future__ import annotations

from pathlib import Path
from typing import Any

from tiny_trigger.models import ActionEvent, Detection
from tiny_trigger.video import _create_browser_mp4_writer, process_video, render_automation_video


class FakeDetector:
    class_names = ["cat"]

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        assert image_size == 960
        assert max_detections == 20
        return [
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label="cat",
                confidence=0.99,
                bbox_xyxy=(2.0, 2.0, 12.0, 12.0),
                bbox_xyxy_norm=(0.1, 0.1, 0.6, 0.6),
            )
        ]


class TrackedDetector:
    class_names = ["cat"]

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        offset = frame_index * 0.01
        return [
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label="cat",
                confidence=0.99,
                bbox_xyxy=(2.0 + frame_index, 2.0, 12.0 + frame_index, 12.0),
                bbox_xyxy_norm=(0.1 + offset, 0.1, 0.2 + offset, 0.2),
                track_id=7,
            )
        ]


class DuplicateDetector:
    class_names = ["cat"]

    def detect(
        self,
        frame: Any,
        *,
        frame_index: int,
        timestamp_sec: float,
        confidence: float,
        image_size: int | None = None,
        max_detections: int | None = None,
    ) -> list[Detection]:
        return [
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label="cat",
                confidence=0.62,
                bbox_xyxy=(2.0, 2.0, 16.0, 16.0),
                bbox_xyxy_norm=(0.1, 0.1, 0.5, 0.5),
            ),
            Detection(
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
                label="cat",
                confidence=0.91,
                bbox_xyxy=(3.0, 3.0, 17.0, 17.0),
                bbox_xyxy_norm=(0.11, 0.11, 0.51, 0.51),
            ),
        ]


def test_process_video_with_fake_detector(tmp_path: Path) -> None:
    cv2 = __import__("cv2")
    video_path = _make_video(tmp_path)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=2,
        max_frames=2,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    assert Path(result.output_video_path).exists()
    assert result.processed_frames == 2
    assert [item.frame_index for item in result.detections] == [0, 2]
    assert result.frame_stride == 2
    assert result.sample_interval_sec is None
    capture = cv2.VideoCapture(result.output_video_path)
    try:
        assert capture.get(cv2.CAP_PROP_FRAME_COUNT) == 4
        assert capture.get(cv2.CAP_PROP_FPS) == 10.0
    finally:
        capture.release()


def test_process_video_does_not_synthesize_track_ids(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=10.0, frames=4)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=1,
        max_frames=3,
        detector=FakeDetector(),
        image_size=960,
        max_detections=20,
        output_dir=str(tmp_path),
    )

    assert [item.track_id for item in result.detections] == [None, None, None]


def test_process_video_preserves_detector_track_ids(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=10.0, frames=4)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=1,
        max_frames=3,
        detector=TrackedDetector(),
        output_dir=str(tmp_path),
    )

    assert [item.track_id for item in result.detections] == [7, 7, 7]


def test_process_video_suppresses_duplicate_same_label_boxes(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=10.0, frames=2)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=1,
        max_frames=1,
        detector=DuplicateDetector(),
        output_dir=str(tmp_path),
    )

    assert len(result.detections) == 1
    assert result.detections[0].confidence == 0.91


def test_process_video_samples_once_per_second(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=30.0, frames=95)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=2,
        sample_interval_sec=1.0,
        max_frames=3,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    assert result.frame_stride == 30
    assert result.sample_interval_sec == 1.0
    assert [item.frame_index for item in result.detections] == [0, 30, 60]


def test_process_video_samples_half_second_intervals(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=10.0, frames=16)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        sample_interval_sec=0.5,
        max_frames=3,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    assert result.frame_stride == 5
    assert [item.frame_index for item in result.detections] == [0, 5, 10]


def test_render_automation_video_with_fired_event(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path)
    detections = [
        Detection(
            frame_index=0,
            timestamp_sec=0.0,
            label="cat",
            confidence=0.99,
            bbox_xyxy=(2.0, 2.0, 12.0, 12.0),
            bbox_xyxy_norm=(0.1, 0.1, 0.6, 0.6),
        )
    ]
    events = [
        ActionEvent(
            rule="cat-rule",
            action="feed cat",
            type="simulate",
            frame_index=0,
            timestamp_sec=0.0,
            status="simulated",
        )
    ]

    output_path = render_automation_video(
        source_video_path=str(video_path),
        detections=detections,
        events=events,
        frame_stride=2,
        max_frames=2,
        output_dir=str(tmp_path),
    )

    assert Path(output_path).exists()
    assert output_path.endswith("-automated.mp4")


def test_render_automation_video_uses_computed_stride(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=30.0, frames=95)
    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        sample_interval_sec=1.0,
        max_frames=3,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    output_path = render_automation_video(
        source_video_path=str(video_path),
        detections=result.detections,
        events=[],
        frame_stride=result.frame_stride,
        max_frames=3,
        output_dir=str(tmp_path),
    )

    assert Path(output_path).exists()


def test_process_video_writes_full_motion_clip_with_sampled_overlays(tmp_path: Path) -> None:
    cv2 = __import__("cv2")
    video_path = _make_video(tmp_path, fps=5.0, frames=15)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=5,
        max_frames=3,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    capture = cv2.VideoCapture(result.output_video_path)
    try:
        assert capture.get(cv2.CAP_PROP_FPS) == 5.0
        assert capture.get(cv2.CAP_PROP_FRAME_COUNT) == 15
    finally:
        capture.release()


def test_process_video_writes_faststart_mp4(tmp_path: Path) -> None:
    video_path = _make_video(tmp_path, fps=5.0, frames=6)

    result = process_video(
        video_path=str(video_path),
        class_prompt="cat",
        frame_stride=2,
        max_frames=2,
        image_size=960,
        max_detections=20,
        detector=FakeDetector(),
        output_dir=str(tmp_path),
    )

    data = Path(result.output_video_path).read_bytes()
    assert data.find(b"moov") < data.find(b"mdat")


def test_browser_mp4_writer_uses_mp4v_only(monkeypatch, tmp_path: Path) -> None:
    calls: list[str] = []

    class Writer:
        def isOpened(self) -> bool:
            return True

        def release(self) -> None:
            return None

    class CV2:
        @staticmethod
        def VideoWriter_fourcc(*codec):
            calls.append("".join(codec))
            return 1234

        @staticmethod
        def VideoWriter(path, fourcc, fps, frame_size):
            return Writer()

    monkeypatch.setitem(__import__("sys").modules, "cv2", CV2)

    writer = _create_browser_mp4_writer(tmp_path / "out.mp4", 8.0, (32, 32))

    assert writer is not None
    assert calls == ["mp4v"]


def _make_video(tmp_path: Path, *, fps: float = 10.0, frames: int = 4) -> Path:
    cv2 = __import__("cv2")
    video_path = tmp_path / "input.mp4"
    writer = cv2.VideoWriter(str(video_path), cv2.VideoWriter_fourcc(*"mp4v"), fps, (32, 32))
    for index in range(frames):
        frame = __import__("numpy").zeros((32, 32, 3), dtype="uint8")
        frame[:] = (index * 20) % 256
        writer.write(frame)
    writer.release()
    return video_path