medias 0.3.4

//! Graphics pipeline for the UI plane (milestone 3): render pass + pipeline +
//! vertex buffer, drawing into the dmabuf `Target`. SPIR-V is compiled offline
//! (see shaders/) and embedded, so there is no build-time or runtime shader
//! dependency.

use std::io::Cursor;

use ash::vk;

use super::{Gpu, Target};

const VERT_SPV: &[u8] = include_bytes!("../../shaders/ui.vert.spv");
const FRAG_SPV: &[u8] = include_bytes!("../../shaders/ui.frag.spv");

/// pos (vec2) + color (vec3), interleaved.
#[repr(C)]
#[derive(Clone, Copy)]
struct Vertex {
    pos: [f32; 2],
    color: [f32; 3],
}

/// Bottom overlay bar (like a control strip); the rest of the UI plane stays
/// transparent so the video plane shows through.
const BAR_TOP: f32 = 0.72; // NDC; bar spans y in [BAR_TOP, 1.0] (bottom in NDC)
const QUAD: [Vertex; 6] = [
    Vertex { pos: [-1.0, BAR_TOP], color: [0.10, 0.12, 0.18] },
    Vertex { pos: [1.0, BAR_TOP], color: [0.10, 0.12, 0.18] },
    Vertex { pos: [1.0, 1.0], color: [0.16, 0.42, 0.90] },
    Vertex { pos: [-1.0, BAR_TOP], color: [0.10, 0.12, 0.18] },
    Vertex { pos: [1.0, 1.0], color: [0.16, 0.42, 0.90] },
    Vertex { pos: [-1.0, 1.0], color: [0.16, 0.42, 0.90] },
];

pub struct Renderer {
    render_pass: vk::RenderPass,
    layout: vk::PipelineLayout,
    pipeline: vk::Pipeline,
    vert: vk::ShaderModule,
    frag: vk::ShaderModule,
    vbuf: vk::Buffer,
    vmem: vk::DeviceMemory,
}

impl Gpu {
    pub fn create_renderer(&self) -> Result<Renderer, String> {
        let d = &self.device;

        let color = vk::AttachmentDescription::default()
            .format(vk::Format::B8G8R8A8_UNORM)
            .samples(vk::SampleCountFlags::TYPE_1)
            .load_op(vk::AttachmentLoadOp::CLEAR)
            .store_op(vk::AttachmentStoreOp::STORE)
            .stencil_load_op(vk::AttachmentLoadOp::DONT_CARE)
            .stencil_store_op(vk::AttachmentStoreOp::DONT_CARE)
            .initial_layout(vk::ImageLayout::UNDEFINED)
            .final_layout(vk::ImageLayout::GENERAL); // ready for KMS scanout
        let color_ref = [vk::AttachmentReference::default()
            .attachment(0)
            .layout(vk::ImageLayout::COLOR_ATTACHMENT_OPTIMAL)];
        let subpass = [vk::SubpassDescription::default()
            .pipeline_bind_point(vk::PipelineBindPoint::GRAPHICS)
            .color_attachments(&color_ref)];
        let dep = [vk::SubpassDependency::default()
            .src_subpass(vk::SUBPASS_EXTERNAL)
            .dst_subpass(0)
            .src_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT)
            .dst_stage_mask(vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT)
            .dst_access_mask(vk::AccessFlags::COLOR_ATTACHMENT_WRITE)];
        let attachments = [color];
        let render_pass = unsafe {
            d.create_render_pass(
                &vk::RenderPassCreateInfo::default()
                    .attachments(&attachments)
                    .subpasses(&subpass)
                    .dependencies(&dep),
                None,
            )
        }
        .map_err(|e| format!("render pass: {e}"))?;

        let vert = shader_module(d, VERT_SPV)?;
        let frag = shader_module(d, FRAG_SPV)?;

        let push = [vk::PushConstantRange::default()
            .stage_flags(vk::ShaderStageFlags::VERTEX)
            .offset(0)
            .size(4)];
        let layout = unsafe {
            d.create_pipeline_layout(
                &vk::PipelineLayoutCreateInfo::default().push_constant_ranges(&push),
                None,
            )
        }
        .map_err(|e| format!("pipeline layout: {e}"))?;

        let pipeline = build_pipeline(d, render_pass, layout, vert, frag)?;

        let (vbuf, vmem) = self.vertex_buffer(&QUAD)?;

