use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet};
use vulkano::device::{Device};
use vulkano::pipeline::{ComputePipeline};
use std::time::SystemTime;
use std::sync::Arc;
use std::path::PathBuf;
use image::{ImageBuffer};
use image::GenericImageView;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSetBuf};
use vulkano::image::attachment::AttachmentImage;
use vulkano::image::{ImageUsage};
use vulkano::format::Format;
use image::Rgba;

/*

Compute Image holds read write swap and settings buffers for the kernel

This is a pretty specific use case. One in for settings. One in for data, two for the transfer.

multiple data inputs might be nice?

*/

#[derive(Clone)]
pub struct ComputeImage {

    device: Arc<Device>,

    compute_graphics_swap_buffer: std::sync::Arc<vulkano::image::attachment::AttachmentImage>,

    pub rw_buffers: Vec<Arc<CpuAccessibleBuffer<[u8]>>>,
    pub settings_buffer: Arc<CpuAccessibleBuffer<[u32]>>,
}


impl ComputeImage {

    pub fn load_raw(filename: String) -> (Vec<u8>, (u32,u32)) {

        let project_root =
            std::env::current_dir()
                .expect("failed to get root directory");

        let mut compute_path = project_root.clone();
        compute_path.push(PathBuf::from("resources/images/"));
        compute_path.push(PathBuf::from(filename.clone()));

        let img = image::open(compute_path).expect("Couldn't find image");

        let xy = img.dimensions();

        let data_length = xy.0 * xy.1 * 4;
        let pixel_count = img.raw_pixels().len();

        let mut image_buffer = Vec::new();

        if pixel_count != data_length as usize {
            println!("Creating apha channel...");
            for i in img.raw_pixels().iter() {
                if (image_buffer.len() + 1) % 4 == 0 {
                    image_buffer.push(255);
                }
                image_buffer.push(*i);
            }
            image_buffer.push(255);
        } else {
            image_buffer = img.raw_pixels();
        }

        (image_buffer, xy)
    }

    pub fn new(device: Arc<Device>, image_filename: String) -> ComputeImage {

        let (image_buffer, xy) = ComputeImage::load_raw(image_filename);

        let compute_graphics_swap_buffer = {

            let mut usage = ImageUsage::none();
            usage.transfer_destination = true;
            usage.storage = true;

            AttachmentImage::with_usage(
                device.clone(),
                [xy.0, xy.1],
                Format::R8G8B8A8Uint,
                usage)
        };

        let data_length = xy.0 * xy.1 * 4;

        // Pull out the image data and place it in a buffer for the kernel to write to and for us to read from
        let write_buffer = {
            let mut buff = image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };


        // Pull out the image data and place it in a buffer for the kernel to read from
        let read_buffer = {
            let mut buff = image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };


        // A buffer to hold many i32 values to use as settings
        let settings_buffer = {
            let vec = vec![xy.0, xy.1];
            let mut buff = vec.iter();
            let data_iter =
                (0..2).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(),
                                           BufferUsage::all(),
                                           data_iter).unwrap()
        };

        ComputeImage{
            device: device.clone(),
            compute_graphics_swap_buffer: compute_graphics_swap_buffer.unwrap(),

            rw_buffers: vec![write_buffer, read_buffer],
            settings_buffer: settings_buffer
        }
    }

    pub fn get_swap_buffer(&mut self) -> Arc<AttachmentImage> {
        self.compute_graphics_swap_buffer.clone()
    }

    pub fn get_size(&self) -> (u32, u32) {
        let xy = self.compute_graphics_swap_buffer.dimensions();
        (xy[0], xy[1])
    }

    pub fn read_read_buffer(&self) -> ImageBuffer<Rgba<u8>, Vec<u8>>{

        let xy = self.get_size();

        self.rw_buffers.get(0).unwrap().write().unwrap().map(|x| x);
        let data_buffer_content = self.rw_buffers.get(0).unwrap().read().unwrap();
        ImageBuffer::from_fn(xy.0, xy.1, |x, y| {
            let r = data_buffer_content[((xy.0 * y + x) * 4 + 0) as usize] as u8;
            let g = data_buffer_content[((xy.0 * y + x) * 4 + 1) as usize] as u8;
            let b = data_buffer_content[((xy.0 * y + x) * 4 + 2) as usize] as u8;
            let a = data_buffer_content[((xy.0 * y + x) * 4 + 3) as usize] as u8;

            image::Rgba([r, g, b, a])
        })
    }

    pub fn save_image(&self) {
        self.read_read_buffer().save(format!("output/{}.jpg", SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs()));
    }

    pub fn get_descriptor_set(&self, compute_pipeline: std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>)
        -> Arc<PersistentDescriptorSet<std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>, ((((),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u32]>>>)>> {

        Arc::new(PersistentDescriptorSet::start(compute_pipeline.clone(), 0)
            .add_buffer(self.rw_buffers.get(0).unwrap().clone()).unwrap()
            .add_buffer(self.rw_buffers.get(1).unwrap().clone()).unwrap()
            .add_buffer(self.settings_buffer.clone()).unwrap()
            .build().unwrap())
    }
}