//! Write with/without rim comparison PNGs for acceptance artifacts.
use paperdoll_renderer::compositor::{
    apply_rim, compose, relight, LayerMaps, Light, RimParams, DEFAULT_DEPTH_EDGE_SCALE,
    DEFAULT_RIM_POWER,
};
use paperdoll_renderer::render_job::ImageBuf;
use std::env;
use std::path::PathBuf;

fn encode_normal(n: [f32; 3]) -> [u8; 4] {
    let len = (n[0] * n[0] + n[1] * n[1] + n[2] * n[2]).sqrt().max(1e-10);
    let n = [n[0] / len, n[1] / len, n[2] / len];
    [
        ((n[0] * 0.5 + 0.5) * 255.0).round() as u8,
        ((n[1] * 0.5 + 0.5) * 255.0).round() as u8,
        ((n[2] * 0.5 + 0.5) * 255.0).round() as u8,
        255,
    ]
}

fn sphere_layer(size: u32) -> LayerMaps {
    let w = size as usize;
    let mut diffuse = vec![0u8; w * w * 4];
    let mut normal = vec![0u8; w * w * 4];
    let mut specular = vec![0u8; w * w * 4];
    let mut depth = vec![0u16; w * w * 2];
    let cx = (size as f32 - 1.0) * 0.5;
    let cy = cx;
    let radius = cx * 0.92;
    for y in 0..w {
        for x in 0..w {
            let i = y * w + x;
            let dx = x as f32 - cx;
            let dy = y as f32 - cy;
            let r2 = dx * dx + dy * dy;
            if r2 > radius * radius {
                continue;
            }
            let nz = (1.0 - (r2 / (radius * radius))).sqrt();
            let nx = dx / radius;
            let ny = -dy / radius;
            let n = encode_normal([nx, ny, nz]);
            let o = i * 4;
            // Warm albedo
            diffuse[o] = 160;
            diffuse[o + 1] = 90;
            diffuse[o + 2] = 70;
            diffuse[o + 3] = 255;
            normal[o..o + 4].copy_from_slice(&n);
            specular[o] = 40;
            specular[o + 1] = 40;
            specular[o + 2] = 40;
            specular[o + 3] = 255;
            // Nearer at center (smaller gray)
            let gray = (5000.0 + (1.0 - nz) * 25000.0) as u16;
            depth[i * 2] = gray;
            depth[i * 2 + 1] = 65535;
        }
    }
    LayerMaps {
        label: "sphere".into(),
        width: size,
        height: size,
        diffuse,
        normal,
        specular,
        depth,
    }
}

fn side_by_side(a: &ImageBuf, b: &ImageBuf) -> ImageBuf {
    assert_eq!(a.width, b.width);
    assert_eq!(a.height, b.height);
    let w = a.width as usize;
    let h = a.height as usize;
    let out_w = w * 2 + 4;
    let mut rgba = vec![30u8; out_w * h * 4];
    for y in 0..h {
        for x in 0..w {
            let si = (y * w + x) * 4;
            let dl = (y * out_w + x) * 4;
            let dr = (y * out_w + (x + w + 4)) * 4;
            rgba[dl..dl + 4].copy_from_slice(&a.rgba[si..si + 4]);
            rgba[dr..dr + 4].copy_from_slice(&b.rgba[si..si + 4]);
        }
    }
    ImageBuf {
        width: out_w as u32,
        height: a.height,
        rgba,
    }
}

fn main() -> Result<(), String> {
    let out = PathBuf::from(
        env::args()
            .nth(1)
            .unwrap_or_else(|| "/workspace/paperdoll-rim-artifacts".into()),
    );
    std::fs::create_dir_all(&out).map_err(|e| e.to_string())?;
    let layer = sphere_layer(96);
    let composed = compose(&[layer]).map_err(|e| e.to_string())?;
    let light = Light::default();
    let lit = relight(&composed, &light).map_err(|e| e.to_string())?;
    let rimmed = apply_rim(
        &composed,
        &lit,
        &RimParams {
            intensity: 0.35,
            power: DEFAULT_RIM_POWER,
            depth_edge_weight: 0.25,
            depth_edge_scale: DEFAULT_DEPTH_EDGE_SCALE,
        },
    )
    .map_err(|e| e.to_string())?;
    let off = apply_rim(&composed, &lit, &RimParams::off()).map_err(|e| e.to_string())?;
    assert_eq!(off.rgba, lit.rgba);
    let compare = side_by_side(&lit, &rimmed);
    for (name, img) in [
        ("rim_off.png", &lit),
        ("rim_on.png", &rimmed),
        ("rim_compare_off_vs_on.png", &compare),
    ] {
        let path = out.join(name);
        let file = std::fs::File::create(&path).map_err(|e| e.to_string())?;
        let mut encoder = png::Encoder::new(file, img.width, img.height);
        encoder.set_color(png::ColorType::Rgba);
        encoder.set_depth(png::BitDepth::Eight);
        let mut writer = encoder.write_header().map_err(|e| e.to_string())?;
        writer
            .write_image_data(&img.rgba)
            .map_err(|e| e.to_string())?;
        writer.finish().map_err(|e| e.to_string())?;
    }
    // Bound check note
    println!("wrote comparison PNGs under {}", out.display());
    Ok(())
}