"""Tests for OverlayRenderer.draw_hazard_freeze and freeze-frame injection."""
import numpy as np
from badas.visualization.overlay import OverlayRenderer


def _blank(h=720, w=1280):
    return np.zeros((h, w, 3), dtype=np.uint8)


# ---------------------------------------------------------------------------
# draw_hazard_freeze — output contract
# ---------------------------------------------------------------------------

def test_draw_hazard_freeze_returns_same_shape():
    """Output frame has same shape as input."""
    r = OverlayRenderer()
    frame = _blank()
    pred = {'probability': 0.9, 'raw_probability': 0.9, 'timestamp': 5.0, 'bbox': (10, 10, 100, 100)}
    out = r.draw_hazard_freeze(frame, pred)
    assert out.shape == frame.shape


def test_draw_hazard_freeze_modifies_frame():
    """Output is not identical to input (overlay was drawn)."""
    r = OverlayRenderer()
    frame = _blank()
    pred = {'probability': 0.9, 'raw_probability': 0.9, 'timestamp': 5.0, 'bbox': None}
    out = r.draw_hazard_freeze(frame, pred)
    assert not np.array_equal(out, frame)


def test_draw_hazard_freeze_no_bbox_no_crash():
    """Works when bbox is None (no description either)."""
    r = OverlayRenderer()
    frame = _blank()
    pred = {'probability': 0.85, 'raw_probability': 0.85, 'timestamp': 3.0, 'bbox': None}
    out = r.draw_hazard_freeze(frame, pred)
    assert out.shape == frame.shape


def test_draw_hazard_freeze_with_description_no_crash():
    """Works when description is present but bbox is None."""
    r = OverlayRenderer()
    frame = _blank()
    pred = {
        'probability': 0.9, 'raw_probability': 0.9, 'timestamp': 5.0,
        'bbox': None, 'description': 'A truck from the left — brake hard!',
    }
    out = r.draw_hazard_freeze(frame, pred)
    assert out.shape == frame.shape


def test_draw_hazard_freeze_with_bbox_and_description_no_crash():
    """Works with both bbox and description."""
    r = OverlayRenderer()
    frame = _blank()
    pred = {
        'probability': 0.95, 'raw_probability': 0.95, 'timestamp': 7.0,
        'bbox': (30, 40, 150, 200), 'description': 'Cyclist on the right — slow down!',
    }
    out = r.draw_hazard_freeze(frame, pred)
    assert out.shape == frame.shape


# ---------------------------------------------------------------------------
# _is_hazard_event helper
# ---------------------------------------------------------------------------

from badas.visualization.overlay import _is_hazard_event  # module-level helper


def test_is_hazard_event_bbox():
    assert _is_hazard_event({'bbox': (0, 0, 10, 10)}) is True


def test_is_hazard_event_description():
    assert _is_hazard_event({'description': 'A truck!'}) is True


def test_is_hazard_event_both():
    assert _is_hazard_event({'bbox': (0, 0, 10, 10), 'description': 'X'}) is True


def test_is_hazard_event_neither():
    assert _is_hazard_event({'bbox': None}) is False
    assert _is_hazard_event({}) is False
    assert _is_hazard_event({'description': ''}) is False
    assert _is_hazard_event({'description': None}) is False


# ---------------------------------------------------------------------------
# Freeze-frame injection: VideoOverlayProcessor with mock video
# ---------------------------------------------------------------------------

import tempfile, os
import subprocess