        Ok(Renderer { render_pass, layout, pipeline, vert, frag, vbuf, vmem })
    }

    /// A framebuffer binding `target`'s view to the render pass. Caller owns it.
    pub fn framebuffer(&self, r: &Renderer, t: &Target) -> Result<vk::Framebuffer, String> {
        let views = [t.view];
        unsafe {
            self.device.create_framebuffer(
                &vk::FramebufferCreateInfo::default()
                    .render_pass(r.render_pass)
                    .attachments(&views)
                    .width(t.width)
                    .height(t.height)
                    .layers(1),
                None,
            )
        }
        .map_err(|e| format!("framebuffer: {e}"))
    }

    /// Draw the quad into `target` and wait for completion (CPU fence; 3b moves
    /// to KMS OUT_FENCE). `t` animates the shader.
    pub fn draw(&self, r: &Renderer, target: &Target, fb: vk::Framebuffer, t: f32) -> Result<(), String> {
        let d = &self.device;
        let cmd = target.cmd;
        let extent = target.extent();
        unsafe {
            d.reset_command_pool(target.pool, vk::CommandPoolResetFlags::empty())
                .map_err(|e| format!("reset pool: {e}"))?;
            d.begin_command_buffer(
                cmd,
                &vk::CommandBufferBeginInfo::default()
                    .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
            )
            .map_err(|e| format!("begin: {e}"))?;

            // Transparent clear: everything outside the bar shows the video plane.
            let clear = [vk::ClearValue {
                color: vk::ClearColorValue { float32: [0.0, 0.0, 0.0, 0.0] },
            }];
            d.cmd_begin_render_pass(
                cmd,
                &vk::RenderPassBeginInfo::default()
                    .render_pass(r.render_pass)
                    .framebuffer(fb)
                    .render_area(vk::Rect2D { offset: vk::Offset2D { x: 0, y: 0 }, extent })
                    .clear_values(&clear),
                vk::SubpassContents::INLINE,
            );
            let viewport = [vk::Viewport {
                x: 0.0,
                y: 0.0,
                width: extent.width as f32,
                height: extent.height as f32,
                min_depth: 0.0,
                max_depth: 1.0,
            }];
            let scissor = [vk::Rect2D { offset: vk::Offset2D { x: 0, y: 0 }, extent }];
            d.cmd_set_viewport(cmd, 0, &viewport);
            d.cmd_set_scissor(cmd, 0, &scissor);
            d.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::GRAPHICS, r.pipeline);
            d.cmd_bind_vertex_buffers(cmd, 0, &[r.vbuf], &[0]);
            d.cmd_push_constants(
                cmd,
                r.layout,
                vk::ShaderStageFlags::VERTEX,
                0,
                &t.to_ne_bytes(),
            );
            d.cmd_draw(cmd, QUAD.len() as u32, 1, 0, 0);
            d.cmd_end_render_pass(cmd);
            d.end_command_buffer(cmd).map_err(|e| format!("end: {e}"))?;

            let cmds = [cmd];
            let submit = vk::SubmitInfo::default().command_buffers(&cmds);
            d.queue_submit(self.queue, &[submit], target.fence)
                .map_err(|e| format!("submit: {e}"))?;
            d.wait_for_fences(&[target.fence], true, u64::MAX)
                .map_err(|e| format!("wait: {e}"))?;
            d.reset_fences(&[target.fence]).map_err(|e| format!("reset fence: {e}"))?;
        }
        Ok(())
    }

    pub fn destroy_renderer(&self, r: Renderer) {
        let d = &self.device;
        unsafe {
            d.destroy_buffer(r.vbuf, None);
            d.free_memory(r.vmem, None);
            d.destroy_pipeline(r.pipeline, None);
            d.destroy_pipeline_layout(r.layout, None);
            d.destroy_shader_module(r.vert, None);
            d.destroy_shader_module(r.frag, None);
            d.destroy_render_pass(r.render_pass, None);
        }
    }

    pub fn destroy_framebuffer(&self, fb: vk::Framebuffer) {
        unsafe { self.device.destroy_framebuffer(fb, None) };
    }

    /// Host-visible vertex buffer filled with `verts`.
    fn vertex_buffer(&self, verts: &[Vertex]) -> Result<(vk::Buffer, vk::DeviceMemory), String> {
        let d = &self.device;
        let size = std::mem::size_of_val(verts) as u64;
        let buf = unsafe {
            d.create_buffer(
                &vk::BufferCreateInfo::default()
                    .size(size)
                    .usage(vk::BufferUsageFlags::VERTEX_BUFFER)
                    .sharing_mode(vk::SharingMode::EXCLUSIVE),
                None,
            )
        }
        .map_err(|e| format!("vbuf: {e}"))?;
        let req = unsafe { d.get_buffer_memory_requirements(buf) };
        let mem_type = self
            .find_memory_type_with(
                req.memory_type_bits,
                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
            )
            .ok_or("no host-visible memory type")?;
        let mem = unsafe {
            d.allocate_memory(
                &vk::MemoryAllocateInfo::default()
                    .allocation_size(req.size)
                    .memory_type_index(mem_type),
                None,
            )
        }
        .map_err(|e| format!("vmem: {e}"))?;
        unsafe {
            d.bind_buffer_memory(buf, mem, 0).map_err(|e| format!("bind vbuf: {e}"))?;
            let ptr = d
                .map_memory(mem, 0, size, vk::MemoryMapFlags::empty())
                .map_err(|e| format!("map: {e}"))? as *mut Vertex;
            ptr.copy_from_nonoverlapping(verts.as_ptr(), verts.len());
            d.unmap_memory(mem);
        }
        Ok((buf, mem))
    }

    pub(crate) fn find_memory_type_with(&self, bits: u32, flags: vk::MemoryPropertyFlags) -> Option<u32> {
        let props = unsafe { self.instance.get_physical_device_memory_properties(self.physical) };
        (0..props.memory_type_count).find(|&i| {
            bits & (1 << i) != 0 && props.memory_types[i as usize].property_flags.contains(flags)
        })
    }
}

