use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet, StdDescriptorPoolAlloc, PersistentDescriptorSetBuilder, FixedSizeDescriptorSetsPool, StdDescriptorPool};
use vulkano::descriptor::descriptor_set::collection::DescriptorSetsCollection;
use vulkano::device::{Device, DeviceExtensions, QueuesIter, Queue};
use vulkano::instance::{Instance, InstanceExtensions, PhysicalDevice, QueueFamily};
use vulkano::pipeline::{ComputePipeline, GraphicsPipeline, GraphicsPipelineAbstract};
use vulkano::sync::{GpuFuture, FlushError};
use vulkano::sync;
use std::time::SystemTime;
use std::sync::Arc;
use std::ffi::CStr;
use std::path::PathBuf;
use shade_runner as sr;
use image::{DynamicImage, ImageBuffer};
use image::GenericImageView;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use image::GenericImage;
use shade_runner::{ComputeLayout, CompileError, FragLayout, FragInput, FragOutput, VertInput, VertOutput, VertLayout, CompiledShaders, Entry};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSetBuf, PersistentDescriptorSetImg, PersistentDescriptorSetSampler};
use shaderc::CompileOptions;
use vulkano::framebuffer::{Subpass, RenderPass, RenderPassAbstract, Framebuffer, FramebufferAbstract};
use vulkano::pipeline::shader::{GraphicsShaderType, ShaderModule, GraphicsEntryPoint, SpecializationConstants, SpecializationMapEntry};
use vulkano::swapchain::{Swapchain, PresentMode, SurfaceTransform, Surface, SwapchainCreationError, AcquireError, Capabilities};
use vulkano::swapchain::acquire_next_image;
use vulkano::image::swapchain::SwapchainImage;
use winit::{EventsLoop, WindowBuilder, Window, Event, WindowEvent};
use vulkano_win::VkSurfaceBuild;
use vulkano::pipeline::vertex::{SingleBufferDefinition, Vertex};
use vulkano::descriptor::{PipelineLayoutAbstract, DescriptorSet};
use std::alloc::Layout;
use vulkano::pipeline::viewport::Viewport;
use image::ImageFormat;
use vulkano::image::immutable::ImmutableImage;
use vulkano::image::attachment::AttachmentImage;
use vulkano::image::{Dimensions, ImageUsage, ImageAccess, ImageDimensions};
use vulkano::format::Format;
use vulkano::format::ClearValue;
use vulkano::sampler::{Sampler, Filter, MipmapMode, SamplerAddressMode};
use image::flat::NormalForm::ColumnMajorPacked;

use crate::util::compute_kernel::ComputeKernel;
use crate::util::shader_kernels::ShaderKernels;
use crate::util::compute_image::ComputeImage;

use vulkano::descriptor::descriptor::DescriptorDesc;
use crate::vertex_2d::ColoredVertex2D;



/// This method is called once during initialization, then again whenever the window is resized
fn window_size_dependent_setup(
    images: &[Arc<SwapchainImage<Window>>],
    render_pass: Arc<dyn RenderPassAbstract + Send + Sync>,
    dynamic_state: &mut DynamicState,
) -> Vec<Arc<dyn FramebufferAbstract + Send + Sync>> {
    let dimensions = images[0].dimensions();

    let viewport = Viewport {
        origin: [0.0, 0.0],
        dimensions: [dimensions.width() as f32, dimensions.height() as f32],
        depth_range: 0.0..1.0,
    };
    dynamic_state.viewports = Some(vec!(viewport));

    images.iter().map(|image| {
        Arc::new(
            Framebuffer::start(render_pass.clone())
                .add(image.clone()).unwrap()
                .build().unwrap()
        ) as Arc<dyn FramebufferAbstract + Send + Sync>
    }).collect::<Vec<_>>()
}

