use vulkano::buffer::{CpuAccessibleBuffer, BufferUsage};
use std::sync::Arc;
use crate::canvas::Drawable;
use vulkano::framebuffer::RenderPassAbstract;
use vulkano::pipeline::{GraphicsPipelineAbstract, ComputePipeline};
use vulkano::device::Device;
use crate::util::compute_kernel::ComputeKernel;
use image::ImageBuffer;
use image::GenericImageView;
use crate::util::compute_image::ComputeImage;
use vulkano::image::{ImageUsage, AttachmentImage};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSetBuf, PersistentDescriptorSet};
use vulkano::format::Format;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use std::borrow::Borrow;
use image::Rgba;

pub struct CompuSprite {
    vertices: [(f32, f32); 6],
    position: (f32, f32),
    size: (f32, f32),
    color: (f32, f32, f32, f32),
    image_handle: Arc<u32>,
}

impl Drawable for CompuSprite {
    fn get_vertices(&self) -> Vec<(f32, f32)> {
        self.vertices.to_vec()
    }

    fn get_color(&self) -> (f32, f32, f32, f32) {
        self.color
    }

    fn get_texture_handle(&self) -> Option<Arc<u32>> {
        None
    }

    fn get_image_handle(&self) -> Option<Arc<u32>> {
        Some(self.image_handle.clone())
    }
}

impl CompuSprite {
    pub fn new(position: (f32, f32),
               size: (f32, f32),
               image_handle: Arc<u32>) -> CompuSprite {
        let fsize = (size.0 as f32, size.1 as f32);

        CompuSprite {
            vertices: [
                (position.0, position.1),                    // top left
                (position.0, position.1 + fsize.1),           // bottom left
                (position.0 + fsize.0, position.1 + fsize.1), // bottom right
                (position.0, position.1),                    // top left
                (position.0 + fsize.0, position.1 + fsize.1), // bottom right
                (position.0 + fsize.0, position.1),           // top right
            ],

            position: position,
            size: size,
            color: (0.0, 0.0, 0.0, 0.0),
            image_handle: image_handle.clone(),
        }
    }

    fn get_vertices(&self) -> Vec<(f32, f32)> {
        self.vertices.to_vec()
    }

    fn get_color(&self) -> (f32, f32, f32, f32) {
        self.color.clone()
    }

    fn get_image_handle(&self) -> Arc<u32> {
        self.image_handle.clone()
    }
}

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

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

            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])
        })
    }
}

// Canvas analog
pub struct CompuState {
    compute_buffers: Vec<CompuBuffers>,
    compute_buffer_handles: Vec<Arc<u32>>,

    kernels: Vec<ComputeKernel>,
    kernel_handles: Vec<Arc<u32>>,
}

impl CompuState {

    pub fn new() -> CompuState {
        CompuState {
            compute_buffers: vec![],
            compute_buffer_handles: vec![],
            kernels: vec![],
            kernel_handles: vec![],
        }
    }

    pub fn new_compute_buffer(&mut self,
                              data: Vec<u8>,
                              dimensions: (u32, u32),
                              stride: u32,
                              device: Arc<Device>) -> Arc<u32> {

        self.compute_buffers.push(
            CompuBuffers::new(device.clone(), data, dimensions, stride));

        let id = Arc::new(self.compute_buffers.len() as u32);
        self.compute_buffer_handles.push(id.clone());
        id
    }

    pub fn read_compute_buffer(&mut self, handle: Arc<u32>) -> Vec<u8> {
        // This is way more difficult than it should be
        //let compute_buffer : CompuBuffers = self.compute_buffers.get(handle.into()).unwrap();
        //compute_buffer.read_output_buffer().to_vec()
        Vec::new()
    }
    pub fn write_compute_buffer(&self, handle: Arc<u32>, data: Vec<u8>) {
        
    }

    pub fn new_kernel(&mut self,
                      filename: String,
                      device: &Arc<Device>) -> Arc<u32> {
        let kernel = ComputeKernel::new(filename, device.clone());

        self.kernels.push(kernel);

        let id = Arc::new(self.kernels.len() as u32);

        self.kernel_handles.push(id.clone());

        id
    }

    pub fn compute(&mut self, compute_frame: ComputeFrame) {

        // i = (Buffer, Kernel)
        for i in compute_frame.pure_compute {
            
        }

    }
}


pub struct ComputeFrame {
    // Vec<(Buffer, Kernel)>
    pure_compute: Vec<(Arc<u32>, Arc<u32>)>,

    // Vec<(Buffer, Image, Kernel)>
    swapped_to_image: Vec<(Arc<u32>, Arc<u32>, Arc<u32>)>,

    // Vec<(Input Buffer, Output Buffer, Kernel)>
    swapped_to_buffer: Vec<(Arc<u32>, Arc<u32>, Arc<u32>)>,
}

impl ComputeFrame {
    pub fn new() -> ComputeFrame {
        ComputeFrame {
            pure_compute: vec![],
            swapped_to_image: vec![],
            swapped_to_buffer: vec![]
        }
    }

    pub fn add(&mut self, buffer: Arc<u32>, kernel: Arc<u32>) {
        self.pure_compute.push((buffer, kernel));
    }

    /*
        INPUT_BUFFER -> input -> kernel -> output
                           v------------------^
        OUTPUT_BUFFER -> input X kernel X output
    */
    pub fn add_chained(&mut self, input_buffer: Arc<u32>, output_buffer: Arc<u32>, kernel: Arc<u32>) {
        self.swapped_to_buffer.push((input_buffer, output_buffer, kernel));
    }

    pub fn add_with_image_swap(&mut self, buffer: Arc<u32>, kernel: Arc<u32>, sprite: &CompuSprite) {
        self.swapped_to_image.push((buffer, sprite.get_image_handle().clone(), kernel))
    }
}
