use std::fs::File;

use png::{BitDepth, ColorType};

use crate::{image::{Image, ImageError}, quality_settings::QualitySettings};

pub fn load_image(filename: String, quality_rgb: QualitySettings, quality_alpha: QualitySettings) -> Result<Image, ImageError> {
    let decoder = png::Decoder::new(File::open(filename).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()];

    assert_eq!(BitDepth::Eight, reader.info().bit_depth);
    assert_eq!(ColorType::Rgb, reader.info().color_type);

    let width = reader.info().width;
    let height = reader.info().height;

    match reader.info().color_type {
        ColorType::Rgb => {
            assert_eq!(3, reader.info().bytes_per_pixel());

            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();

            let mut image = Image::new(width as usize, height as usize);
            image.add(quality_rgb, r)?;
            image.add(quality_rgb, g)?;
            image.add(quality_rgb, b)?;
            return Ok(image);
        }
        ColorType::Rgba => {
            assert_eq!(4, reader.info().bytes_per_pixel());

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

            let mut image = Image::new(width as usize, height as usize);
            image.add(quality_rgb, r)?;
            image.add(quality_rgb, g)?;
            image.add(quality_rgb, b)?;
            image.add(quality_alpha, a)?;
            return Ok(image);
        }
        ct => {
            return Err(ImageError::InvalidInputColorType(ct))
        }
    }
}

pub fn save_image(filename: String, image: Image) -> Result<(), ImageError> {
    let n_layers = image.layers.len();

    let color_type = match n_layers {
        3 => ColorType::Rgb,
        4 => ColorType::Rgba,
        _ => return Err(ImageError::WrongNumberOfLayersToSave(n_layers)),
    };

    let mut encoder = png::Encoder::new(
        File::create(filename)?,
        image.width as u32, image.height as u32, 
    );
    encoder.set_color(color_type);

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

    let mut writer = encoder.write_header()?;
    writer.write_image_data(&idata)?;
    Ok(())
}