path: root/src/platform/vulkan/vulkan_common.h

#pragma once

// Shared internals of the Vulkan RHI backend: the enum-mapping helpers, the
// opaque resource classes (Buffer/Texture/RenderTarget/Pipeline/CommandList), and
// the VulkanDevice declaration. The implementation is split across vulkan_device
// (lifecycle + frame loop), vulkan_swapchain (instance/device/swapchain/passes),
// vulkan_resources (buffer/texture/target/pipeline) and vulkan_cubemap. This
// header is private to platform/vulkan/ — nothing outside includes it.

#include "vulkan_device.h"   // rendering/rhi.h (base classes + factory decls)

#define GLFW_INCLUDE_VULKAN
#include <GLFW/glfw3.h>
#include <glm/glm.hpp>

#include <vector>
#include <array>
#include <string>
#include <unordered_map>
#include <cstdint>
#include <cstring>
#include <algorithm>

#include "core/log.h"

namespace Donut::RHI
{
    constexpr int MAX_FRAMES_IN_FLIGHT = 2;
    constexpr uint32_t MAX_BINDINGS = 8;

    #define VKD_CHECK(expr)                                            \
        do {                                                           \
            VkResult _r = (expr);                                      \
            if (_r != VK_SUCCESS) {                                    \
                DONUT_ERROR("Vulkan RHI: {} failed ({})", #expr, (int)_r); \
                return false;                                          \
            }                                                          \
        } while (0)

    inline auto vk_format(Format f) -> VkFormat
    {
        switch (f) {
            case Format::RGBA16F: return VK_FORMAT_R16G16B16A16_SFLOAT;
            case Format::D32:     return VK_FORMAT_D32_SFLOAT;
            default:              return VK_FORMAT_R8G8B8A8_UNORM;
        }
    }
    inline auto vk_attr_format(uint32_t comps) -> VkFormat
    {
        switch (comps) {
            case 1:  return VK_FORMAT_R32_SFLOAT;
            case 2:  return VK_FORMAT_R32G32_SFLOAT;
            case 3:  return VK_FORMAT_R32G32B32_SFLOAT;
            default: return VK_FORMAT_R32G32B32A32_SFLOAT;
        }
    }
    inline auto vk_topology(Topology t) -> VkPrimitiveTopology
    { return t == Topology::Lines ? VK_PRIMITIVE_TOPOLOGY_LINE_LIST : VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST; }
    inline auto vk_compare(CompareOp o) -> VkCompareOp
    { return o == CompareOp::Always ? VK_COMPARE_OP_ALWAYS : o == CompareOp::LessEqual ? VK_COMPARE_OP_LESS_OR_EQUAL : VK_COMPARE_OP_LESS; }
    inline auto vk_filter(Filter f) -> VkFilter { return f == Filter::Nearest ? VK_FILTER_NEAREST : VK_FILTER_LINEAR; }
    inline auto vk_cull(CullMode c) -> VkCullModeFlags
    { return c == CullMode::None ? VK_CULL_MODE_NONE : c == CullMode::Back ? VK_CULL_MODE_BACK_BIT : VK_CULL_MODE_FRONT_BIT; }

    // Buffer: host-visible + coherent, persistently mapped.
    class VkBufferR : public Buffer
    {
    public:
        VkBufferR(VkDevice d, VkBuffer b, VkDeviceMemory m, void* mapped, size_t size)
            : m_device(d), m_buf(b), m_mem(m), m_mapped(mapped), m_size(size) {}
        ~VkBufferR() override
        {
            if (m_mapped) vkUnmapMemory(m_device, m_mem);
            if (m_buf) vkDestroyBuffer(m_device, m_buf, nullptr);
            if (m_mem) vkFreeMemory(m_device, m_mem, nullptr);
        }
        auto update(const void* data, size_t size) -> void override
        { if (m_mapped) std::memcpy(m_mapped, data, std::min(size, m_size)); }

        VkDevice m_device; VkBuffer m_buf; VkDeviceMemory m_mem; void* m_mapped; size_t m_size;
    };

    // Texture: sampled image (2D or cube). Owns its handles unless it is a borrowed
    // wrapper around a render-target view.
    class VkTextureR : public Texture
    {
    public:
        VkTextureR() = default;
        ~VkTextureR() override
        {
            if (!m_owns) return;
            if (m_sampler) vkDestroySampler(m_device, m_sampler, nullptr);
            if (m_view)    vkDestroyImageView(m_device, m_view, nullptr);
            if (m_image)   vkDestroyImage(m_device, m_image, nullptr);
            if (m_mem)     vkFreeMemory(m_device, m_mem, nullptr);
        }
        VkDevice      m_device  = VK_NULL_HANDLE;
        VkImage       m_image   = VK_NULL_HANDLE;
        VkDeviceMemory m_mem    = VK_NULL_HANDLE;
        VkImageView   m_view    = VK_NULL_HANDLE;
        VkSampler     m_sampler = VK_NULL_HANDLE;
        bool          m_owns    = true;
    };

    // RenderTarget: off-screen colour image (+ optional depth) + framebuffer.
    class VkRenderTargetR : public RenderTarget
    {
    public:
        ~VkRenderTargetR() override
        {
            if (m_fb)          vkDestroyFramebuffer(m_device, m_fb, nullptr);
            if (m_sampler)     vkDestroySampler(m_device, m_sampler, nullptr);
            if (m_view)        vkDestroyImageView(m_device, m_view, nullptr);
            if (m_image)       vkDestroyImage(m_device, m_image, nullptr);
            if (m_mem)         vkFreeMemory(m_device, m_mem, nullptr);
            if (m_depth_view)  vkDestroyImageView(m_device, m_depth_view, nullptr);
            if (m_depth_image) vkDestroyImage(m_device, m_depth_image, nullptr);
            if (m_depth_mem)   vkFreeMemory(m_device, m_depth_mem, nullptr);
            // m_pass is owned by the device's render-pass cache, not by us.
        }
        auto width()  const -> int override { return m_w; }
        auto height() const -> int override { return m_h; }
        auto color_texture() -> Texture* override { return &m_color; }

        VkDevice       m_device  = VK_NULL_HANDLE;
        int            m_w = 0, m_h = 0;
        VkImage        m_image   = VK_NULL_HANDLE;
        VkDeviceMemory m_mem     = VK_NULL_HANDLE;
        VkImageView    m_view    = VK_NULL_HANDLE;
        VkSampler      m_sampler = VK_NULL_HANDLE;
        VkFramebuffer  m_fb      = VK_NULL_HANDLE;
        VkRenderPass   m_pass    = VK_NULL_HANDLE;   // borrowed (device pass cache)
        bool           m_has_depth = false;
        VkImage        m_depth_image = VK_NULL_HANDLE;
        VkDeviceMemory m_depth_mem   = VK_NULL_HANDLE;
        VkImageView    m_depth_view  = VK_NULL_HANDLE;
        VkTextureR     m_color;   // borrowed wrapper (view+sampler) for sampling
    };

    // Pipeline: the graphics pipeline + its (set-0) descriptor layout.
    class VkPipelineR : public Pipeline
    {
    public:
        ~VkPipelineR() override
        {
            if (m_pipeline)   vkDestroyPipeline(m_device, m_pipeline, nullptr);
            if (m_layout)     vkDestroyPipelineLayout(m_device, m_layout, nullptr);
            if (m_set_layout) vkDestroyDescriptorSetLayout(m_device, m_set_layout, nullptr);
        }
        VkDevice              m_device      = VK_NULL_HANDLE;
        VkPipeline            m_pipeline    = VK_NULL_HANDLE;
        VkPipelineLayout      m_layout      = VK_NULL_HANDLE;
        VkDescriptorSetLayout m_set_layout  = VK_NULL_HANDLE;
        std::vector<ResourceSlot> m_resources;
    };

    // CommandList: records into the frame's command buffer. Per-draw descriptor
    // sets are allocated from a per-frame pool (no push-descriptor extension).
    class VkCommandListR : public CommandList
    {
    public:
        auto begin_render_pass(RenderTarget* target, const glm::vec4& clear) -> void override
        {
            VkClearValue cvs[2]{};
            cvs[0].color = { { clear.r, clear.g, clear.b, clear.a } };
            cvs[1].depthStencil = { 1.0f, 0 };
            VkRenderPassBeginInfo rpbi{ VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO };
            if (target)
            {
                auto* rt = static_cast<VkRenderTargetR*>(target);
                rpbi.renderPass = rt->m_pass; rpbi.framebuffer = rt->m_fb;
                rpbi.renderArea = { { 0, 0 }, { (uint32_t)rt->m_w, (uint32_t)rt->m_h } };
                rpbi.clearValueCount = rt->m_has_depth ? 2 : 1; rpbi.pClearValues = cvs;
            }
            else
            {
                rpbi.renderPass = m_swapchain_rp; rpbi.framebuffer = m_swapchain_fb;
                rpbi.renderArea = { { 0, 0 }, m_extent };
                rpbi.clearValueCount = 2; rpbi.pClearValues = cvs;
            }
            vkCmdBeginRenderPass(m_cmd, &rpbi, VK_SUBPASS_CONTENTS_INLINE);
        }
        auto end_render_pass() -> void override { vkCmdEndRenderPass(m_cmd); }

        auto bind_pipeline(Pipeline* p) -> void override
        {
            m_pipe = static_cast<VkPipelineR*>(p);
            vkCmdBindPipeline(m_cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipe->m_pipeline);
        }
        auto set_viewport(int x, int y, int w, int h, bool flip_y) -> void override
        {
            VkViewport vp{ (float)x, flip_y ? (float)(y + h) : (float)y,
                           (float)w, flip_y ? -(float)h : (float)h, 0.0f, 1.0f };
            VkRect2D sc{ { x, y }, { (uint32_t)w, (uint32_t)h } };
            vkCmdSetViewport(m_cmd, 0, 1, &vp);
            vkCmdSetScissor(m_cmd, 0, 1, &sc);
        }
        auto bind_uniform(uint32_t binding, Buffer* ubo) -> void override
        {
            if (binding >= MAX_BINDINGS) return;
            m_buf_info[binding] = { static_cast<VkBufferR*>(ubo)->m_buf, 0, VK_WHOLE_SIZE };
        }
        auto bind_texture(uint32_t binding, Texture* texture) -> void override
        {
            if (binding >= MAX_BINDINGS) return;
            auto* t = static_cast<VkTextureR*>(texture);
            m_img_info[binding] = { t->m_sampler, t->m_view, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL };
        }
        auto bind_vertex_buffer(Buffer* vb) -> void override
        {
            VkBuffer b = static_cast<VkBufferR*>(vb)->m_buf; VkDeviceSize off = 0;
            vkCmdBindVertexBuffers(m_cmd, 0, 1, &b, &off);
        }
        auto bind_index_buffer(Buffer* ib) -> void override
        { vkCmdBindIndexBuffer(m_cmd, static_cast<VkBufferR*>(ib)->m_buf, 0, VK_INDEX_TYPE_UINT32); }
        auto draw(uint32_t vertex_count) -> void override
        { flush_descriptors(); vkCmdDraw(m_cmd, vertex_count, 1, 0, 0); }
        auto draw_indexed(uint32_t index_count) -> void override
        { flush_descriptors(); vkCmdDrawIndexed(m_cmd, index_count, 1, 0, 0, 0); }

        // Allocate + write + bind a descriptor set for the current pipeline's
        // declared resources, using whatever was bound since bind_pipeline.
        auto flush_descriptors() -> void
        {
            if (!m_pipe || m_pipe->m_resources.empty()) return;
            VkDescriptorSetAllocateInfo ai{ VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO };
            ai.descriptorPool = m_frame_pool; ai.descriptorSetCount = 1; ai.pSetLayouts = &m_pipe->m_set_layout;
            VkDescriptorSet set = VK_NULL_HANDLE;
            if (vkAllocateDescriptorSets(m_device, &ai, &set) != VK_SUCCESS)
            { DONUT_ERROR("Vulkan RHI: descriptor set allocation failed"); return; }

            std::array<VkWriteDescriptorSet, MAX_BINDINGS> writes{};
            uint32_t n = 0;
            for (const auto& r : m_pipe->m_resources)
            {
                VkWriteDescriptorSet& w = writes[n++];
                w.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
                w.dstSet = set; w.dstBinding = r.binding; w.descriptorCount = 1;
                if (r.kind == ResourceKind::UniformBuffer)
                { w.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; w.pBufferInfo = &m_buf_info[r.binding]; }
                else
                { w.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER; w.pImageInfo = &m_img_info[r.binding]; }
            }
            vkUpdateDescriptorSets(m_device, n, writes.data(), 0, nullptr);
            vkCmdBindDescriptorSets(m_cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipe->m_layout, 0, 1, &set, 0, nullptr);
        }

        // Set by the device at begin_frame:
        VkDevice        m_device       = VK_NULL_HANDLE;
        VkCommandBuffer m_cmd          = VK_NULL_HANDLE;
        VkRenderPass    m_swapchain_rp = VK_NULL_HANDLE;
        VkFramebuffer   m_swapchain_fb = VK_NULL_HANDLE;
        VkExtent2D      m_extent{};
        VkDescriptorPool m_frame_pool  = VK_NULL_HANDLE;

        VkPipelineR* m_pipe = nullptr;
        VkDescriptorBufferInfo m_buf_info[MAX_BINDINGS]{};
        VkDescriptorImageInfo  m_img_info[MAX_BINDINGS]{};
    };

    // The backend device. Its methods are defined across vulkan_device.cpp,
    // vulkan_swapchain.cpp, vulkan_resources.cpp and vulkan_cubemap.cpp.
    class VulkanDevice : public Device
    {
    public:
        auto init(const NativeWindow& window) -> bool override;
        auto shutdown() -> void override;
        auto resize(int width, int height) -> void override { m_framebuffer_resized = true; m_width = width; m_height = height; }
        auto wait_idle() -> void override { if (m_device) vkDeviceWaitIdle(m_device); }

        auto create_buffer(BufferType type, size_t size, const void* data) -> Ref<Buffer> override;
        auto create_texture(int w, int h, Format format, Filter filter, const void* data) -> Ref<Texture> override;
        auto create_cubemap_from_hdri(const std::string& path) -> Ref<Texture> override;
        auto create_render_target(int w, int h, Format color, Format depth, Filter filter, int mips) -> Ref<RenderTarget> override;
        auto create_pipeline(const PipelineDesc& desc) -> Ref<Pipeline> override;

        auto begin_frame(const glm::vec4& clear) -> CommandList* override;
        auto end_frame() -> void override;

        auto init_imgui() -> void override;
        auto imgui_new_frame() -> void override;
        auto imgui_render(CommandList& cmds) -> void override;

        auto device_name() const -> const std::string& override { return m_gpu_name; }

        // internals (implemented across the vulkan_*.cpp files)
        auto find_memory_type(uint32_t filter, VkMemoryPropertyFlags flags) const -> uint32_t;
        auto create_buffer_raw(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags props, VkBuffer& buf, VkDeviceMemory& mem) const -> bool;
        auto load_spirv(const std::string& path) const -> std::vector<uint32_t>;
        auto create_shader_module(const std::string& path, VkShaderModule& out) const -> bool;

        auto create_instance() -> bool;
        auto pick_physical_and_device() -> bool;
        auto create_swapchain() -> bool;
        auto create_image_views() -> bool;
        // Render passes are format-driven and cached: a pipeline/target's
        // attachment signature (colour + optional depth, present vs sampled)
        // maps to one pass. `present` = presented swapchain image.
        auto get_render_pass(VkFormat color, VkFormat depth, bool present) -> VkRenderPass;
        auto create_depth_and_framebuffers() -> bool;
        auto create_command_and_sync() -> bool;
        auto recreate_swapchain() -> bool;
        auto cleanup_swapchain() -> void;

        GLFWwindow* m_window = nullptr;
        int m_width = 0, m_height = 0;
        bool m_framebuffer_resized = false;
        std::string m_gpu_name;

        VkInstance       m_instance = VK_NULL_HANDLE;
        VkSurfaceKHR     m_surface  = VK_NULL_HANDLE;
        VkPhysicalDevice m_physical = VK_NULL_HANDLE;
        VkDevice         m_device   = VK_NULL_HANDLE;
        uint32_t         m_graphics_family = 0, m_present_family = 0;
        VkQueue          m_graphics_queue = VK_NULL_HANDLE, m_present_queue = VK_NULL_HANDLE;
        VkPhysicalDeviceMemoryProperties m_mem_props{};

        VkSwapchainKHR m_swapchain = VK_NULL_HANDLE;
        VkFormat       m_swapchain_format = VK_FORMAT_B8G8R8A8_UNORM;
        VkExtent2D     m_extent{};
        std::vector<VkImage>       m_images;
        std::vector<VkImageView>   m_image_views;
        VkRenderPass               m_swapchain_rp = VK_NULL_HANDLE;   // also lives in m_pass_cache
        std::unordered_map<uint64_t, VkRenderPass> m_pass_cache;      // keyed by (color,depth,present)
        std::vector<VkFramebuffer> m_framebuffers;
        VkImage m_depth_image = VK_NULL_HANDLE; VkDeviceMemory m_depth_mem = VK_NULL_HANDLE; VkImageView m_depth_view = VK_NULL_HANDLE;

        VkCommandPool m_command_pool = VK_NULL_HANDLE;
        std::vector<VkCommandBuffer> m_command_buffers;
        std::vector<VkSemaphore> m_image_available;
        std::vector<VkSemaphore> m_render_finished;
        std::vector<VkFence>     m_in_flight;
        std::vector<VkFence>     m_images_in_flight;
        VkFence m_geo_in_use = VK_NULL_HANDLE;
        uint32_t m_current_frame = 0, m_image_index = 0;

        std::vector<VkDescriptorPool> m_frame_pools;    // one per frame in flight
        VkDescriptorPool m_imgui_pool = VK_NULL_HANDLE;
        bool m_imgui = false;

        VkCommandListR m_cmds;
    };
}