use vulkano::device::{Device};
use vulkano::pipeline::{ComputePipeline};
use std::sync::Arc;
use std::ffi::CStr;
use std::path::PathBuf;
use shade_runner as sr;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use shade_runner::{CompileError, Layout, Input, Output, CompiledShaders, Entry, CompiledShader};
use shaderc::CompileOptions;
use vulkano::pipeline::shader::{ShaderModule, GraphicsEntryPoint, SpecializationConstants, SpecializationMapEntry};
use crate::compute::compu_buffer::{CompuBuffers, CompuBufferHandle};

#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct CompuKernelHandle {
    pub handle: u32,
}

#[derive(Clone)]
pub struct CompuKernel {

    handle: Arc<CompuKernelHandle>,

    compute_pipeline: Option<std::sync::Arc<ComputePipeline<PipelineLayout<Layout>>>>,
    compute_kernel_path: PathBuf,

    name: String,

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

impl CompuKernel {

    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>,
               handle: Arc<CompuKernelHandle>) -> CompuKernel {

        let compute_path = CompuKernel::get_path(filename.clone());

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


        CompuKernel {
            name: filename,
            handle: handle,
            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<Layout>>> {

        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.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::Layout>>>  {
        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<Layout>>> {

        let mut options = CompileOptions::new().ok_or(CompileError::CreateCompiler).unwrap();
        self.compute_kernel_path = CompuKernel::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.spriv)
        }.unwrap();

        self.get_pipeline()
    }

    pub fn get_handle(&self) -> Arc<CompuKernelHandle> {
        self.handle.clone()
    }
    pub fn get_name(&self) -> String {
        self.name.clone()
    }
}

#[repr(C)]
#[derive(Default, Debug, Clone)]
pub 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
    }
}
