from capture_va import (
    crop_scale_i420_drop_bottom,
    jpeg_mcu_bottom_crop,
    needs_jpeg_mcu_fix,
    repeat_i420_bottom_from_last_good,
)


def _i420(width, height, y=16, u=128, v=128):
    y_size = width * height
    uv = (width // 2) * (height // 2)
    buf = bytearray(y_size + 2 * uv)
    buf[:y_size] = bytes([y]) * y_size
    buf[y_size:y_size + uv] = bytes([u]) * uv
    buf[y_size + uv:] = bytes([v]) * uv
    return buf


def test_repeat_clones_last_good_y_and_chroma_over_black_pad():
    w, h, crop = 640, 360, 8
    buf = _i420(w, h, y=16, u=128, v=128)
    last_good = h - crop - 1
    y_stride = w
    # last good Y row = 90; last good UV row = 40 (chroma 4:2:0)
    src = last_good * y_stride
    buf[src:src + y_stride] = bytes([90]) * y_stride
    uv_w = w // 2
    uv_h = h // 2
    uv_crop = crop // 2
    last_uv = uv_h - uv_crop - 1
    y_size = w * h
    u0 = y_size + last_uv * uv_w
    v0 = y_size + uv_w * uv_h + last_uv * uv_w
    buf[u0:u0 + uv_w] = bytes([40]) * uv_w
    buf[v0:v0 + uv_w] = bytes([200]) * uv_w

    repeat_i420_bottom_from_last_good(buf, w, h, crop)

    for r in range(h - crop, h):
        row = buf[r * y_stride:(r + 1) * y_stride]
        assert set(row) == {90}, r
    for r in range(uv_h - uv_crop, uv_h):
        u = buf[y_size + r * uv_w:y_size + (r + 1) * uv_w]
        v = buf[y_size + uv_w * uv_h + r * uv_w:
                y_size + uv_w * uv_h + (r + 1) * uv_w]
        assert set(u) == {40}, r
        assert set(v) == {200}, r
    # rows above the pad stay put
    mid = (h // 2) * y_stride
    assert set(buf[mid:mid + y_stride]) == {16}


def test_360p_crops_partial_mcu_and_last_full_mcu():
    """360 % 16 = 8 leftover; last full 16-row MCU is also garbage on 07/14/15."""
    assert jpeg_mcu_bottom_crop(360) == 24
    assert jpeg_mcu_bottom_crop(480) == 0
    assert jpeg_mcu_bottom_crop(352) == 0
    assert jpeg_mcu_bottom_crop(8) == 0
    buf = _i420(640, 480, y=77)
    before = bytes(buf)
    repeat_i420_bottom_from_last_good(buf, 640, 480, jpeg_mcu_bottom_crop(480))
    assert bytes(buf) == before


def test_mcu_fix_cams_off():
    assert not needs_jpeg_mcu_fix("cam-07")
    assert not needs_jpeg_mcu_fix("cam-12")
    assert not needs_jpeg_mcu_fix("cam-14")
    assert not needs_jpeg_mcu_fix("cam-15")
    assert not needs_jpeg_mcu_fix("cam-01")
    assert not needs_jpeg_mcu_fix("cam-20")


def test_crop_scale_drops_garbage_without_cloned_stripe():
    w, h = 640, 360
    crop = 24
    buf = _i420(w, h, y=40, u=100, v=140)
    y_size = w * h
    # garbage last 24 Y rows
    buf[(h - crop) * w:y_size] = bytes([250]) * (crop * w)
    out = crop_scale_i420_drop_bottom(buf, w, h, crop)
    assert len(out) == len(buf)
    y = memoryview(out)[:y_size]
    last = bytes(y[(h - 8) * w:h * w])
    assert last.count(250) == 0
    assert set(last) == {16}


def test_live_path_clones_not_studio_black():
    w, h, crop = 640, 360, 24
    buf = _i420(w, h, y=40, u=100, v=140)
    last_good = h - crop - 1
    buf[last_good * w:(last_good + 1) * w] = bytes([90]) * w
    buf[(h - crop) * w:h * w] = bytes([250]) * (crop * w)
    repeat_i420_bottom_from_last_good(buf, w, h, crop)
    visible = bytes(buf[(h - crop) * w:h * w])
    assert set(visible) == {90}
    assert 16 not in visible
    assert 250 not in visible