use image::io::Reader as ImageReader;
use rand::seq::SliceRandom;
use rand::thread_rng;

fn main() -> anyhow::Result<()> {
    if let Some(filename) = std::env::args().nth(1) {
        println!("{}", filename);

        let img = ImageReader::open(filename)?.decode()?;
        let rimage = img.to_rgb8();

        let all_pixels: Vec<(u32, u32, &image::Rgb<u8>)> = rimage.enumerate_pixels().collect();
        let mut iter = 0;
        println!("Splitting into chunks...");
        let mut chunks = all_pixels.chunks(100_000).collect::<Vec<_>>();
        chunks.shuffle(&mut thread_rng());
        for chunk in &chunks {
            save(rimage.width(), rimage.height(), chunk, iter)?;
            println!("Saved image");
            iter += 1;
        }
    } else {
        println!("Please provide a valid filename.");
    }

    Ok(())
}

fn save(
    width: u32,
    height: u32,
    chunk: &[(u32, u32, &image::Rgb<u8>)],
    iter: u32,
) -> anyhow::Result<()> {
    let white = image::Rgb([255, 255, 255]);
    let mut out = image::ImageBuffer::from_pixel(width, height, white);

    println!("Created image");

    for (x, y, px) in chunk {
        out.put_pixel(*x, *y, **px);
    }

    println!("Applied pixels to new image");

    let bytes: Vec<u8> = out.into_raw();

    if let Err(e) = std::panic::catch_unwind(|| {
        let mut comp = mozjpeg::Compress::new(mozjpeg::ColorSpace::JCS_RGB);

        comp.set_size(width as usize, height as usize);
        comp.set_mem_dest();
        comp.start_compress();

        assert!(comp.write_scanlines(&bytes[..]));

        comp.finish_compress();
        if let Ok(bytes) = comp.data_to_vec() {
            println!("Compressed image");
            if let Ok(()) = std::fs::write(format!("output/{}.jpg", iter), bytes) {
                println!("Wrote new file to disk");
            } else {
                panic!("Error: could not save the file to disk.");
            }
        }
    }) {
        panic!("Error: {:?}", e);
    }

    Ok(())
}