//! Port of `pipeline_360.py rig --self-test`: validates stage-4 geometry &
//! schema without any video or GPU. Pure `std` so it can run offline.

use std::path::Path;

use crate::json;
use crate::math;
use crate::stages::rig::{self, Layout, LinkMode};

/// Write a tiny valid (stored/deflate) single-colour PNG so rig's pure-std
/// dimension reader and copy-linker can operate on it without a codec.
fn write_dummy_png(path: &Path, width: u32, height: u32) -> std::io::Result<()> {
    // Build an uncompressed-all-zero RGB image: each scanline prefixed with a
    // filter byte 0, then width RGB triples.
    let mut raw: Vec<u8> = Vec::with_capacity((height as usize) * (1 + 3 * width as usize));
    for _ in 0..height {
        raw.push(0);
        raw.extend(std::iter::repeat(0u8).take((3 * width) as usize));
    }
    // zlib stream: a single "stored" deflate block (BTYPE=00) with the full
    // uncompressed length, no compression. Valid zlib with header 0x78 0x01.
    let len = raw.len() as u32;
    let mut deflate = Vec::new();
    deflate.push(0x01); // BFINAL=1, BTYPE=00
    deflate.extend_from_slice(&len.to_le_bytes());
    deflate.extend_from_slice(&(!len).to_le_bytes());
    deflate.extend_from_slice(&raw);
    let mut idat = vec![0x78, 0x01];
    // zlib adler32 of `raw`.
    let adler = adler32(&raw);
    idat.extend_from_slice(&deflate);
    idat.extend_from_slice(&adler.to_be_bytes());

    let ihdr_data = make_ihdr(width, height);

    let mut out = Vec::new();
    out.extend_from_slice(b"\x89PNG\r\n\x1a\n");
    png_chunk(&mut out, b"IHDR", &ihdr_data);
    png_chunk(&mut out, b"IDAT", &idat);
    png_chunk(&mut out, b"IEND", &[]);
    std::fs::create_dir_all(path.parent().unwrap_or(Path::new(".")))?;
    std::fs::write(path, out)
}

fn make_ihdr(width: u32, height: u32) -> Vec<u8> {
    let mut d = Vec::new();
    d.extend_from_slice(&width.to_be_bytes());
    d.extend_from_slice(&height.to_be_bytes());
    d.extend_from_slice(&[8, 2, 0, 0, 0]); // bit depth 8, colour type 2 (RGB)
    d
}

fn png_chunk(out: &mut Vec<u8>, tag: &[u8; 4], data: &[u8]) {
    out.extend_from_slice(&(data.len() as u32).to_be_bytes());
    out.extend_from_slice(tag);
    out.extend_from_slice(data);
    // CRC32 of (tag + data).
    let mut crc_input = Vec::new();
    crc_input.extend_from_slice(tag);
    crc_input.extend_from_slice(data);
    out.extend_from_slice(&crc32(&crc_input).to_be_bytes());
}

fn crc32(data: &[u8]) -> u32 {
    let mut table = [0u32; 256];
    for (n, item) in table.iter_mut().enumerate() {
        let mut c = n as u32;
        for _ in 0..8 {
            c = if c & 1 != 0 {
                0xEDB8_8320 ^ (c >> 1)
            } else {
                c >> 1
            };
        }
        *item = c;
    }
    let mut c = 0xFFFF_FFFFu32;
    for &b in data {
        c = table[((c ^ b as u32) & 0xFF) as usize] ^ (c >> 8);
    }
    c ^ 0xFFFF_FFFF
}

fn adler32(data: &[u8]) -> u32 {
    let mut a = 1u32;
    let mut b = 0u32;
    for &byte in data {
        a = (a + byte as u32) % 65521;
        b = (b + a) % 65521;
    }
    (b << 16) | a
}

fn approx_eq_vec(a: &[f64], b: &[f64], tol: f64) -> bool {
    a.len() == b.len() && a.iter().zip(b).all(|(x, y)| (x - y).abs() <= tol)
}

