monolith

use crate::Error;

use crate::device::Shared;

use ash::vk;

use std::collections::HashMap;

use std::sync::Arc;

/// Push constants shared by every shader, see `shaders/push.glsl`.

/// `dst` is the quad in normalized device coordinates as (x, y, w, h),

/// `uv` the texture coordinate at each corner (TL, TR, BL, BR).

#[repr(C)]

#[derive(Clone, Copy, Debug, Default)]

pub(crate) struct Push {

    pub color: [f32; 4],

    pub dst: [f32; 4],

    pub uv: [[f32; 2]; 4],

}

impl Push {

    pub fn bytes(&self) -> &[u8] {

        unsafe { std::slice::from_raw_parts((self as *const Self).cast::<u8>(), size_of::<Self>()) }

    }

}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]

pub(crate) enum Kind {

    Solid,

    Texture,

}

pub(crate) struct Pipelines {

    shared: Arc<Shared>,

    pub set_layout: vk::DescriptorSetLayout,

    pub layout: vk::PipelineLayout,

    vert: vk::ShaderModule,

    solid: vk::ShaderModule,

    texture: vk::ShaderModule,

    cache: HashMap<(vk::Format, Kind), vk::Pipeline>,

}

const QUAD_VERT: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/quad.vert.spv"));

const SOLID_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/solid.frag.spv"));

const TEXTURE_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/texture.frag.spv"));

impl Pipelines {

    pub fn new(shared: &Arc<Shared>) -> Result<Self, Error> {

        let d = &shared.device.device;

        let bindings = [vk::DescriptorSetLayoutBinding::default()

            .binding(0)

            .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)

            .descriptor_count(1)

            .stage_flags(vk::ShaderStageFlags::FRAGMENT)];

        let set_layout = unsafe {

            d.create_descriptor_set_layout(

                &vk::DescriptorSetLayoutCreateInfo::default()

                    .flags(vk::DescriptorSetLayoutCreateFlags::PUSH_DESCRIPTOR_KHR)

                    .bindings(&bindings),

                None,

            )?

        };

        let ranges = [vk::PushConstantRange::default()

            .stage_flags(vk::ShaderStageFlags::VERTEX | vk::ShaderStageFlags::FRAGMENT)

            .size(size_of::<Push>() as u32)];

        let set_layouts = [set_layout];

        let layout = unsafe {

            d.create_pipeline_layout(

                &vk::PipelineLayoutCreateInfo::default()

                    .set_layouts(&set_layouts)

                    .push_constant_ranges(&ranges),

                None,

            )?

        };

        let module = |code: &[u8]| -> Result<vk::ShaderModule, Error> {

            let words = ash::util::read_spv(&mut std::io::Cursor::new(code))

                .map_err(|_| vk::Result::ERROR_INVALID_SHADER_NV)?;

            Ok(unsafe {

                d.create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&words), None)?

            })

        };

        Ok(Self {

            shared: shared.clone(),

            set_layout,

            layout,

            vert: module(QUAD_VERT)?,

            solid: module(SOLID_FRAG)?,

            texture: module(TEXTURE_FRAG)?,

            cache: HashMap::new(),

        })

    }

    pub fn get(&mut self, format: vk::Format, kind: Kind) -> Result<vk::Pipeline, Error> {

        if let Some(&p) = self.cache.get(&(format, kind)) {

            return Ok(p);

        }

        let d = &self.shared.device.device;

        let frag = match kind {

            Kind::Solid => self.solid,

            Kind::Texture => self.texture,

        };

        let stages = [

            vk::PipelineShaderStageCreateInfo::default()

                .stage(vk::ShaderStageFlags::VERTEX)

                .module(self.vert)

                .name(c"main"),

            vk::PipelineShaderStageCreateInfo::default()

                .stage(vk::ShaderStageFlags::FRAGMENT)

                .module(frag)

                .name(c"main"),

        ];

        let vertex = vk::PipelineVertexInputStateCreateInfo::default();

        let assembly = vk::PipelineInputAssemblyStateCreateInfo::default()

            .topology(vk::PrimitiveTopology::TRIANGLE_STRIP);

        let viewport = vk::PipelineViewportStateCreateInfo::default()

            .viewport_count(1)

            .scissor_count(1);

        let raster = vk::PipelineRasterizationStateCreateInfo::default()

            .polygon_mode(vk::PolygonMode::FILL)

            .cull_mode(vk::CullModeFlags::NONE)

            .line_width(1.0);

        let multisample = vk::PipelineMultisampleStateCreateInfo::default()

            .rasterization_samples(vk::SampleCountFlags::TYPE_1);

        let attachments = [vk::PipelineColorBlendAttachmentState::default()

            .blend_enable(true)

            .src_color_blend_factor(vk::BlendFactor::ONE)

            .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)

            .color_blend_op(vk::BlendOp::ADD)

            .src_alpha_blend_factor(vk::BlendFactor::ONE)

            .dst_alpha_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)

            .alpha_blend_op(vk::BlendOp::ADD)

            .color_write_mask(vk::ColorComponentFlags::RGBA)];

        let blend = vk::PipelineColorBlendStateCreateInfo::default().attachments(&attachments);

        let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];

        let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);

        let formats = [format];

        let mut rendering =

            vk::PipelineRenderingCreateInfo::default().color_attachment_formats(&formats);

        let info = vk::GraphicsPipelineCreateInfo::default()

            .stages(&stages)

            .vertex_input_state(&vertex)

            .input_assembly_state(&assembly)

            .viewport_state(&viewport)

            .rasterization_state(&raster)

            .multisample_state(&multisample)

            .color_blend_state(&blend)

            .dynamic_state(&dynamic)

            .layout(self.layout)

            .push_next(&mut rendering);

        let pipeline =

            unsafe { d.create_graphics_pipelines(vk::PipelineCache::null(), &[info], None) }

                .map_err(|(_, e)| e)?[0];

        self.cache.insert((format, kind), pipeline);

        Ok(pipeline)

    }

}

impl Drop for Pipelines {

    fn drop(&mut self) {

        let d = &self.shared.device.device;

        unsafe {

            for (_, p) in self.cache.drain() {

                d.destroy_pipeline(p, None);

            }

            d.destroy_shader_module(self.vert, None);

            d.destroy_shader_module(self.solid, None);

            d.destroy_shader_module(self.texture, None);

            d.destroy_pipeline_layout(self.layout, None);

            d.destroy_descriptor_set_layout(self.set_layout, None);

        }

    }

}