pub struct VkProcessor<'a> {

    // Vulkan state fields
    pub instance: Arc<Instance>,
    pub physical: PhysicalDevice<'a>,
    pub device: Arc<Device>,
    pub queues: QueuesIter,
    pub queue: Arc<Queue>,
    pub dynamic_state: DynamicState,

    // TODO: This will need to handle multiple of each type
    pub shader_kernels: Option<ShaderKernels>,
    pub compute_kernel: Option<ComputeKernel>,

    // TODO: Move this into canvas
    pub vertex_buffer: Option<Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync + 'static)>>,
    pub vertex_buffer2: Option<Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync + 'static)>>,

    pub textures: Vec<Arc<ImmutableImage<Format>>>,

    pub compute_image: Option<ComputeImage>,

    pub swapchain: Option<Arc<Swapchain<Window>>>,
    pub swapchain_images: Option<Vec<Arc<SwapchainImage<Window>>>>,

    swapchain_recreate_needed: bool,

    capabilities: Capabilities
}


impl<'a> VkProcessor<'a> {

    pub fn new(instance: &'a Arc<Instance>, surface: &'a Arc<Surface<Window>>) -> VkProcessor<'a> {



        let physical = PhysicalDevice::enumerate(instance).next().unwrap();

        let queue_family = physical.queue_families().find(|&q| {
            // We take the first queue that supports drawing to our window.
            q.supports_graphics() &&
                surface.is_supported(q).unwrap_or(false) &&
                q.supports_compute()
        }).unwrap();

        let device_ext = DeviceExtensions { khr_swapchain: true, ..DeviceExtensions::none() };

        let (device, mut queues) = Device::new(physical,
                                               physical.supported_features(),
                                               &device_ext,
                                               [(queue_family, 0.5)].iter().cloned()).unwrap();
        let queue = queues.next().unwrap();

