//! Color-space conversion for the H.264 encode path (Task 7).
//!
//! Screen capture produces BGRA (4 bytes/pixel, B,G,R,A order — the DXGI/GDI
//! native layout). Media Foundation hardware H.264 encoders want NV12: a full-
//! resolution 8-bit Y (luma) plane followed by an interleaved half-resolution
//! U/V (chroma) plane. This module does that conversion in software.
//!
//! NV12 memory layout for a `width x height` frame (width/height assumed even):
//! - Y plane:  `width * height` bytes, row-major.
//! - UV plane: `width * (height / 2)` bytes — for each 2x2 luma block one
//!   (U, V) pair, so the plane is `(width/2)` (U,V) pairs per row over
//!   `height/2` rows, i.e. `width` bytes per chroma row.
//!
//! Total size = `width * height * 3 / 2`.
//!
//! The coefficients are BT.601 "studio swing" (limited range, 16..235 luma),
//! which is what MF H.264 encoders expect by default. Chroma is computed by
//! averaging the 2x2 BGRA block before conversion (box downsample) to reduce
//! aliasing.

/// Size in bytes of an NV12 buffer for `width` x `height` (both even).
#[inline]
pub fn nv12_size(width: u32, height: u32) -> usize {
    (width as usize * height as usize) * 3 / 2
}

/// BT.601 limited-range luma from 8-bit R,G,B.
#[inline]
fn rgb_to_y(r: i32, g: i32, b: i32) -> u8 {
    // Y = 16 + (65.481*R + 128.553*G + 24.966*B) / 255, fixed-point.
    // Using the common integer approximation:
    //   Y = ((66*R + 129*G + 25*B + 128) >> 8) + 16
    let y = ((66 * r + 129 * g + 25 * b + 128) >> 8) + 16;
    y.clamp(0, 255) as u8
}

