//! Encode correctness + honest timing on synthetic MOVING content (independent
//! of whether the live desktop is active). Encodes N frames, decodes them back
//! with the real viewer decoder, and asserts the roundtrip reproduces the right
//! dimensions — proving the Annex-B the host emits is valid for the viewer.
//!
//! Usage: netdesk-host-encbench [sw|mf]   (default sw = openh264 software)

use std::time::Instant;

fn make_frame(w: usize, h: usize, t: usize) -> Vec<u8> {
    let mut buf = vec![0u8; w * h * 4];
    let bx = (t * 13) % (w.saturating_sub(120).max(1));
    let by = (t * 7) % (h.saturating_sub(120).max(1));
    for y in 0..h {
        for x in 0..w {
            let i = (y * w + x) * 4;
            let inbox = x >= bx && x < bx + 120 && y >= by && y < by + 120;
            if inbox {
                buf[i] = 255;
                buf[i + 1] = 255;
                buf[i + 2] = 255;
                buf[i + 3] = 255;
            } else {
                buf[i] = ((x + t * 3) & 0xff) as u8;
                buf[i + 1] = ((y + t * 2) & 0xff) as u8;
                buf[i + 2] = ((x + y + t) & 0xff) as u8;
                buf[i + 3] = 255;
            }
        }
    }
    buf
}

/// Mean (R,G,B) of a decoded row.
fn row_means(rgb: &[u8], w: usize, row: usize) -> (f64, f64, f64) {
    let base = row * w * 3;
    let (mut r, mut g, mut b) = (0u64, 0u64, 0u64);
    for col in 0..w {
        let o = base + col * 3;
        r += rgb[o] as u64;
        g += rgb[o + 1] as u64;
        b += rgb[o + 2] as u64;
    }
    let n = w as f64;
    (r as f64 / n, g as f64 / n, b as f64 / n)
}

/// Solid gray frame (128,128,128) with a moving white box, so a correct decode
/// is gray everywhere and any green row stands out immediately.
fn gray_frame(w: usize, h: usize, t: usize) -> Vec<u8> {
    let mut buf = vec![128u8; w * h * 4];
    for p in buf.chunks_exact_mut(4) {
        p[3] = 255;
    }
    let bx = (t * 11) % (w.saturating_sub(100).max(1));
    for y in 0..100.min(h) {
        for x in bx..(bx + 100).min(w) {
            let i = (y * w + x) * 4;
            buf[i] = 255;
            buf[i + 1] = 255;
            buf[i + 2] = 255;
        }
    }
    buf
}