        VkProcessor {
            instance: instance.clone(),
            physical: physical.clone(),
            device: device.clone(),
            queue: queue,
            queues: queues,
            dynamic_state: DynamicState { line_width: None, viewports: None, scissors: None },
            shader_kernels: None,
            compute_kernel: None,
            vertex_buffer: None,
            vertex_buffer2: None,
            textures: vec![],
            compute_image: None,
            swapchain: None,
            swapchain_images: None,
            swapchain_recreate_needed: false,
            capabilities: surface.capabilities(physical).unwrap()
        }
    }


    pub fn compile_kernel(&mut self, filename: String) {
        self.compute_kernel   = Some(ComputeKernel::new(filename, self.device.clone()));
    }

    pub fn compile_shaders(&mut self, filename: String, surface: &'a Arc<Surface<Window>>) {
        self.shader_kernels = Some(
            ShaderKernels::new(filename.clone(),
                               self.capabilities.clone(), self.queue.clone(),
                               self.physical,
                               self.device.clone())
        );
    }

    pub fn create_swapchain(&mut self, surface: &'a Arc<Surface<Window>>) {

        let (mut swapchain, images) = {
            let capabilities = surface.capabilities(self.physical).unwrap();
            let usage = capabilities.supported_usage_flags;
            let alpha = capabilities.supported_composite_alpha.iter().next().unwrap();
            // Choosing the internal format that the images will have.
            let format = capabilities.supported_formats[0].0;

            // Set the swapchains window dimensions
            let initial_dimensions = if let Some(dimensions) = surface.window().get_inner_size() {
                // convert to physical pixels
                let dimensions: (u32, u32) = dimensions.to_physical(surface.window().get_hidpi_factor()).into();
                [dimensions.0, dimensions.1]
            } else {
                // The window no longer exists so exit the application.
                panic!("window closed");
            };

            Swapchain::new(self.device.clone(),
                           surface.clone(),
                           capabilities.min_image_count, // number of attachment images
                           format,
                           initial_dimensions,
                           1, // Layers
                           usage,
                           &self.queue,
                           SurfaceTransform::Identity,
                           alpha,
                           PresentMode::Fifo, true, None).unwrap()
        };

        self.swapchain = Some(swapchain);
        self.swapchain_images = Some(images);

    }

    // On resizes we have to recreate the swapchain
    pub fn recreate_swapchain(&mut self, surface: &'a Arc<Surface<Window>>) {

        let dimensions = if let Some(dimensions) = surface.window().get_inner_size() {
            let dimensions: (u32, u32) = dimensions.to_physical(surface.window().get_hidpi_factor()).into();
            [dimensions.0, dimensions.1]
        } else {
            return;
        };

        let (new_swapchain, new_images) = match self.swapchain.clone().unwrap().clone().recreate_with_dimension(dimensions) {
            Ok(r) => r,
            // This error tends to happen when the user is manually resizing the window.
            // Simply restarting the loop is the easiest way to fix this issue.
            Err(SwapchainCreationError::UnsupportedDimensions) => panic!("Uh oh"),
            Err(err) => panic!("{:?}", err)
        };

        self.swapchain = Some(new_swapchain);
        self.swapchain_images = Some(new_images);

    }

    fn get_texture_from_file(image_filename: String, queue: Arc<Queue>) -> Arc<ImmutableImage<Format>> {

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

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

        let (texture, tex_future) = ImmutableImage::from_iter(
            image_buffer.iter().cloned(),
            Dimensions::Dim2d { width: xy.0, height: xy.1 },
            Format::R8G8B8A8Srgb,
            queue.clone()
        ).unwrap();

        texture
    }

    pub fn load_compute_image(&mut self, image_filename: String) {
        self.compute_image = Some(ComputeImage::new(self.device.clone(), image_filename.clone()));
    }

    pub fn load_buffers(&mut self, image_filename: String)
    {
        self.load_compute_image(image_filename.clone());

        let color = [1.,0.,0.,0.];

        let vertex_buffer = {

            CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::all(), [
                ColoredVertex2D { position: [ 1.0, 1.0 ], color },
                ColoredVertex2D { position: [ 1.0, 0.5 ], color },
                ColoredVertex2D { position: [ 0.5, 0.5 ], color },
                ColoredVertex2D { position: [ 0.5, 1.0 ], color },
            ].iter().cloned()).unwrap()
        };

        let vertex_buffer2 = {

            CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::all(), [
                ColoredVertex2D { position: [-1.0, -1.0 ], color },
                ColoredVertex2D { position: [-1.0, -0.5 ], color },
                ColoredVertex2D { position: [-0.5, -0.5 ], color },
                ColoredVertex2D { position: [-0.5, -1.0 ], color },
            ].iter().cloned()).unwrap()
        };




        self.vertex_buffer = Some(vertex_buffer);
        self.vertex_buffer2 = Some(vertex_buffer2);

        let texture = VkProcessor::get_texture_from_file(image_filename.clone(), self.queue.clone());
        self.textures.push(texture);

        let texture1 = VkProcessor::get_texture_from_file(String::from("button.png"), self.queue.clone());
        self.textures.push(texture1);
    }


    // The image set is the containing object for all texture and image hooks.
    fn get_image_set(&mut self) -> Box<DescriptorSet + Send + Sync> {

        let sampler = Sampler::new(self.device.clone(), Filter::Linear, Filter::Linear,
                                   MipmapMode::Nearest, SamplerAddressMode::Repeat, SamplerAddressMode::Repeat,
                                   SamplerAddressMode::Repeat, 0.0, 1.0, 0.0, 0.0).unwrap();

        let o : Box<DescriptorSet + Send + Sync> = Box::new(
        PersistentDescriptorSet::start(
            self.shader_kernels.clone().unwrap().get_pipeline(), 0
        )
            .add_sampled_image(self.textures.get(0).unwrap().clone(), sampler.clone()).unwrap()
            .add_image(self.compute_image.clone().unwrap().clone().get_swap_buffer().clone()).unwrap()
            .build().unwrap());
        o
    }

    // The image set is the containing object for all texture and image hooks.
    fn get_gui_image_set(&mut self) -> Box<DescriptorSet + Send + Sync> {

        let sampler = Sampler::new(self.device.clone(), Filter::Linear, Filter::Linear,
                                   MipmapMode::Nearest, SamplerAddressMode::Repeat, SamplerAddressMode::Repeat,
                                   SamplerAddressMode::Repeat, 0.0, 1.0, 0.0, 0.0).unwrap();

        let o : Box<DescriptorSet + Send + Sync> = Box::new(
            PersistentDescriptorSet::start(
                self.shader_kernels.clone().unwrap().get_pipeline(), 0
            )
                .add_sampled_image(self.textures.get(1).unwrap().clone(), sampler.clone()).unwrap()
                .add_image(self.compute_image.clone().unwrap().clone().get_swap_buffer().clone()).unwrap()
                .build().unwrap());
        o
    }

    pub fn save_edges_image(&mut self){
        self.compute_image.clone().unwrap().clone().save_image();
    }



    pub fn run(&mut self, surface: &'a Arc<Surface<Window>>, mut frame_future: Box<dyn GpuFuture>) -> Box<dyn GpuFuture> {

        let mut framebuffers = window_size_dependent_setup(&self.swapchain_images.clone().unwrap().clone(),
                                                           self.shader_kernels.clone().unwrap().render_pass.clone(),
                                                           &mut self.dynamic_state);

        // The docs said to call this on each loop.
        frame_future.cleanup_finished();

        // Whenever the window resizes we need to recreate everything dependent on the window size.
        // In this example that includes the swapchain, the framebuffers and the dynamic state viewport.
        if self.swapchain_recreate_needed {
            self.recreate_swapchain(surface);
            framebuffers = window_size_dependent_setup(&self.swapchain_images.clone().unwrap().clone(),
                                                       self.shader_kernels.clone().unwrap().render_pass.clone(),
                                                       &mut self.dynamic_state);
            self.swapchain_recreate_needed = false;
        }

        // This function can block if no image is available. The parameter is an optional timeout
        // after which the function call will return an error.
        let (image_num, acquire_future) = match vulkano::swapchain::acquire_next_image(self.swapchain.clone().unwrap().clone(), None) {
            Ok(r) => r,
            Err(AcquireError::OutOfDate) => {
                self.swapchain_recreate_needed = true;
                return Box::new(sync::now(self.device.clone())) as Box<_>;
            }
            Err(err) => panic!("{:?}", err)
        };

        // Specify the color to clear the framebuffer with i.e. blue
        let clear_values = vec!(ClearValue::Float([0.0, 0.0, 1.0, 1.0]));

        let mut v = Vec::new();
        v.push(self.vertex_buffer.clone().unwrap().clone());

        let mut v2 = Vec::new();
        v2.push(self.vertex_buffer2.clone().unwrap().clone());

        let xy = self.compute_image.clone().unwrap().get_size();

        let mut command_buffer =
            AutoCommandBufferBuilder::primary_one_time_submit(self.device.clone(), self.queue.family())
                .unwrap()

                .dispatch([xy.0, xy.1, 1],
                          self.compute_kernel.clone().unwrap().clone().get_pipeline(),
                          self.compute_image.clone().unwrap().clone()
                              .get_descriptor_set(self.compute_kernel.clone().unwrap().clone().get_pipeline()).clone(), ()).unwrap()

                .copy_buffer_to_image(self.compute_image.clone().unwrap().clone().rw_buffers.get(0).unwrap().clone(),
                    self.compute_image.clone().unwrap().clone().get_swap_buffer().clone()).unwrap()

                .begin_render_pass(framebuffers[image_num].clone(), false, clear_values.clone())
                .unwrap()

                .draw(self.shader_kernels.clone().unwrap().get_pipeline(),
                      &self.dynamic_state.clone(), v,
                      vec![self.get_image_set()], ())
                .unwrap();

        command_buffer = command_buffer.draw(self.shader_kernels.clone().unwrap().get_pipeline(),
                      &self.dynamic_state.clone(), v2,
                      vec![self.get_gui_image_set()], ())
                .unwrap();

        let command_buffer = command_buffer
                .end_render_pass()
                .unwrap()

                .build().unwrap();

        // Wait on the previous frame, then execute the command buffer and present the image
        let future = frame_future.join(acquire_future)
            .then_execute(self.queue.clone(), command_buffer).unwrap()
            .then_swapchain_present(self.queue.clone(), self.swapchain.clone().unwrap().clone(), image_num)
            .then_signal_fence_and_flush();

        match future {
            Ok(future) => {
                (Box::new(future) as Box<_>)
            }
            Err(FlushError::OutOfDate) => {
                self.swapchain_recreate_needed = true;
                (Box::new(sync::now(self.device.clone())) as Box<_>)
            }
            Err(e) => {
                println!("{:?}", e);
                (Box::new(sync::now(self.device.clone())) as Box<_>)
            }
        }
    }
}