/// Run the six check groups. Returns `(ok_count, first_error)`.
pub fn run_selftest() -> Result<usize, String> {
    let n = 64u32;
    let mut ok = 0usize;

    // Group 1: every FACE_AXES rotation is valid.
    for face in math::FACES {
        let m = math::face_rotation_matrix(face);
        math::assert_valid_rotation(&m, &format!("FACE_AXES[{face}]"), 1e-9)
            .map_err(|e| format!("group 1: {e}"))?;
    }
    ok += 1;

    // Group 2: R<-F quaternion and 90° angle.
    let q_r = math::cam_from_rig_rotation("R", "F");
    let sqrt2_2 = std::f64::consts::SQRT_2 / 2.0;
    let expected_r = [sqrt2_2, 0.0, -sqrt2_2, 0.0];
    if !approx_eq_vec(&[q_r.0, q_r.1, q_r.2, q_r.3], &expected_r, 1e-12) {
        return Err(format!("group 2: R quaternion wrong: {q_r:?}"));
    }
    if (math::rotation_angle_deg(q_r) - 90.0).abs() >= 1e-9 {
        return Err("group 2: rotation angle not 90°".into());
    }
    ok += 1;

    // Group 3: forward directions from each face's transpose mapping +Z.
    let r_mat = math::face_rotation_matrix("R");
    let fwd_f = math::apply(&math::mat_transpose(&r_mat), &[0.0, 0.0, 1.0]);
    if !approx_eq_vec(&fwd_f, &[1.0, 0.0, 0.0], 1e-12) {
        return Err(format!("group 3: R forward wrong: {fwd_f:?}"));
    }
    for (face, expected) in [
        ("B", [0.0, 0.0, -1.0]),
        ("L", [-1.0, 0.0, 0.0]),
        ("U", [0.0, -1.0, 0.0]),
        ("D", [0.0, 1.0, 0.0]),
    ] {
        let got = math::apply(
            &math::mat_transpose(&math::face_rotation_matrix(face)),
            &[0.0, 0.0, 1.0],
        );
        if !approx_eq_vec(&got, &expected, 1e-12) {
            return Err(format!("group 3: face {face} forward wrong: {got:?}"));
        }
    }
    ok += 1;

    // Group 4: intrinsics.
    let (fx, fy, cx, cy) = math::compute_intrinsics(n, "linspace");
    if !approx_eq_vec(
        &[fx, fy, cx, cy],
        &[
            (n - 1) as f64 / 2.0,
            (n - 1) as f64 / 2.0,
            n as f64 / 2.0,
            n as f64 / 2.0,
        ],
        0.0,
    ) {
        return Err(format!("group 4: intrinsics wrong: {fx},{fy},{cx},{cy}"));
    }
    ok += 1;

    // Group 5: full rig run on a built-up dataset.
    let tmp = std::env::temp_dir().join(format!("rig_selftest_{}", std::process::id()));
    let root = tmp.as_path();
    let _ = std::fs::remove_dir_all(root);
    let flat = root.join("images");
    let masks_flat = root.join("masks");

    for i in 1..=2u32 {
        for face in math::FACES {
            write_dummy_png(&flat.join(format!("frame_{i:05}_{face}.png")), n, n)
                .map_err(|e| format!("group 5: write image: {e}"))?;
        }
    }
    for face in ["F", "R", "B", "L"] {
        write_dummy_png(&flat.join(format!("frame_00003_{face}.png")), n, n)
            .map_err(|e| format!("group 5: write partial image: {e}"))?;
    }
    for face in math::FACES {
        write_dummy_png(
            &masks_flat.join(format!("frame_00001_{face}.png.png")),
            n,
            n,
        )
        .map_err(|e| format!("group 5: write mask: {e}"))?;
    }
    for face in ["F", "R"] {
        write_dummy_png(
            &masks_flat.join(format!("frame_00003_{face}.png.png")),
            n,
            n,
        )
        .map_err(|e| format!("group 5: write partial mask: {e}"))?;
    }

    let out = root.join("rig.json");
    let link_tree = root.join("nested");
    let masks_link_tree = root.join("masks_nested");
    let _res = rig::run_rig_stage(
        &flat,
        &out,
        Some(&masks_flat),
        None,
        Layout::Auto,
        "F",
        "linspace",
        Some(&link_tree),
        Some(&masks_link_tree),
        LinkMode::Copy,
    )
    .map_err(|e| format!("group 5: run_rig_stage: {e}"))?;

    let text = std::fs::read_to_string(&out).map_err(|e| format!("group 5: read rig: {e}"))?;
    let config = json::parse(&text).map_err(|e| format!("group 5: parse rig json: {e}"))?;
    let cams = config
        .as_array()
        .and_then(|a| a.first())
        .and_then(|v| v.find("cameras"))
        .and_then(|v| v.as_array())
        .ok_or("group 5: no cameras array")?;
    if cams.len() != 6 {
        return Err(format!("group 5: expected 6 cameras, got {}", cams.len()));
    }
    let refs: Vec<&json::Value> = cams
        .iter()
        .filter(|c| matches!(c.find("ref_sensor"), Some(crate::json::Value::Bool(true))))
        .collect();
    if refs.len() != 1 {
        return Err(format!(
            "group 5: expected exactly 1 ref_sensor, got {}",
            refs.len()
        ));
    }
    if refs[0].find("image_prefix").and_then(|v| v.as_str()) != Some("F/") {
        return Err("group 5: ref prefix not F/".into());
    }
    let mut by_prefix: std::collections::HashMap<&str, &json::Value> =
        std::collections::HashMap::new();
    for c in cams {
        if let Some(p) = c.find("image_prefix").and_then(|v| v.as_str()) {
            by_prefix.insert(p, c);
        }
    }
    let got_r: Vec<f64> = by_prefix
        .get("R/")
        .and_then(|c| c.find("cam_from_rig_rotation"))
        .and_then(|v| v.as_array())
        .map(|a| a.iter().filter_map(|x| x.as_f64()).collect())
        .ok_or("group 5: no R rotation")?;
    if !approx_eq_vec(&got_r, &expected_r, 1e-12) {
        return Err(format!("group 5: R rotation wrong: {got_r:?}"));
    }
    if !link_tree.join("F").join("frame_00001.png").is_file() {
        return Err("group 5: linked image missing".into());
    }
    if !masks_link_tree
        .join("F")
        .join("frame_00001.png.png")
        .is_file()
    {
        return Err("group 5: linked mask missing".into());
    }
    if masks_link_tree
        .join("F")
        .join("frame_00002.png.png")
        .exists()
    {
        return Err("group 5: unexpected frame_00002 mask".into());
    }
    if !masks_link_tree
        .join("R")
        .join("frame_00003.png.png")
        .is_file()
    {
        return Err("group 5: partial frame_00003 mask missing".into());
    }
    ok += 1;

    // Group 6: partial rig (only 4 faces).
    let partial_dir = root.join("partial");
    for i in 1..=2u32 {
        for face in ["F", "R", "B", "L"] {
            write_dummy_png(&partial_dir.join(format!("frame_{i:05}_{face}.png")), n, n)
                .map_err(|e| format!("group 6: write partial: {e}"))?;
        }
    }
    let out2 = root.join("rig_partial.json");
    let _ = rig::run_rig_stage(
        &partial_dir,
        &out2,
        None,
        None,
        Layout::Auto,
        "F",
        "linspace",
        None,
        None,
        LinkMode::Copy,
    )
    .map_err(|e| format!("group 6: run_rig_stage: {e}"))?;
    let text2 = std::fs::read_to_string(&out2).map_err(|e| format!("group 6: read rig: {e}"))?;
    let cfg2 = json::parse(&text2).map_err(|e| format!("group 6: parse: {e}"))?;
    let cams2 = cfg2
        .as_array()
        .and_then(|a| a.first())
        .and_then(|v| v.find("cameras"))
        .and_then(|v| v.as_array())
        .ok_or("group 6: no cameras")?;
    if cams2.len() != 4 {
        return Err(format!("group 6: expected 4 cameras, got {}", cams2.len()));
    }
    ok += 1;

    let _ = std::fs::remove_dir_all(root);
    Ok(ok)
}

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

    #[test]
    fn rig_selftest_passes() {
        let ok = run_selftest().expect("self-test must pass");
        assert_eq!(ok, 6, "expected 6 check groups");
    }
}