monolith

use crate::Error;

use crate::device::Shared;

use crate::format::{SUPPORTED, vk_format};

use ash::vk;

use std::collections::HashMap;

use std::path::PathBuf;

use std::sync::Arc;

use std::sync::atomic::{AtomicU32, Ordering};

use std::time::Instant;

/// 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],

    /// Transition progress and kind, read only by `shaders/effect.frag`.

    pub effect: [f32; 4],

    /// See `shaders/clip.glsl`. A zero width turns clipping off.

    pub clip: [f32; 4],

    pub radii: [f32; 4],

    pub ring: [f32; 4],

}

impl Push {

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

        // SAFETY: `Push` is `repr(C)` and all f32, so it has no padding bytes.

        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,

    Effect,

}

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,

    effect: vk::ShaderModule,

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

    vk_cache: vk::PipelineCache,

    cache_path: Option<PathBuf>,

    loaded_len: usize,

}

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"));

const EFFECT_FRAG: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/effect.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)?

            })

        };

        let cache_path = cache_path();

        let props = unsafe {

            shared

                .device

                .instance

                .get_physical_device_properties(shared.device.physical)

        };

        let loaded = cache_path

            .as_ref()

            .and_then(|p| std::fs::read(p).ok())

            .filter(|data| header_matches(data, &props))

            .unwrap_or_default();

        let create = |data: &[u8]| unsafe {

            d.create_pipeline_cache(

                &vk::PipelineCacheCreateInfo::default().initial_data(data),

                None,

            )

        };

        let vk_cache = create(&loaded).or_else(|_| create(&[])).unwrap_or_default();

        Ok(Self {

            shared: shared.clone(),

            set_layout,

            layout,

            vert: module(QUAD_VERT)?,

            solid: module(SOLID_FRAG)?,

            texture: module(TEXTURE_FRAG)?,

            effect: module(EFFECT_FRAG)?,

            cache: HashMap::new(),

            vk_cache,

            cache_path,

            loaded_len: loaded.len(),

        })

    }

    /// Compiles the pipelines a frame draws with for every render format, so

    /// the first spawn animation does not stall on the driver compiler.

    pub fn warm(&mut self) {

        let began = Instant::now();

        for fourcc in SUPPORTED {

            let Some(format) = vk_format(fourcc) else {

                continue;

            };

            for kind in [Kind::Solid, Kind::Texture, Kind::Effect] {

                if let Err(e) = self.get(format.format, kind) {

                    tracing::warn!(?kind, "cannot warm a pipeline: {e}");

                    return;

                }

            }

        }

        tracing::debug!(elapsed = ?began.elapsed(), "warmed pipelines");

    }

    fn save_cache(&self) {

        let Some(path) = &self.cache_path else {

            return;

        };

        if self.vk_cache == vk::PipelineCache::null() {

            return;

        }

        let Ok(data) = (unsafe {

            self.shared

                .device

                .device

                .get_pipeline_cache_data(self.vk_cache)

        }) else {

            return;

        };

        if data.len() == self.loaded_len {

            return;

        }

        static WRITERS: AtomicU32 = AtomicU32::new(0);

        let tmp = path.with_extension(format!(

            "{}.{}.tmp",

            std::process::id(),

            WRITERS.fetch_add(1, Ordering::Relaxed)

        ));

        let written = path

            .parent()

            .map_or(Ok(()), std::fs::create_dir_all)

            .and_then(|()| std::fs::write(&tmp, &data))

            .and_then(|()| std::fs::rename(&tmp, path));

        if let Err(e) = written {

            tracing::debug!("cannot save the pipeline cache: {e}");

            let _ = std::fs::remove_file(&tmp);

        }

    }

    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,

            Kind::Effect => self.effect,

        };

        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(self.vk_cache, &[info], None) }

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

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

        Ok(pipeline)

    }

}

impl Drop for Pipelines {

    fn drop(&mut self) {

        self.save_cache();

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

        unsafe {

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

                d.destroy_pipeline(p, None);

            }

            d.destroy_pipeline_cache(self.vk_cache, None);

            d.destroy_shader_module(self.vert, None);

            d.destroy_shader_module(self.solid, None);

            d.destroy_shader_module(self.texture, None);

            d.destroy_shader_module(self.effect, None);

            d.destroy_pipeline_layout(self.layout, None);

            d.destroy_descriptor_set_layout(self.set_layout, None);

        }

    }

}

fn cache_path() -> Option<PathBuf> {

    let base = std::env::var_os("XDG_CACHE_HOME")

        .filter(|v| !v.is_empty())

        .map(PathBuf::from)

        .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".cache")))?;

    Some(base.join("monolith").join("pipelines.bin"))

}

/// Whether `data` opens with a cache header written by this driver on this

/// device. The spec lets a driver reject a foreign cache but some crash on

/// one instead, so the header is checked first.

fn header_matches(data: &[u8], props: &vk::PhysicalDeviceProperties) -> bool {

    let word = |at: usize| {

        data.get(at..at + 4)

            .and_then(|b| b.try_into().ok())

            .map(u32::from_ne_bytes)

    };

    word(0).is_some_and(|size| size >= 32 && size as usize <= data.len())

        && word(4) == Some(vk::PipelineCacheHeaderVersion::ONE.as_raw() as u32)

        && word(8) == Some(props.vendor_id)

        && word(12) == Some(props.device_id)

        && data.get(16..32) == Some(&props.pipeline_cache_uuid[..])

}

#[cfg(test)]

mod tests {

    use super::*;

    fn props() -> vk::PhysicalDeviceProperties {

        vk::PhysicalDeviceProperties {

            vendor_id: 0x10de,

            device_id: 0x2487,

            pipeline_cache_uuid: [7; 16],

            ..Default::default()

        }

    }

    fn header(vendor: u32, uuid: u8) -> Vec<u8> {

        let words = [32u32, 1, vendor, 0x2487];

        let mut data: Vec<u8> = words.iter().flat_map(|w| w.to_ne_bytes()).collect();

        data.extend([uuid; 16]);

        data

    }

    #[test]

    fn cache_header_must_match_the_device_and_driver() {

        assert!(header_matches(&header(0x10de, 7), &props()));

        assert!(!header_matches(&header(0x1002, 7), &props()));

        assert!(!header_matches(&header(0x10de, 8), &props()));

        assert!(!header_matches(&header(0x10de, 7)[..20], &props()));

        assert!(!header_matches(&[], &props()));

    }

}