from __future__ import annotations

import av
import numpy as np
import rerun as rr

W, H = 320, 240
N_FRAMES = 30

def _h264_encoder():
    """Open an all-intra H.264 encoder with the same settings used elsewhere."""
    container = av.open("/dev/null", "w", format="h264")
    stream = container.add_stream("libx264", rate=30)
    stream.width = W
    stream.height = H
    stream.pix_fmt = "yuv420p"
    stream.max_b_frames = 0
    stream.options = {
        "preset": "ultrafast",
        "crf": "30",
        "g": "1",
        "x264-params": "scenecut=0",
    }
    return container, stream


def encode_h264_all_intra(band_color: tuple[int, int, int]) -> list[bytes]:
    """Encode `N_FRAMES` H.264 I-frames; one packet per frame."""
    container, stream = _h264_encoder()
    out: list[bytes] = []
    try:
        for i in range(N_FRAMES):
            arr = np.full((H, W, 3), 32, dtype=np.uint8)  # dark grey background
            band_x = int((i / N_FRAMES) * W)
            arr[:, band_x : band_x + 24] = band_color
            frame = av.VideoFrame.from_ndarray(arr, format="rgb24").reformat(format="yuv420p")
            for pkt in stream.encode(frame):
                out.append(bytes(pkt))
        for pkt in stream.encode(None):
            out.append(bytes(pkt))
    finally:
        container.close()
    return out


rr.init("rerun_example_video_stream_clear", spawn=True)

ENTITIES = ["/baseline", "/clear", "/empty_sample", "/empty_codec"]

# Anchor the codec at step 0 (non-static) on every entity, before any samples.
rr.set_time("step", sequence=0)
for entity in ENTITIES:
    rr.log(entity, rr.VideoStream.from_fields(codec=rr.components.VideoCodec.H264))

packets_red = encode_h264_all_intra(band_color=(220, 60, 60))
packets_green = encode_h264_all_intra(band_color=(60, 220, 60))

# Burst A (red) starts at step 10 (not 0) so we can see what the viewer renders
# *before* any sample is logged — useful baseline for the "blanking" attempts.
BURST_A_START = 10
for i, pkt in enumerate(packets_red):
    rr.set_time("step", sequence=BURST_A_START + i)
    for entity in ENTITIES:
        rr.log(entity, rr.VideoStream.from_fields(sample=pkt))

# At step 40, do nothing on /baseline.  Then apply each attempt to its entity:
rr.set_time("step", sequence=40)
rr.log("/clear", rr.Clear(recursive=True))
rr.log("/empty_sample", rr.VideoStream.from_fields(sample=[]))
rr.log("/empty_codec", rr.VideoStream.from_fields(codec=[]))

# Burst B (green) on every entity at steps 60..89.
for i, pkt in enumerate(packets_green):
    rr.set_time("step", sequence=60 + i)
    for entity in ENTITIES:
        rr.log(entity, rr.VideoStream.from_fields(sample=pkt))