from i420_nv12 import (
    coded_dim,
    drop_va_pad_i420,
    i420_to_nv12_into,
    nv12_bytes,
    va_visible_height,
)

def test_coded_dim_pads_up():
    assert coded_dim(360) == 368
    assert coded_dim(1080) == 1088
    assert coded_dim(640) == 640
    assert coded_dim(352) == 352


def test_i420_to_nv12_pads_last_row():
    w, h = 640, 360
    ysz = w * h
    uv = (w // 2) * (h // 2)
    src = bytearray(ysz + 2 * uv)
    src[:ysz] = b"\x80" * ysz
    src[(h - 1) * w:ysz] = b"\x90" * w
    src[ysz:ysz + uv] = b"\x40" * uv
    src[ysz + uv:] = b"\xc0" * uv
    cw, ch = 640, 368
    out = bytearray(nv12_bytes(cw, ch))
    n = i420_to_nv12_into(src, w, h, out)
    assert n == nv12_bytes(cw, ch)
    assert out[0] == 0x80
    # cloned pad rows 360-367
    assert out[360 * cw] == 0x90
    assert out[367 * cw] == 0x90
    assert out[cw * ch] == 0x40
    assert out[cw * ch + 1] == 0xc0


def test_drop_va_pad_restores_360():
    w, h = 640, 368
    vis = 360
    ysz = w * h
    uv = (w // 2) * (h // 2)
    src = bytearray(ysz + 2 * uv)
    src[: w * vis] = b"\x77" * (w * vis)
    src[w * vis:ysz] = b"\x11" * (w * 8)
    out, ow, oh = drop_va_pad_i420(src, w, h)
    assert (ow, oh) == (w, vis)
    assert len(out) == w * vis * 3 // 2
    assert out[0] == 0x77
    assert 0x11 not in out[: w * vis]


def test_va_visible_height():
    assert va_visible_height(640, 368) == 360
    assert va_visible_height(1920, 1088) == 1080
    assert va_visible_height(640, 360) == 360
    assert va_visible_height(640, 352) == 352