use std::sync::Arc;
use vulkano::device::Device;
use vulkano::buffer::{CpuAccessibleBuffer, BufferUsage};
use vulkano::pipeline::ComputePipeline;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use vulkano::descriptor::descriptor_set::PersistentDescriptorSet;
use image::ImageBuffer;

#[derive(Clone)]
pub struct CompuBuffers {
    dimensions: (u32, u32),
    device: Arc<Device>,
    handle: Arc<CompuBufferHandle>,

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

impl CompuBuffers {
    pub fn new(device: Arc<Device>, data: Vec<u8>,
               dimensions: (u32, u32), stride: u32,
               handle: Arc<CompuBufferHandle>) -> CompuBuffers {

        let data_length = dimensions.0 * dimensions.1 * stride;

        let input_buffer = {
            let mut buff = data.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };

        let output_buffer = {
            let mut buff = data.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };

        // Settings buffer which holds i32's
        // Compile macros into the kernel eventually to index them
        let settings_buffer = {
            let vec = vec![dimensions.0, dimensions.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()
        };

        CompuBuffers {
            dimensions: dimensions,
            device: device.clone(),

            handle: handle,
            io_buffers: vec![input_buffer, output_buffer],
            settings_buffer: settings_buffer,
        }
    }

    pub fn get_size(&self) -> (u32, u32) {
        self.dimensions
    }

    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.io_buffers.get(0).unwrap().clone()).unwrap()
            .add_buffer(self.io_buffers.get(1).unwrap().clone()).unwrap()
            .add_buffer(self.settings_buffer.clone()).unwrap()
            .build().unwrap())
    }

    pub fn read_output_buffer(&self) -> ImageBuffer<Rgba<u8>, Vec<u8>> {
        let xy = self.get_size();

        self.io_buffers.get(1).unwrap().write().unwrap().map(|x| x);
        let data_buffer_content = self.io_buffers.get(1).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])
        })
    }
}

#[derive(Clone)]
pub struct CompuBufferHandle {
    handle: u32,
}
#[derive(Clone)]
pub struct CompuKernelHandle {
    handle: u32,
}