use std::{fs::File, io::{Cursor, Write}};

use png::{BitDepth, ColorType};

use crate::protocol::{ProtocolReader, ProtocolWriter};

mod compression;
mod constants;
mod decompression;
mod protocol;
mod quality_settings;
mod quantization;
mod transform;


fn main() {
    for chunk in vec![
        [0, 0, 0, 0, 0, 0, 0, 0],
        [0, 1, 2, 3, 4, 5, 6, 7],
        [255, 255, 255, 255, 255, 255, 255, 80],
        [255, 255, 255, 255, 255, 255, 255, 0],
        [255, 253, 254, 252, 251, 252, 255, 254]
    ] {
        let orig = chunk;
        let mut enc = chunk.clone();
        transform::encode(&mut enc, 0, 1);
        let mut dec = enc.clone();
        transform::decode(&mut dec, 0, 1);
        dbg!(orig, enc, dec);
    }

    hard_main();
}
fn hard_main() {
    let (width, height, r, g, b) = load_image();
    let mut writer = ProtocolWriter::new(vec![]);

    let quality_settings = quality_settings::QualitySettings::new(3);

    compression::compress(
        width as u32, height as u32, &[(quality_settings, &r), (quality_settings, &g), (quality_settings, &b)],
        &mut writer
    ).unwrap();
    let compressed = writer.destroy();

    let mut output_file = File::create("outputs/avatar2.rxi").unwrap();
    output_file.write_all(&compressed).unwrap();

    let mut reader = ProtocolReader::new(Cursor::new(compressed));
    let (width2, height2, decompressed) = 
        decompression::decompress(&mut reader).unwrap();

    assert_eq!(3, decompressed.len());
    save_image(width2 as usize, height2 as usize, decompressed);
}


fn load_image() -> (usize, usize, Vec<u8>, Vec<u8>, Vec<u8>) {
    let decoder = png::Decoder::new(File::open("inputs/avatar2.png").unwrap());
    let mut reader = decoder.read_info().unwrap();

    let mut buf = vec![0; reader.output_buffer_size()];
    let info = reader.next_frame(&mut buf).unwrap();
    let bytes = &buf[..info.buffer_size()];

    let width = reader.info().width;
    let height = reader.info().height;
    assert_eq!(BitDepth::Eight, reader.info().bit_depth);
    assert_eq!(3, reader.info().bytes_per_pixel());
    assert_eq!(ColorType::Rgb, reader.info().color_type);

    let r: Vec<u8> = bytes[0..].iter().cloned().step_by(3).collect();
    let g: Vec<u8> = bytes[1..].iter().cloned().step_by(3).collect();
    let b: Vec<u8> = bytes[2..].iter().cloned().step_by(3).collect();

    assert_eq!(r.len(), (width * height) as usize);
    assert_eq!(g.len(), (width * height) as usize);
    assert_eq!(b.len(), (width * height) as usize);

    (width as usize, height as usize, r, g, b)
}

fn save_image(width: usize, height: usize, image: Vec<Vec<u8>>) {
    assert_eq!(image.len(), 3);

    let mut encoder = png::Encoder::new(
        File::create("outputs/avatar2_out.png").unwrap(),
        width as u32, height as u32, 
    );
    encoder.set_color(ColorType::Rgb);

    let mut idata: Vec<u8> = vec![0; width * height * 3];
    for i in [0, 1, 2] {
        for (dst, src) in 
            idata[i..].iter_mut().step_by(3).zip(image[i].iter()) 
        {  
            *dst = *src;
        }
    }

    let mut writer = encoder.write_header().unwrap();
    writer.write_image_data(&idata).unwrap();

}