fn shader_module(d: &ash::Device, spv: &[u8]) -> Result<vk::ShaderModule, String> {
    let code = ash::util::read_spv(&mut Cursor::new(spv)).map_err(|e| format!("spv: {e}"))?;
    unsafe { d.create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&code), None) }
        .map_err(|e| format!("shader module: {e}"))
}

fn build_pipeline(
    d: &ash::Device,
    render_pass: vk::RenderPass,
    layout: vk::PipelineLayout,
    vert: vk::ShaderModule,
    frag: vk::ShaderModule,
) -> Result<vk::Pipeline, String> {
    let entry = c"main";
    let stages = [
        vk::PipelineShaderStageCreateInfo::default()
            .stage(vk::ShaderStageFlags::VERTEX)
            .module(vert)
            .name(entry),
        vk::PipelineShaderStageCreateInfo::default()
            .stage(vk::ShaderStageFlags::FRAGMENT)
            .module(frag)
            .name(entry),
    ];
    let bindings = [vk::VertexInputBindingDescription::default()
        .binding(0)
        .stride(std::mem::size_of::<Vertex>() as u32)
        .input_rate(vk::VertexInputRate::VERTEX)];
    let attrs = [
        vk::VertexInputAttributeDescription::default()
            .location(0)
            .binding(0)
            .format(vk::Format::R32G32_SFLOAT)
            .offset(0),
        vk::VertexInputAttributeDescription::default()
            .location(1)
            .binding(0)
            .format(vk::Format::R32G32B32_SFLOAT)
            .offset(8),
    ];
    let vertex_input = vk::PipelineVertexInputStateCreateInfo::default()
        .vertex_binding_descriptions(&bindings)
        .vertex_attribute_descriptions(&attrs);
    let input_assembly = vk::PipelineInputAssemblyStateCreateInfo::default()
        .topology(vk::PrimitiveTopology::TRIANGLE_LIST);
    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)
        .front_face(vk::FrontFace::COUNTER_CLOCKWISE)
        .line_width(1.0);
    let multisample = vk::PipelineMultisampleStateCreateInfo::default()
        .rasterization_samples(vk::SampleCountFlags::TYPE_1);
    let blend_attach = [vk::PipelineColorBlendAttachmentState::default()
        .color_write_mask(vk::ColorComponentFlags::RGBA)
        .blend_enable(false)];
    let blend = vk::PipelineColorBlendStateCreateInfo::default().attachments(&blend_attach);
    let dyn_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
    let dynamic = vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dyn_states);

    let info = [vk::GraphicsPipelineCreateInfo::default()
        .stages(&stages)
        .vertex_input_state(&vertex_input)
        .input_assembly_state(&input_assembly)
        .viewport_state(&viewport)
        .rasterization_state(&raster)
        .multisample_state(&multisample)
        .color_blend_state(&blend)
        .dynamic_state(&dynamic)
        .layout(layout)
        .render_pass(render_pass)
        .subpass(0)];
    unsafe { d.create_graphics_pipelines(vk::PipelineCache::null(), &info, None) }
        .map_err(|(_, e)| format!("graphics pipeline: {e}"))
        .map(|v| v[0])
}