"""Потеря реального H.264 кадра и восстановление декодера на следующем IDR."""

from __future__ import annotations

import av
import numpy as np

from protocol.video_stream_transport import FrameType, VideoFrame
from tools.video002_simulation import simulate
from tools.video_compare import new_h264_encoder


def test_h264_recovers_after_lost_keyframe_and_delta() -> None:
    fps = 5
    encoder, _ = new_h264_encoder(64, 48, fps, 160_000)
    frames = []
    for index in range(15):
        pixels = np.zeros((48, 64, 3), dtype=np.uint8)
        pixels[:, :, 1] = (index * 13) % 255
        pixels[10:22, index * 3:index * 3 + 12, 0] = 230
        raw = av.VideoFrame.from_ndarray(pixels, format="rgb24").reformat(format="yuv420p")
        raw.pts = index
        packets = encoder.encode(raw)
        assert len(packets) == 1
        packet = packets[0]
        frames.append(VideoFrame(17, index,
                                 FrameType.KEYFRAME if packet.is_keyframe
                                 else FrameType.DELTA_FRAME,
                                 index * 200, bytes(packet)))
    assert not encoder.encode(None)
    assert [item.frame_id for item in frames if item.frame_type == FrameType.KEYFRAME] == [0, 5, 10]
    link = ("SOFTWARE_2MBPS", None, 2_000_000, 1024, 0)
    clean = simulate(frames, fps, link, ("random", 0))
    key_lost = simulate(frames, fps, link, ("keyframe", 1))
    delta_lost = simulate(frames, fps, link, ("delta", 1))
    assert clean["displayed_frame_ids"] == list(range(15))
    assert key_lost["displayed_frame_ids"] == list(range(5)) + list(range(10, 15))
    assert delta_lost["displayed_frame_ids"] == list(range(6)) + list(range(10, 15))
    assert key_lost["keyframe_sha256_by_id"]["10"] == clean["keyframe_sha256_by_id"]["10"]
    assert key_lost["unique_decoded_images"] == key_lost["displayed_frames"]
    assert key_lost["no_old_frame_replay"] and delta_lost["no_old_frame_replay"]