fn main() -> anyhow::Result<()> {
    let mode = std::env::args().nth(1).unwrap_or_else(|| "sw".into());
    let (w, h) = (1920usize, 1080usize);
    let n = 120usize;
    println!("encode roundtrip [{}]: {}x{}, {} frames", mode, w, h, n);

    // Green-line hunt: QSV-encode gray frames, decode through the HARDWARE
    // decoder (best_decoder), and scan every decoded row for anomalous green.
    #[cfg(windows)]
    if mode == "mfdec" {
        let _ = tracing_subscriber::fmt().with_max_level(tracing::Level::INFO).try_init();
        let w = std::env::var("ND_W").ok().and_then(|s| s.parse().ok()).unwrap_or(w);
        let h = std::env::var("ND_H").ok().and_then(|s| s.parse().ok()).unwrap_or(h);
        println!("  mfdec dims: {}x{}", w, h);
        let mut enc = netdesk_host::encode_mf::MfEncoder::new(w, h, 60, 8_000_000)?;
        let mut dec = netdesk_core::decode::best_decoder();
        println!("encoder: {}", enc.name());
        let mut aus: Vec<Vec<u8>> = Vec::new();
        for t in 0..30 {
            for ef in enc.encode(&gray_frame(w, h, t))? {
                aus.push(ef.data);
            }
        }
        for ef in enc.drain()? {
            aus.push(ef.data);
        }
        let mut frames_checked = 0usize;
        let mut green_rows_total = 0usize;
        let mut first_report = String::new();
        for au in &aus {
            if let Some(fr) = dec.decode(au)? {
                frames_checked += 1;
                let mut green_here = 0usize;
                let mut first = None;
                for row in 0..fr.h {
                    let (r, g, b) = row_means(&fr.rgb, fr.w, row);
                    // "green" = G clearly above both R and B (gray/white have R≈G≈B).
                    if g - r > 40.0 && g - b > 40.0 {
                        green_here += 1;
                        if first.is_none() {
                            first = Some((row, r, g, b));
                        }
                    }
                }
                green_rows_total += green_here;
                if green_here > 0 && first_report.is_empty() {
                    let (row, r, g, b) = first.unwrap();
                    first_report = format!(
                        "frame {}x{}: {} green rows; first row {} rgb=({:.0},{:.0},{:.0})",
                        fr.w, fr.h, green_here, row, r, g, b
                    );
                }
            }
        }
        println!("decoded {} frames; total green rows: {}", frames_checked, green_rows_total);
        if green_rows_total == 0 {
            println!("NO GREEN LINES via hardware decode");
        } else {
            println!("GREEN DETECTED: {}", first_report);
        }
        return Ok(());
    }

    let mut dec = netdesk_core::decode::H264Decoder::new()?;
    let mut keyframes = 0usize;
    let mut decoded_ok = 0usize;
    let mut total_bytes = 0usize;
    let mut enc_ns = 0u128;
    let mut emitted = 0usize;

    // Collect all AUs (with timing) then decode, so the two encoders share one path.
    let mut aus: Vec<(Vec<u8>, bool)> = Vec::new();

    #[cfg(windows)]
    if mode == "mf" {
        let mut enc = match netdesk_host::encode_mf::MfEncoder::new(w, h, 60, 8_000_000) {
            Ok(e) => e,
            Err(e) => {
                println!("no hardware encoder ({e}); skipping mf test");
                return Ok(());
            }
        };
        println!("hardware encoder: {} (hw={})", enc.name(), enc.is_hardware());
        for t in 0..n {
            let frame = make_frame(w, h, t);
            let t0 = Instant::now();
            let outs = enc.encode(&frame)?;
            enc_ns += t0.elapsed().as_nanos();
            for ef in outs {
                aus.push((ef.data, ef.keyframe));
            }
        }
        for ef in enc.drain()? {
            aus.push((ef.data, ef.keyframe));
        }
    }

    if mode != "mf" {
        let mut enc = netdesk_host::encode::H264Encoder::new(w, h, 60, 8_000_000)?;
        for t in 0..n {
            let frame = make_frame(w, h, t);
            let t0 = Instant::now();
            let out = enc.encode(&frame)?;
            enc_ns += t0.elapsed().as_nanos();
            if let Some(ef) = out {
                aus.push((ef.data, ef.keyframe));
            }
        }
    }

    for (data, keyframe) in &aus {
        emitted += 1;
        total_bytes += data.len();
        if *keyframe {
            keyframes += 1;
        }
        if let Some(fr) = dec.decode(data)? {
            if fr.w == w && fr.h == h {
                decoded_ok += 1;
            }
        }
    }

    let enc_ms = enc_ns as f64 / n as f64 / 1e6;
    println!("  encode avg: {:.2} ms/frame  (~{:.0} fps)", enc_ms, 1000.0 / enc_ms);
    println!("  AUs emitted: {}  keyframes: {}  decoded_ok: {}", emitted, keyframes, decoded_ok);
    println!("  avg AU: {:.1} KB", total_bytes as f64 / emitted.max(1) as f64 / 1024.0);

    if keyframes == 0 {
        anyhow::bail!("no keyframe emitted");
    }
    if decoded_ok < emitted.saturating_sub(2) {
        anyhow::bail!("roundtrip decode failed: {}/{} decoded", decoded_ok, emitted);
    }
    println!("ROUNDTRIP OK");
    Ok(())
}