yscv-video 0.1.8

use proptest::prelude::*;
use std::sync::atomic::{AtomicUsize, Ordering};

use super::av1_obu::{read_leb128, write_leb128};
use super::frame_pipeline::{FramePipeline, SlotMut, SlotRef, run_pipeline};
use super::h264_encoder::{forward_dct_4x4, rgb8_to_yuv420};
use super::h264_yuv::{nv12_to_rgb8, yuv420_to_rgb8};
use super::mjpeg::decode_mjpeg_to_rgb8;
use super::overlay::draw_rect;

/// Strategy: generate YUV420 data with valid dimensions.
/// Width and height must be even for YUV420.
fn arb_yuv420_data() -> impl Strategy<Value = (Vec<u8>, Vec<u8>, Vec<u8>, usize, usize)> {
    (1usize..=16, 1usize..=16).prop_flat_map(|(half_w, half_h)| {
        let w = half_w * 2;
        let h = half_h * 2;
        let y_len = w * h;
        let uv_len = half_w * half_h;
        (
            proptest::collection::vec(0u8..=255, y_len),
            proptest::collection::vec(0u8..=255, uv_len),
            proptest::collection::vec(0u8..=255, uv_len),
            Just(w),
            Just(h),
        )
    })
}

proptest! {
    // ── YUV420 -> RGB8 range ───────────────────────────────────────────
    #[test]
    fn yuv420_to_rgb8_output_in_range(
        (y, u, v, w, h) in arb_yuv420_data()
    ) {
        let rgb = yuv420_to_rgb8(&y, &u, &v, w, h).expect("yuv420_to_rgb8");
        prop_assert_eq!(rgb.len(), w * h * 3, "output length mismatch");
        // All bytes are inherently in [0, 255] by type,
        // but this validates no panics from overflow in conversion math.
    }

    // ── NV12 -> RGB8 matches YUV420 -> RGB8 ────────────────────────────
    #[test]
    fn nv12_matches_yuv420(
        (y, u, v, w, h) in arb_yuv420_data()
    ) {
        let rgb_yuv420 = yuv420_to_rgb8(&y, &u, &v, w, h).expect("yuv420");

        // Build NV12 interleaved UV plane from separate U/V
        let half_w = w / 2;
        let half_h = h / 2;
        let mut uv_interleaved = vec![0u8; w * half_h];
        for row in 0..half_h {
            for col in 0..half_w {
                uv_interleaved[row * w + col * 2] = u[row * half_w + col];
                uv_interleaved[row * w + col * 2 + 1] = v[row * half_w + col];
            }
        }

        let mut rgb_nv12 = vec![0u8; w * h * 3];
        nv12_to_rgb8(&y, &uv_interleaved, w, h, &mut rgb_nv12).expect("nv12");

        // Both use BT.601 coefficients but may use different fixed-point
        // rounding (SIMD vs scalar). Allow small tolerance.
        let mut max_diff = 0i32;
        for (i, (&a, &b)) in rgb_yuv420.iter().zip(rgb_nv12.iter()).enumerate() {
            let diff = (a as i32 - b as i32).abs();
            if diff > max_diff {
                max_diff = diff;
            }
            prop_assert!(
                diff <= 2,
                "nv12 vs yuv420 mismatch at byte {i}: yuv420={a}, nv12={b}, diff={diff}"
            );
        }
    }

    // ── MAVLink CRC deterministic ──────────────────────────────────────
    #[test]
    fn mavlink_crc_is_deterministic(
        data in proptest::collection::vec(0u8..=255, 1..=64),
        crc_extra in 0u8..=255,
    ) {
        // Compute CRC twice with same input
        let crc1 = mavlink_crc_wrapper(&data, crc_extra);
        let crc2 = mavlink_crc_wrapper(&data, crc_extra);
        prop_assert_eq!(crc1, crc2, "CRC not deterministic");
    }

    // ── LEB128 round-trip ──────────────────────────────────────────────
    #[test]
    fn leb128_roundtrip(value in 0u64..=(1u64 << 56) - 1) {
        let encoded = write_leb128(value);
        let (decoded, bytes_read) = read_leb128(&encoded).expect("leb128 decode");
        prop_assert_eq!(decoded, value, "leb128 roundtrip failed");
        prop_assert_eq!(bytes_read, encoded.len(), "leb128 bytes consumed mismatch");
    }

    // ── (a) H.264 forward DCT bounded output ─────────────────────────
    #[test]
    fn h264_forward_dct_bounded_output(
        vals in proptest::collection::vec(-255i32..=255, 16)
    ) {
        let mut block = [0i32; 16];
        block.copy_from_slice(&vals);

        // Forward transform
        forward_dct_4x4(&mut block);

        // The forward 4x4 integer DCT transform coefficients should be bounded.
        // For 8-bit residuals ([-255,255]), the maximum coefficient magnitude
        // after the H.264 forward transform is bounded by:
        // DC (position 0): sum of all 16 values * max scale = 16 * 255 * 4 = 16320
        // AC: bounded proportionally
        // We verify no overflow and reasonable bounds.
        for (i, &v) in block.iter().enumerate() {
            prop_assert!(
                v.abs() < 100_000,
                "forward DCT coefficient out of range at {}: {}", i, v
            );
        }

        // Verify energy is preserved or increased (DCT is energy-concentrating).
        // The H.264 integer DCT is not orthogonal, but total coefficient
        // energy should be proportional to input energy (within scale factor).
        let input_energy: i64 = vals.iter().map(|&v| (v as i64) * (v as i64)).sum();
        let output_energy: i64 = block.iter().map(|&v| (v as i64) * (v as i64)).sum();

        // Output energy should be non-zero if input is non-zero
        if input_energy > 0 {
            prop_assert!(
                output_energy > 0,
                "forward DCT zeroed out non-zero input: input_energy={}", input_energy
            );
        }
    }

    // ── (b) RGB -> YUV420 -> RGB roundtrip ────────────────────────────
    #[test]
    fn rgb_yuv420_rgb_roundtrip(
        r in 0u8..=255,
        g in 0u8..=255,
        b in 0u8..=255,
    ) {
        // Build a 2x2 RGB block (minimum valid YUV420 size)
        let rgb = vec![r, g, b, r, g, b, r, g, b, r, g, b];
        let yuv = rgb8_to_yuv420(&rgb, 2, 2);

        // YUV420 layout: Y (4 bytes) + U (1 byte) + V (1 byte)
        let y_plane = &yuv[..4];
        let u_plane = &yuv[4..5];
        let v_plane = &yuv[5..6];

        let recovered = yuv420_to_rgb8(y_plane, u_plane, v_plane, 2, 2)
            .expect("yuv420_to_rgb8 roundtrip");

        // Check first pixel (all four should be identical for uniform input).
        // BT.601 studio-range uses integer fixed-point arithmetic with
        // Y_offset=+16, UV_offset=+128, chroma subsampling, and >>8 shifts in
        // both forward and inverse paths. Worst-case accumulated rounding is
        // around 25 (saturated channels like r=254, g=251, b=92). The key
        // invariant is that the roundtrip is bounded, not exact.
        let rr = recovered[0];
        let rg = recovered[1];
        let rb = recovered[2];

        prop_assert!(
            (r as i32 - rr as i32).abs() <= 25,
            "R channel: orig={}, recovered={}", r, rr
        );
        prop_assert!(
            (g as i32 - rg as i32).abs() <= 25,
            "G channel: orig={}, recovered={}", g, rg
        );
        prop_assert!(
            (b as i32 - rb as i32).abs() <= 25,
            "B channel: orig={}, recovered={}", b, rb
        );
    }

    // ── (c) MJPEG decode never panics ─────────────────────────────────
    #[test]
    fn mjpeg_decode_random_never_panics(
        data in proptest::collection::vec(0u8..=255, 0..=1024)
    ) {
        let mut output = Vec::new();
        let result = decode_mjpeg_to_rgb8(&data, &mut output);
        // Random bytes are almost certainly not valid JPEG; should return Err
        prop_assert!(
            result.is_err(),
            "random bytes unexpectedly decoded as valid MJPEG"
        );
    }

    // ── (d) Frame pipeline no data loss ───────────────────────────────
    #[test]
    fn frame_pipeline_no_data_loss(
        num_frames in 5usize..=20,
    ) {
        let pipeline = FramePipeline::new(4, 1024);
        let capture_idx = AtomicUsize::new(0);
        let output_count = AtomicUsize::new(0);

        // Track timestamps to verify ordering
        let timestamps = std::sync::Mutex::new(Vec::new());

        run_pipeline(
            &pipeline,
            |slot: &mut SlotMut<'_>| {
                let i = capture_idx.fetch_add(1, Ordering::Relaxed);
                if i >= num_frames {
                    return false;
                }
                slot.set_timestamp_us(i as u64);
                slot.data_mut()[0] = (i & 0xFF) as u8;
                true
            },
            |_slot: &mut SlotMut<'_>| {
                // Pass-through processing
            },
            |slot: &SlotRef<'_>| {
                timestamps.lock().expect("lock").push(slot.timestamp_us());
                output_count.fetch_add(1, Ordering::Relaxed);
            },
            num_frames,
        );

        let count = output_count.load(Ordering::Relaxed);
        prop_assert_eq!(count, num_frames, "expected {} frames, got {}", num_frames, count);

        // Verify timestamps are sequential
        let ts = timestamps.lock().expect("lock");
        for (i, &t) in ts.iter().enumerate() {
            prop_assert_eq!(t, i as u64, "timestamp mismatch at position {}", i);
        }
    }

    // ── (e) MAVLink parse roundtrip ───────────────────────────────────
    #[test]
    fn mavlink_heartbeat_parse_roundtrip(
        seq in 0u8..=255,
        sysid in 0u8..=255,
        compid in 0u8..=255,
    ) {
        // Build a valid MAVLink v2 HEARTBEAT frame manually
        // HEARTBEAT payload: custom_mode(u32) + type(u8) + autopilot(u8) +
        //                    base_mode(u8) + system_status(u8) + mavlink_version(u8)
        // = 9 bytes
        let payload_len: u8 = 9;
        let payload = [
            0, 0, 0, 0, // custom_mode = 0
            1, // type = 1 (fixed wing)
            3, // autopilot = 3 (ardupilot)
            0, // base_mode = 0
            4, // system_status = 4 (ACTIVE)
            3, // mavlink_version = 3
        ];

        // Header: STX + len + incompat_flags + compat_flags + seq + sysid + compid + msgid(3 bytes)
        let mut frame = vec![
            0xFD,        // STX
            payload_len, // payload length
            0,           // incompat_flags
            0,           // compat_flags
            seq,
            sysid,
            compid,
            0, 0, 0, // msgid = 0 (HEARTBEAT)
        ];
        frame.extend_from_slice(&payload);

        // Compute CRC over bytes 1..10+payload_len, then fold in CRC_EXTRA=50
        let crc_data = &frame[1..10 + payload_len as usize];
        let crc = mavlink_crc_wrapper(crc_data, 50);
        frame.extend_from_slice(&crc.to_le_bytes());

        let result = super::mavlink::parse_mavlink_frame(&frame);
        prop_assert!(result.is_ok(), "parse failed: {:?}", result.err());

        let (msg, consumed) = result.expect("already checked");
        prop_assert_eq!(consumed, frame.len(), "bytes consumed mismatch");

        match msg {
            super::mavlink::MavlinkMessage::Heartbeat {
                autopilot, mav_type, system_status,
            } => {
                prop_assert_eq!(autopilot, 3, "autopilot mismatch");
                prop_assert_eq!(mav_type, 1, "mav_type mismatch");
                prop_assert_eq!(system_status, 4, "system_status mismatch");
            }
            other => {
                prop_assert!(false, "expected Heartbeat, got {:?}", other);
            }
        }
    }

    // ── (f) Overlay draw_rect never panics ────────────────────────────
    #[test]
    fn overlay_draw_rect_never_panics(
        w in 1usize..=200,
        h in 1usize..=200,
        x in 0u32..=250,
        y in 0u32..=250,
        bw in 0u32..=250,
        bh in 0u32..=250,
    ) {
        let mut frame = vec![0u8; w * h * 3];
        // Should never panic regardless of rect position/size relative to frame
        draw_rect(&mut frame, w, h, x, y, bw, bh, 255, 0, 0, 2);
    }

    // ── (g) AV1 LEB128 extended roundtrip (wider range) ───────────────
    #[test]
    fn leb128_extended_roundtrip(value in 0u64..=(1u64 << 28)) {
        let encoded = write_leb128(value);
        prop_assert!(
            encoded.len() <= 8,
            "LEB128 encoding too long: {} bytes for value {value}",
            encoded.len()
        );
        let (decoded, bytes_read) = read_leb128(&encoded).expect("leb128 decode");
        prop_assert_eq!(decoded, value, "leb128 extended roundtrip failed");
        prop_assert_eq!(bytes_read, encoded.len(), "bytes consumed mismatch");

        // Verify the encoding is minimal: the last byte must have its high bit clear
        if let Some(&last) = encoded.last() {
            prop_assert!(
                last & 0x80 == 0,
                "LEB128 last byte should have high bit clear"
            );
        }
    }
}

/// Re-implement MAVLink CRC locally since the function is `fn` (not `pub fn`)
/// in the mavlink module. Uses the same X.25 CRC-16/MCRF4XX algorithm.
fn mavlink_crc_wrapper(data: &[u8], crc_extra: u8) -> u16 {
    let mut crc: u16 = 0xFFFF;
    for &b in data {
        let mut tmp = b ^ (crc as u8);
        tmp ^= tmp.wrapping_shl(4);
        let t16 = tmp as u16;
        crc = (crc >> 8) ^ (t16 << 8) ^ (t16 << 3) ^ (t16 >> 4);
    }
    let mut tmp = crc_extra ^ (crc as u8);
    tmp ^= tmp.wrapping_shl(4);
    let t16 = tmp as u16;
    crc = (crc >> 8) ^ (t16 << 8) ^ (t16 << 3) ^ (t16 >> 4);
    crc
}