use image::*;
use std::fs::File;
use std::io::Read;
use std::path::Path;

use crate::{Error, CHAN};

pub struct RawImage {
    pub rows: usize,
    pub cols: usize,
    pub pixels: Vec<u8>,
}

fn read_bin_u32(f: &mut File) -> Result<u32, Error> {
    let mut buffer = [0; 4];
    f.read_exact(&mut buffer)?;
    Ok(u32::from_le_bytes(buffer))
}

pub fn read_raw<P: AsRef<Path>>(
    image: P,
    crop_rows: Option<usize>,
    crop_cols: Option<usize>,
) -> Result<RawImage, Error> {
    let mut file = File::open(image)?;
    let in_rows = read_bin_u32(&mut file)? as usize;
    let in_cols = read_bin_u32(&mut file)? as usize;
    let out_rows = crop_rows.unwrap_or(in_rows);
    let out_cols = crop_cols.unwrap_or(in_cols);
    let chans = read_bin_u32(&mut file)? as usize;

    assert!(
        chans == CHAN,
        "Channel size read from the binary file doesn't match the default value"
    );
    assert!(in_rows >= out_rows);
    assert!(in_cols >= out_cols);

    let in_size: usize = in_rows * in_cols * CHAN;
    let mut buffer: Vec<u8> = vec![0; in_size];
    // The file has pixels is a 3D array with dimensions height, width, channel
    file.read_exact(&mut buffer)?;

    let chunked_channels = buffer.chunks(CHAN).collect::<Vec<_>>();
    let chunked = chunked_channels.chunks(in_cols).collect::<Vec<_>>();
    let input: &[&[&[u8]]] = chunked.as_slice();

    // We need the image in a 3D array with dimensions channel, height, width
    let out_size: usize = out_rows * out_cols * CHAN;
    let mut pixels: Vec<u8> = vec![0; out_size];
    let mut pixels_columns = pixels.chunks_mut(out_cols).collect::<Vec<_>>();
    let mut pixels_chunked = pixels_columns.chunks_mut(out_rows).collect::<Vec<_>>();
    let result: &mut [&mut [&mut [u8]]] = pixels_chunked.as_mut_slice();

    for row in 0..out_rows {
        for col in 0..out_cols {
            for chan in 0..CHAN {
                result[chan][row][col] = input[row][col][chan];
            }
        }
    }

    Ok(RawImage {
        rows: out_rows,
        cols: out_cols,
        pixels,
    })
}

pub fn extern_image(rows: usize, cols: usize, image: &[u8]) -> RgbImage {
    let chunked_columns = image.chunks(cols).collect::<Vec<_>>();
    let chunked = chunked_columns.chunks(rows).collect::<Vec<_>>();
    let pixels: &[&[&[u8]]] = chunked.as_slice();

    RgbImage::from_fn(cols as u32, rows as u32, |c, r| {
        image::Rgb([
            pixels[0][r as usize][c as usize],
            pixels[1][r as usize][c as usize],
            pixels[2][r as usize][c as usize],
        ])
    })
}