//! Sweep a range mask's band and show what it selects.
//!
//! A diagnostic for FR-DEV-10. The other sweep example drives global
//! parameters through `EditGraph`; a band is not one of those — it lives on a
//! `MaskLayer`, is rasterised by its own pass, and only becomes visible
//! through whatever adjustment the layer carries.
//!
//! So the layer here is given a deliberately blunt adjustment — two stops down
//! — because the question this answers is *what does the band select*, not
//! *what would a photographer do with it*. A subtle edit would show a subtle
//! selection and prove nothing.
//!
//! ```sh
//! cargo run -p dr-gpu --release --example rangesweep -- IMG.CR2 out/ luminance 21
//! cargo run -p dr-gpu --release --example rangesweep -- IMG.CR2 out/ hue 21
//! ```

use dr_gpu::{AdjustPass, DemosaicedImage, Demosaicer, GpuContext, MaskPass};
use dr_pipeline::mask::{MaskLayer, MaskSource, MaskStack};
use dr_pipeline::operation::compose_full;
use dr_pipeline::ops;
use dr_pipeline::spot::SpotSet;
use dr_pipeline::{Framing, ParamId};
use dr_types::ColourSpace;

fn main() {
    env_logger::init();

    let mut args = std::env::args().skip(1);
    let (Some(input), Some(out_dir), Some(mode)) = (args.next(), args.next(), args.next()) else {
        eprintln!("usage: rangesweep <file.cr2> <out_dir> <luminance|hue|width> [steps]");
        std::process::exit(2);
    };
    let steps: usize = args
        .next()
        .and_then(|s| s.parse().ok())
        .filter(|n| *n >= 2)
        .unwrap_or(21);

    let ctx = pollster::block_on(GpuContext::new_headless()).expect("gpu");
    let image = load(&ctx, &input);
    std::fs::create_dir_all(&out_dir).expect("create out dir");

    let (full_w, full_h) = image.size();
    let longest = std::env::var("SWEEP_MAX_PX")
        .ok()
        .and_then(|s| s.parse::<u32>().ok())
        .unwrap_or(1000);
    let scale = (longest as f32 / full_w.max(full_h) as f32).min(1.0);
    let w = ((full_w as f32 * scale) as u32).max(1);
    let h = ((full_h as f32 * scale) as u32).max(1);
    println!("rendering  {w} x {h}");

    let mut masks = MaskPass::new(&ctx).expect("mask pass");
    let mut adjust = AdjustPass::new(&ctx);

    for i in 0..steps {
        let t = i as f32 / (steps - 1) as f32;

        // What moves, and what the caption should say about it.
        let (source, label) = match mode.as_str() {
            // A half-wide band walking from black to white, so the selection
            // sweeps across the tonal scale rather than merely widening.
            "luminance" => {
                let centre = t;
                let half = 0.15;
                (
                    MaskSource::luminance_range(centre - half, centre + half, 0.10),
                    format!("luminance band centred {:.2}", centre),
                )
            }
            // The hue circle, at a fixed arc and a chroma floor that keeps the
            // near-neutral parts of the picture out of it.
            "hue" => (
                MaskSource::colour_range(t, 0.08, 0.05, 1.0, 0.10),
                format!("hue {:.0} deg", t * 360.0),
            ),
            // The arc opening from nothing to everything, at a fixed hue.
            "width" => (
                MaskSource::colour_range(0.08, t * 0.5, 0.05, 1.0, 0.10),
                format!("hue width {:.0} deg", t * 0.5 * 360.0),
            ),
            other => {
                eprintln!("unknown mode `{other}`");
                std::process::exit(2);
            }
        };

        let mut layer = MaskLayer::new("sweep", source);
        // Two stops down: blunt on purpose. See the module docs.
        layer.set_param("exposure", ParamId("exposure"), -2.0);

        let mut stack = MaskStack::new();
        stack.push(layer);

        // No label field and no subject masks: a range needs neither. It is a
        // weighting over the picture's own values, so the only input it wants
        // is the picture, which is the `Some(&image)` below.
        let array = masks
            .render(&stack, None, None, Some(&image), w, h)
            .expect("rasterise");

        let shader = compose_full(
            &ops::chain(),
            &Framing::new(),
            ColourSpace::Srgb,
            &stack,
            &SpotSet::new(),
            &[],
        );
        adjust
            .render_masked(&image, &shader, w, h, Some(array))
            .expect("render");
        let (pixels, pw, ph) = adjust.export_pixels().expect("readback");

        let path = format!("{out_dir}/{mode}_{i:03}.ppm");
        write_ppm(&path, &pixels, pw, ph);
        std::fs::write(format!("{out_dir}/{mode}_{i:03}.txt"), format!("{label}\n"))
            .expect("write label");
        println!("{mode}[{i}] {label}");
    }
}

fn write_ppm(path: &str, rgba: &[u8], w: u32, h: u32) {
    use std::io::Write;
    let mut out = Vec::with_capacity((w * h * 3) as usize + 32);
    out.extend_from_slice(format!("P6\n{w} {h}\n255\n").as_bytes());
    for px in rgba.chunks_exact(4) {
        out.extend_from_slice(&px[..3]);
    }
    std::fs::File::create(path)
        .expect("create output")
        .write_all(&out)
        .expect("write output");
}

/// Load either a RAW or an already-rendered image.
///
/// A JPEG takes the path `DemosaicedImage::from_rgba8` documents: no CFA to
/// interpolate, identity colour matrix, neutral white balance, and the shader
/// linearises the gamma-encoded pixels. The controls all still work; their
/// neutral is "as the camera left it" rather than "as the sensor recorded it",
/// which is worth knowing when reading a sweep made from one.
fn load(ctx: &GpuContext, path: &str) -> DemosaicedImage {
    let bytes = std::fs::read(path).expect("read file");
    match dr_decode::probe(&bytes) {
        Some(dr_types::Format::Jpeg) => {
            let p = dr_decode::decode_jpeg(&bytes).expect("decode jpeg");
            println!("loaded     {} x {} (rendered, not raw)", p.width, p.height);
            DemosaicedImage::from_rgba8(ctx, &p.rgba, p.width, p.height).expect("upload")
        }
        _ => {
            let raw = dr_decode::decode(&bytes).expect("decode raw");
            println!("loaded     {} x {} (raw)", raw.crop.width, raw.crop.height);
            let demosaic = Demosaicer::new(ctx).expect("demosaicer");
            demosaic.run(&raw).expect("demosaic")
        }
    }
}