use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet, StdDescriptorPoolAlloc};
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};
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;
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};
use vulkano::format::Format;
use vulkano::sampler::{Sampler, Filter, MipmapMode, SamplerAddressMode};
use image::flat::NormalForm::ColumnMajorPacked;


#[derive(Clone)]
pub struct ComputeKernel {

    compute_pipeline: Option<std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>>,
    compute_kernel_path: PathBuf,

    shader: CompiledShaders,
    entry: Entry,
    shader_module: Arc<ShaderModule>,
    device: Arc<Device>,
    specialization_constants: ComputeSpecializationConstants,
}

impl ComputeKernel {

    fn get_path(filename: String) -> PathBuf {

        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/shaders/"));
        compute_path.push(PathBuf::from(filename));

        compute_path
    }

    pub fn new(filename: String, device: Arc<Device>) -> ComputeKernel {

        let compute_path = ComputeKernel::get_path(filename);

        let mut options  = CompileOptions::new().ok_or(CompileError::CreateCompiler).unwrap();

        let shader = sr::load_compute_with_options(compute_path.clone(), options)
            .expect("Failed to compile");

        let entry = sr::parse_compute(&shader)
            .expect("Failed to parse");

        let shader_module = unsafe {
            vulkano::pipeline::shader::ShaderModule::from_words(device.clone(), &shader.compute)
        }.unwrap();


        ComputeKernel {
            device: device,
            shader: shader,
            compute_pipeline: Option::None,
            compute_kernel_path: compute_path,
            entry: entry,
            shader_module: shader_module,
            specialization_constants: ComputeSpecializationConstants {
                first_constant: 0,
                second_constant: 0,
                third_constant: 0.0
            }
        }
    }

    pub fn get_pipeline(&mut self) -> std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>> {

        match self.compute_pipeline.clone() {
            Some(t) => t,
            None => {
                self.compute_pipeline = Some(Arc::new({
                    unsafe {
                        ComputePipeline::new(self.device.clone(), &self.shader_module.compute_entry_point(
                            CStr::from_bytes_with_nul_unchecked(b"main\0"),
                            self.entry.compute_layout.clone()), &self.specialization_constants,
                        ).unwrap()
                    }
                }));
                self.compute_pipeline.clone().unwrap()
            }
        }
    }

    pub fn recompile_kernel(&mut self) -> std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>  {
        self.compile_kernel(String::from(self.compute_kernel_path.clone().to_str().unwrap()))
    }

    pub fn compile_kernel(&mut self, filename: String) -> std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>> {

        let mut options = CompileOptions::new().ok_or(CompileError::CreateCompiler).unwrap();
        self.compute_kernel_path = ComputeKernel::get_path(filename);

        self.shader =
            sr::load_compute_with_options(self.compute_kernel_path.clone(), options)
                .expect("Failed to compile");

        self.entry =
            sr::parse_compute(&self.shader)
                .expect("Failed to parse");

        self.shader_module = unsafe {
            vulkano::pipeline::shader::ShaderModule::from_words(self.device.clone(), &self.shader.compute)
        }.unwrap();

        self.get_pipeline()
    }
}


#[repr(C)]
#[derive(Default, Debug, Clone)]
struct ComputeSpecializationConstants {
    first_constant: i32,
    second_constant: u32,
    third_constant: f32,
}

unsafe impl SpecializationConstants for ComputeSpecializationConstants {
    fn descriptors() -> &'static [SpecializationMapEntry] {
        static DESCRIPTORS: [SpecializationMapEntry; 3] = [
            SpecializationMapEntry {
                constant_id: 0,
                offset: 0,
                size: 4,
            },
            SpecializationMapEntry {
                constant_id: 1,
                offset: 4,
                size: 4,
            },
            SpecializationMapEntry {
                constant_id: 2,
                offset: 8,
                size: 4,
            },
        ];

        &DESCRIPTORS
    }
}