def _make_test_video(path, n_frames=30, fps=10, w=128, h=72):
    """Create a tiny solid-colour video with ffmpeg."""
    cmd = [
        'ffmpeg', '-y',
        '-f', 'lavfi', '-i', f'color=c=blue:size={w}x{h}:rate={fps}:duration={n_frames/fps}',
        '-c:v', 'libx264', '-pix_fmt', 'yuv420p', path
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def test_freeze_extends_output_duration():
    """
    With pause_on_hazard=True and one hazard event, output video is longer
    than input by approximately pause_duration seconds.
    """
    import cv2
    from badas.visualization.overlay import VideoOverlayProcessor

    with tempfile.TemporaryDirectory() as tmp:
        in_path = os.path.join(tmp, 'in.mp4')
        out_path = os.path.join(tmp, 'out.mp4')
        _make_test_video(in_path, n_frames=30, fps=10)

        predictions = [
            {
                'timestamp': 1.0, 'probability': 0.9, 'raw_probability': 0.9,
                'bbox': (10, 10, 100, 100), 'description': 'Test hazard',
            },
            {'timestamp': 2.0, 'probability': 0.3, 'raw_probability': 0.3},
        ]

        proc = VideoOverlayProcessor()
        proc.create_overlay_video(
            in_path, predictions, out_path, quiet=True,
            pause_on_hazard=True, pause_duration=2.0,
        )

        cap_in = cv2.VideoCapture(in_path)
        cap_out = cv2.VideoCapture(out_path)
        in_frames = int(cap_in.get(cv2.CAP_PROP_FRAME_COUNT))
        out_frames = int(cap_out.get(cv2.CAP_PROP_FRAME_COUNT))
        cap_in.release()
        cap_out.release()

        extra = int(10 * 2.0)  # 10fps × 2.0s pause
        assert out_frames == in_frames + extra, (
            f"Expected {in_frames + extra} frames, got {out_frames}"
        )


def test_freeze_off_does_not_extend():
    """With pause_on_hazard=False, output frame count equals input."""
    import cv2
    from badas.visualization.overlay import VideoOverlayProcessor

    with tempfile.TemporaryDirectory() as tmp:
        in_path = os.path.join(tmp, 'in.mp4')
        out_path = os.path.join(tmp, 'out.mp4')
        _make_test_video(in_path, n_frames=30, fps=10)

        predictions = [
            {
                'timestamp': 1.0, 'probability': 0.9, 'raw_probability': 0.9,
                'bbox': (10, 10, 100, 100),
            },
        ]

        proc = VideoOverlayProcessor()
        proc.create_overlay_video(
            in_path, predictions, out_path, quiet=True,
            pause_on_hazard=False,
        )

        cap_in = cv2.VideoCapture(in_path)
        cap_out = cv2.VideoCapture(out_path)
        in_frames = int(cap_in.get(cv2.CAP_PROP_FRAME_COUNT))
        out_frames = int(cap_out.get(cv2.CAP_PROP_FRAME_COUNT))
        cap_in.release()
        cap_out.release()

        assert out_frames == in_frames


def test_freeze_fires_once_per_event():
    """Each hazard event triggers exactly one pause, not one per frame."""
    import cv2
    from badas.visualization.overlay import VideoOverlayProcessor

    with tempfile.TemporaryDirectory() as tmp:
        in_path = os.path.join(tmp, 'in.mp4')
        out_path = os.path.join(tmp, 'out.mp4')
        _make_test_video(in_path, n_frames=50, fps=10)

        # Two distinct hazard events
        predictions = [
            {'timestamp': 1.0, 'probability': 0.9, 'raw_probability': 0.9, 'bbox': (0, 0, 64, 64)},
            {'timestamp': 2.0, 'probability': 0.5, 'raw_probability': 0.5},
            {'timestamp': 3.0, 'probability': 0.85, 'raw_probability': 0.85, 'bbox': (10, 10, 80, 80)},
        ]

        proc = VideoOverlayProcessor()
        proc.create_overlay_video(
            in_path, predictions, out_path, quiet=True,
            pause_on_hazard=True, pause_duration=1.0,
        )

        cap_in = cv2.VideoCapture(in_path)
        cap_out = cv2.VideoCapture(out_path)
        in_frames = int(cap_in.get(cv2.CAP_PROP_FRAME_COUNT))
        out_frames = int(cap_out.get(cv2.CAP_PROP_FRAME_COUNT))
        cap_in.release()
        cap_out.release()

        extra = 2 * int(10 * 1.0)  # 2 events × 10fps × 1.0s
        assert out_frames == in_frames + extra