/// BT.601 limited-range Cb (U) from 8-bit R,G,B.
#[inline]
fn rgb_to_u(r: i32, g: i32, b: i32) -> u8 {
    let u = ((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128;
    u.clamp(0, 255) as u8
}

/// BT.601 limited-range Cr (V) from 8-bit R,G,B.
#[inline]
fn rgb_to_v(r: i32, g: i32, b: i32) -> u8 {
    let v = ((112 * r - 94 * g - 18 * b + 128) >> 8) + 128;
    v.clamp(0, 255) as u8
}

/// Convert a tightly-packed BGRA frame into NV12, writing into `out`.
///
/// `bgra` must be at least `width * height * 4` bytes; `out` must be at least
/// `nv12_size(width, height)` bytes. `width` and `height` MUST be even (H.264
/// 4:2:0 requires even dimensions — the caller pads odd capture sizes). Returns
/// an error rather than panicking on a short buffer or odd dimension so the
/// encoder can fall back to raw.
pub fn bgra_to_nv12(
    bgra: &[u8],
    width: u32,
    height: u32,
    out: &mut [u8],
) -> Result<(), ColorConvertError> {
    if width == 0 || height == 0 {
        return Err(ColorConvertError::ZeroDimension);
    }
    if !width.is_multiple_of(2) || !height.is_multiple_of(2) {
        return Err(ColorConvertError::OddDimension { width, height });
    }

    let w = width as usize;
    let h = height as usize;
    let expected_src = w * h * 4;
    if bgra.len() < expected_src {
        return Err(ColorConvertError::SrcTooSmall {
            got: bgra.len(),
            need: expected_src,
        });
    }
    let need_out = nv12_size(width, height);
    if out.len() < need_out {
        return Err(ColorConvertError::DstTooSmall {
            got: out.len(),
            need: need_out,
        });
    }

    let (y_plane, uv_plane) = out.split_at_mut(w * h);

    // Luma: one sample per pixel.
    for row in 0..h {
        let src_row = row * w * 4;
        let dst_row = row * w;
        for col in 0..w {
            let px = src_row + col * 4;
            // BGRA order.
            let b = bgra[px] as i32;
            let g = bgra[px + 1] as i32;
            let r = bgra[px + 2] as i32;
            y_plane[dst_row + col] = rgb_to_y(r, g, b);
        }
    }

    // Chroma: one (U,V) pair per 2x2 block, box-averaged.
    let chroma_rows = h / 2;
    let chroma_cols = w / 2;
    for cy in 0..chroma_rows {
        for cx in 0..chroma_cols {
            let x0 = cx * 2;
            let y0 = cy * 2;

            let mut r_sum = 0i32;
            let mut g_sum = 0i32;
            let mut b_sum = 0i32;
            for dy in 0..2 {
                for dx in 0..2 {
                    let px = ((y0 + dy) * w + (x0 + dx)) * 4;
                    b_sum += bgra[px] as i32;
                    g_sum += bgra[px + 1] as i32;
                    r_sum += bgra[px + 2] as i32;
                }
            }
            let r = r_sum / 4;
            let g = g_sum / 4;
            let b = b_sum / 4;

            let uv_idx = (cy * chroma_cols + cx) * 2;
            uv_plane[uv_idx] = rgb_to_u(r, g, b);
            uv_plane[uv_idx + 1] = rgb_to_v(r, g, b);
        }
    }

    Ok(())
}

/// Errors from BGRA->NV12 conversion. Surfaced (not panicked) so the H.264
/// encoder can downgrade to raw.
#[derive(Debug, thiserror::Error)]
pub enum ColorConvertError {
    #[error("frame dimension is zero")]
    ZeroDimension,
    #[error("NV12 requires even dimensions, got {width}x{height}")]
    OddDimension { width: u32, height: u32 },
    #[error("source BGRA buffer too small: {got} < {need}")]
    SrcTooSmall { got: usize, need: usize },
    #[error("destination NV12 buffer too small: {got} < {need}")]
    DstTooSmall { got: usize, need: usize },
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn nv12_size_is_3half() {
        assert_eq!(nv12_size(2, 2), 6);
        assert_eq!(nv12_size(4, 4), 24);
        assert_eq!(nv12_size(1920, 1080), 1920 * 1080 * 3 / 2);
    }

    #[test]
    fn rejects_odd_dimensions() {
        let bgra = vec![0u8; 3 * 3 * 4];
        let mut out = vec![0u8; nv12_size(4, 4)];
        assert!(matches!(
            bgra_to_nv12(&bgra, 3, 2, &mut out),
            Err(ColorConvertError::OddDimension { .. })
        ));
        assert!(matches!(
            bgra_to_nv12(&bgra, 2, 3, &mut out),
            Err(ColorConvertError::OddDimension { .. })
        ));
    }

    #[test]
    fn rejects_short_source() {
        let bgra = vec![0u8; 4]; // way too small for 2x2
        let mut out = vec![0u8; nv12_size(2, 2)];
        assert!(matches!(
            bgra_to_nv12(&bgra, 2, 2, &mut out),
            Err(ColorConvertError::SrcTooSmall { .. })
        ));
    }

    #[test]
    fn rejects_short_dest() {
        let bgra = vec![0u8; 2 * 2 * 4];
        let mut out = vec![0u8; 1];
        assert!(matches!(
            bgra_to_nv12(&bgra, 2, 2, &mut out),
            Err(ColorConvertError::DstTooSmall { .. })
        ));
    }

    /// A pure-black BGRA frame -> Y = 16 (limited-range black), U = V = 128.
    #[test]
    fn black_frame_maps_to_limited_range_black() {
        let bgra = vec![0u8; 4 * 4 * 4]; // all zero => black, alpha 0
        let mut out = vec![0u8; nv12_size(4, 4)];
        bgra_to_nv12(&bgra, 4, 4, &mut out).unwrap();

        // Y plane (first 16 bytes) all 16.
        for &y in &out[..16] {
            assert_eq!(y, 16, "black luma must be 16 (limited range)");
        }
        // UV plane all 128 (neutral chroma).
        for &c in &out[16..] {
            assert_eq!(c, 128, "black chroma must be neutral 128");
        }
    }

    /// A pure-white BGRA frame -> Y = 235 (limited-range white), U = V = 128.
    #[test]
    fn white_frame_maps_to_limited_range_white() {
        // B=255, G=255, R=255, A=255 for every pixel.
        let bgra = vec![255u8; 2 * 2 * 4];
        let mut out = vec![0u8; nv12_size(2, 2)];
        bgra_to_nv12(&bgra, 2, 2, &mut out).unwrap();

        // Y = ((66+129+25)*255 + 128) >> 8 + 16 = 235.
        for &y in &out[..4] {
            assert_eq!(y, 235, "white luma must be 235 (limited range)");
        }
        // Neutral chroma for a gray/white pixel.
        assert_eq!(out[4], 128);
        assert_eq!(out[5], 128);
    }

    /// A pure-red frame: luma below mid, V (Cr) well above 128, U (Cb) below 128.
    #[test]
    fn red_frame_has_high_cr_low_cb() {
        // BGRA red: B=0, G=0, R=255, A=255.
        let mut bgra = vec![0u8; 2 * 2 * 4];
        for px in bgra.chunks_mut(4) {
            px[0] = 0; // B
            px[1] = 0; // G
            px[2] = 255; // R
            px[3] = 255; // A
        }
        let mut out = vec![0u8; nv12_size(2, 2)];
        bgra_to_nv12(&bgra, 2, 2, &mut out).unwrap();

        let u = out[4];
        let v = out[5];
        assert!(v > 200, "red must have high Cr (V), got {v}");
        assert!(u < 128, "red must have Cb (U) below neutral, got {u}");
    }

    /// Conversion fills the whole NV12 buffer (no leftover zeros where data is
    /// expected) for a non-trivial gradient — a sanity check on plane indexing.
    #[test]
    fn plane_indexing_covers_full_buffer() {
        let w = 8u32;
        let h = 8u32;
        let mut bgra = vec![0u8; (w * h * 4) as usize];
        for (i, px) in bgra.chunks_mut(4).enumerate() {
            let v = (i % 256) as u8;
            px[0] = v;
            px[1] = v;
            px[2] = v;
            px[3] = 255;
        }
        let mut out = vec![0xAAu8; nv12_size(w, h)];
        bgra_to_nv12(&bgra, w, h, &mut out).unwrap();
        // Y plane should be fully written (gray ramp -> non-constant).
        let y_plane = &out[..(w * h) as usize];
        assert!(y_plane.windows(2).any(|p| p[0] != p[1]), "Y plane varies");
    }
}
