"""16-align I420→NV12 for Intel VA (NV12 only). Pad up, never crop."""
from __future__ import annotations

import numpy as np


def coded_dim(n: int) -> int:
    n = int(n)
    if n <= 0:
        return 0
    return (n + 15) // 16 * 16


def nv12_bytes(width: int, height: int) -> int:
    return int(width) * int(height) * 3 // 2


def va_visible_height(width: int, height: int) -> int:
    """Drop the 8-row VA pad when the remainder is 16:9 (360/1080)."""
    w, h = int(width), int(height)
    vis = h - 8
    if w > 0 and vis > 0 and vis * 16 == w * 9:
        return vis
    return h


def drop_va_pad_i420(src, width: int, height: int):
    """Crop coded pad off I420. Returns (payload, w, visible_h)."""
    w, h = int(width), int(height)
    vis = va_visible_height(w, h)
    if vis == h:
        return src, w, h
    ysz = w * h
    uvw, uvh = w // 2, h // 2
    uvsz = uvw * uvh
    raw = np.frombuffer(src, dtype=np.uint8, count=ysz + 2 * uvsz)
    y = raw[:ysz].reshape(h, w)[:vis, :]
    u = raw[ysz:ysz + uvsz].reshape(uvh, uvw)[: vis // 2, :]
    v = raw[ysz + uvsz:ysz + 2 * uvsz].reshape(uvh, uvw)[: vis // 2, :]
    out = np.empty(w * vis * 3 // 2, dtype=np.uint8)
    out[: w * vis] = y.reshape(-1)
    out[w * vis:w * vis + (w // 2) * (vis // 2)] = u.reshape(-1)
    out[w * vis + (w // 2) * (vis // 2):] = v.reshape(-1)
    return out.tobytes(), w, vis


def i420_to_nv12_into(src, width: int, height: int, out: bytearray) -> int:
    """Pad to 16-aligned WxH and write NV12 into out. Returns nbytes."""
    w, h = int(width), int(height)
    cw, ch = coded_dim(w), coded_dim(h)
    need = nv12_bytes(cw, ch)
    if len(out) < need:
        raise ValueError("NV12 dest too small")
    ysz = w * h
    uvw, uvh = w // 2, h // 2
    uvsz = uvw * uvh
    raw = np.frombuffer(src, dtype=np.uint8, count=ysz + 2 * uvsz)
    y = raw[:ysz].reshape(h, w)
    u = raw[ysz:ysz + uvsz].reshape(uvh, uvw)
    v = raw[ysz + uvsz:ysz + 2 * uvsz].reshape(uvh, uvw)
    nv = np.frombuffer(out, dtype=np.uint8, count=need).reshape(ch + ch // 2, cw)
    nv[:h, :w] = y
    if ch > h:
        nv[h:ch, :w] = y[-1]
    if cw > w:
        nv[:ch, w:cw] = nv[:ch, w - 1:w]
    uv = nv[ch:, :].reshape(ch // 2, cw // 2, 2)
    uh, uw = h // 2, w // 2
    uv[:uh, :uw, 0] = u
    uv[:uh, :uw, 1] = v
    if ch // 2 > uh:
        uv[uh:, :uw, 0] = u[-1]
        uv[uh:, :uw, 1] = v[-1]
    if cw // 2 > uw:
        uv[:, uw:, :] = uv[:, uw - 1:uw, :]
    return need