#include "vulkan_common.h"

namespace Donut::RHI
{
    auto VulkanDevice::create_buffer(BufferType type, size_t size, const void* data) -> Ref<Buffer>
    {
        VkBufferUsageFlags usage = type == BufferType::Index ? VK_BUFFER_USAGE_INDEX_BUFFER_BIT
                                 : type == BufferType::Uniform ? VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT
                                 : VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
        VkBuffer buf = VK_NULL_HANDLE; VkDeviceMemory mem = VK_NULL_HANDLE;
        create_buffer_raw(size, usage, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, buf, mem);
        void* mapped = nullptr; vkMapMemory(m_device, mem, 0, size, 0, &mapped);
        if (data && mapped) std::memcpy(mapped, data, size);
        return create_ref<VkBufferR>(m_device, buf, mem, mapped, size);
    }

    auto VulkanDevice::create_texture(int w, int h, Format format, Filter filter, const void* data) -> Ref<Texture>
    {
        auto tex = create_ref<VkTextureR>(); tex->m_device = m_device;
        VkFormat fmt = vk_format(format);
        VkImageCreateInfo ici{ VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO };
        ici.imageType = VK_IMAGE_TYPE_2D; ici.format = fmt; ici.extent = { (uint32_t)w, (uint32_t)h, 1 };
        ici.mipLevels = 1; ici.arrayLayers = 1; ici.samples = VK_SAMPLE_COUNT_1_BIT;
        ici.tiling = VK_IMAGE_TILING_OPTIMAL; ici.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
        vkCreateImage(m_device, &ici, nullptr, &tex->m_image);
        VkMemoryRequirements req{}; vkGetImageMemoryRequirements(m_device, tex->m_image, &req);
        VkMemoryAllocateInfo ai{ VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
        ai.allocationSize = req.size; ai.memoryTypeIndex = find_memory_type(req.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        vkAllocateMemory(m_device, &ai, nullptr, &tex->m_mem);
        vkBindImageMemory(m_device, tex->m_image, tex->m_mem, 0);

        size_t bpp = format == Format::RGBA16F ? 8 : 4;
        VkDeviceSize sz = (VkDeviceSize)w * h * bpp;
        VkBuffer staging = VK_NULL_HANDLE; VkDeviceMemory staging_mem = VK_NULL_HANDLE;
        create_buffer_raw(sz, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, staging, staging_mem);
        void* mp = nullptr; vkMapMemory(m_device, staging_mem, 0, sz, 0, &mp);
        if (data) std::memcpy(mp, data, sz); else std::memset(mp, 0, sz);
        vkUnmapMemory(m_device, staging_mem);

        VkCommandBufferAllocateInfo cbai{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
        cbai.commandPool = m_command_pool; cbai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cbai.commandBufferCount = 1;
        VkCommandBuffer cmd; vkAllocateCommandBuffers(m_device, &cbai, &cmd);
        VkCommandBufferBeginInfo bi{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO }; bi.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
        vkBeginCommandBuffer(cmd, &bi);
        VkImageMemoryBarrier b{ VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER };
        b.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED; b.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        b.image = tex->m_image; b.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
        b.srcAccessMask = 0; b.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &b);
        VkBufferImageCopy copy{}; copy.imageSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 }; copy.imageExtent = { (uint32_t)w, (uint32_t)h, 1 };
        vkCmdCopyBufferToImage(cmd, staging, tex->m_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copy);
        VkImageMemoryBarrier r = b; r.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL; r.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        r.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; r.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, nullptr, 0, nullptr, 1, &r);
        vkEndCommandBuffer(cmd);
        VkSubmitInfo si{ VK_STRUCTURE_TYPE_SUBMIT_INFO }; si.commandBufferCount = 1; si.pCommandBuffers = &cmd;
        vkQueueSubmit(m_graphics_queue, 1, &si, VK_NULL_HANDLE); vkQueueWaitIdle(m_graphics_queue);
        vkFreeCommandBuffers(m_device, m_command_pool, 1, &cmd);
        vkDestroyBuffer(m_device, staging, nullptr); vkFreeMemory(m_device, staging_mem, nullptr);

        VkImageViewCreateInfo vci{ VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO };
        vci.image = tex->m_image; vci.viewType = VK_IMAGE_VIEW_TYPE_2D; vci.format = fmt;
        vci.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
        vkCreateImageView(m_device, &vci, nullptr, &tex->m_view);
        VkSamplerCreateInfo smci{ VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO };
        smci.magFilter = vk_filter(filter); smci.minFilter = vk_filter(filter);
        smci.addressModeU = smci.addressModeV = smci.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
        vkCreateSampler(m_device, &smci, nullptr, &tex->m_sampler);
        return tex;
    }

    auto VulkanDevice::create_render_target(int w, int h, Format color, Format depth, Filter filter, int /*mips*/) -> Ref<RenderTarget>
    {
        auto rt = create_ref<VkRenderTargetR>();
        rt->m_device = m_device; rt->m_w = w; rt->m_h = h;
        VkFormat cfmt = (color == Format::Swapchain) ? m_swapchain_format : vk_format(color);
        VkImageCreateInfo ici{ VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO };
        ici.imageType = VK_IMAGE_TYPE_2D; ici.format = cfmt; ici.extent = { (uint32_t)w, (uint32_t)h, 1 };
        ici.mipLevels = 1; ici.arrayLayers = 1; ici.samples = VK_SAMPLE_COUNT_1_BIT;
        // TRANSFER_SRC because export reads targets back with a copy. MoltenVK never
        // minded it missing; a desktop driver is within its rights to hand back garbage.
        ici.tiling = VK_IMAGE_TILING_OPTIMAL;
        ici.usage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
        vkCreateImage(m_device, &ici, nullptr, &rt->m_image);
        VkMemoryRequirements req{}; vkGetImageMemoryRequirements(m_device, rt->m_image, &req);
        VkMemoryAllocateInfo ai{ VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
        ai.allocationSize = req.size; ai.memoryTypeIndex = find_memory_type(req.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        vkAllocateMemory(m_device, &ai, nullptr, &rt->m_mem);
        vkBindImageMemory(m_device, rt->m_image, rt->m_mem, 0);
        VkImageViewCreateInfo vci{ VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO };
        vci.image = rt->m_image; vci.viewType = VK_IMAGE_VIEW_TYPE_2D; vci.format = cfmt;
        vci.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
        vkCreateImageView(m_device, &vci, nullptr, &rt->m_view);
        VkSamplerCreateInfo smci{ VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO };
        smci.magFilter = vk_filter(filter); smci.minFilter = vk_filter(filter);
        smci.addressModeU = smci.addressModeV = smci.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE;
        vkCreateSampler(m_device, &smci, nullptr, &rt->m_sampler);

        // optional depth attachment (for off-screen passes that need a depth test).
        rt->m_has_depth = depth != Format::None;
        VkFormat dfmt = VK_FORMAT_UNDEFINED;
        if (rt->m_has_depth)
        {
            dfmt = vk_format(depth);
            VkImageCreateInfo dici{ VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO };
            dici.imageType = VK_IMAGE_TYPE_2D; dici.format = dfmt; dici.extent = { (uint32_t)w, (uint32_t)h, 1 };
            dici.mipLevels = 1; dici.arrayLayers = 1; dici.samples = VK_SAMPLE_COUNT_1_BIT;
            dici.tiling = VK_IMAGE_TILING_OPTIMAL; dici.usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
            vkCreateImage(m_device, &dici, nullptr, &rt->m_depth_image);
            VkMemoryRequirements dreq{}; vkGetImageMemoryRequirements(m_device, rt->m_depth_image, &dreq);
            VkMemoryAllocateInfo dai{ VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
            dai.allocationSize = dreq.size; dai.memoryTypeIndex = find_memory_type(dreq.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            vkAllocateMemory(m_device, &dai, nullptr, &rt->m_depth_mem);
            vkBindImageMemory(m_device, rt->m_depth_image, rt->m_depth_mem, 0);
            VkImageViewCreateInfo dvci{ VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO };
            dvci.image = rt->m_depth_image; dvci.viewType = VK_IMAGE_VIEW_TYPE_2D; dvci.format = dfmt;
            dvci.subresourceRange = { VK_IMAGE_ASPECT_DEPTH_BIT, 0, 1, 0, 1 };
            vkCreateImageView(m_device, &dvci, nullptr, &rt->m_depth_view);
        }

        rt->m_pass = get_render_pass(cfmt, dfmt, false);   // off-screen (sampled) target
        VkImageView atts[2] = { rt->m_view, rt->m_depth_view };
        VkFramebufferCreateInfo fbci{ VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO };
        fbci.renderPass = rt->m_pass; fbci.attachmentCount = rt->m_has_depth ? 2u : 1u; fbci.pAttachments = atts;
        fbci.width = w; fbci.height = h; fbci.layers = 1;
        vkCreateFramebuffer(m_device, &fbci, nullptr, &rt->m_fb);

        rt->m_color.m_device = m_device; rt->m_color.m_view = rt->m_view; rt->m_color.m_sampler = rt->m_sampler; rt->m_color.m_owns = false;
        return rt;
    }

    // builds an environment cubemap from an equirectangular HDRI: render the 6
    // faces with the EquirectToCubemap pipeline, then a full mip chain by
    // linear down-blits (so divergence-based LOD reads a blurred sky). returns
    // a Texture owning the cube image/view/sampler.

    auto VulkanDevice::create_pipeline(const PipelineDesc& desc) -> Ref<Pipeline>
    {
        auto p = create_ref<VkPipelineR>(); p->m_device = m_device; p->m_resources = desc.resources;

        std::vector<VkDescriptorSetLayoutBinding> binds;
        for (const auto& r : desc.resources)
        {
            VkDescriptorSetLayoutBinding b{};
            b.binding = r.binding; b.descriptorCount = 1;
            b.descriptorType = r.kind == ResourceKind::UniformBuffer ? VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            b.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
            binds.push_back(b);
        }
        VkDescriptorSetLayoutCreateInfo dslci{ VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO };
        dslci.bindingCount = (uint32_t)binds.size(); dslci.pBindings = binds.data();
        vkCreateDescriptorSetLayout(m_device, &dslci, nullptr, &p->m_set_layout);
        VkPipelineLayoutCreateInfo plci{ VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO };
        plci.setLayoutCount = 1; plci.pSetLayouts = &p->m_set_layout;
        vkCreatePipelineLayout(m_device, &plci, nullptr, &p->m_layout);

        VkShaderModule vmod = VK_NULL_HANDLE, fmod = VK_NULL_HANDLE;
        if (!create_shader_module("assets/shaders/generated/" + desc.shader + ".vertexMain.spv", vmod) ||
            !create_shader_module("assets/shaders/generated/" + desc.shader + ".fragmentMain.spv", fmod))
        { DONUT_ERROR("Vulkan RHI: shader '{}' modules failed", desc.shader); return p; }
        VkPipelineShaderStageCreateInfo stages[2]{};
        stages[0].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; stages[0].stage = VK_SHADER_STAGE_VERTEX_BIT;   stages[0].module = vmod; stages[0].pName = "main";
        stages[1].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO; stages[1].stage = VK_SHADER_STAGE_FRAGMENT_BIT; stages[1].module = fmod; stages[1].pName = "main";

        VkVertexInputBindingDescription vib{ 0, desc.vertex_layout.stride, VK_VERTEX_INPUT_RATE_VERTEX };
        std::vector<VkVertexInputAttributeDescription> vias;
        for (const auto& a : desc.vertex_layout.attributes)
            vias.push_back({ a.location, 0, vk_attr_format(a.components), a.offset });
        VkPipelineVertexInputStateCreateInfo vin{ VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO };
        vin.vertexBindingDescriptionCount = desc.vertex_layout.stride ? 1 : 0; vin.pVertexBindingDescriptions = &vib;
        vin.vertexAttributeDescriptionCount = (uint32_t)vias.size(); vin.pVertexAttributeDescriptions = vias.data();

        VkPipelineInputAssemblyStateCreateInfo ia{ VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO }; ia.topology = vk_topology(desc.topology);
        VkPipelineViewportStateCreateInfo vps{ VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO }; vps.viewportCount = 1; vps.scissorCount = 1;
        VkDynamicState dyn[2] = { VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR };
        VkPipelineDynamicStateCreateInfo dsci{ VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO }; dsci.dynamicStateCount = 2; dsci.pDynamicStates = dyn;
        VkPipelineRasterizationStateCreateInfo rs{ VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO };
        rs.polygonMode = VK_POLYGON_MODE_FILL; rs.cullMode = vk_cull(desc.cull); rs.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE; rs.lineWidth = 1.0f;
        VkPipelineMultisampleStateCreateInfo ms{ VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO }; ms.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
        VkPipelineDepthStencilStateCreateInfo ds{ VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO };
        ds.depthTestEnable = desc.depth_test ? VK_TRUE : VK_FALSE; ds.depthWriteEnable = desc.depth_write ? VK_TRUE : VK_FALSE; ds.depthCompareOp = vk_compare(desc.depth_op);
        VkPipelineColorBlendAttachmentState cba{};
        cba.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
        if (desc.blend == BlendMode::AlphaBlend)
        {
            cba.blendEnable = VK_TRUE;
            cba.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA; cba.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA; cba.colorBlendOp = VK_BLEND_OP_ADD;
            cba.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE; cba.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO; cba.alphaBlendOp = VK_BLEND_OP_ADD;
        }
        VkPipelineColorBlendStateCreateInfo cb{ VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO }; cb.attachmentCount = 1; cb.pAttachments = &cba;

        VkGraphicsPipelineCreateInfo gpci{ VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO };
        gpci.stageCount = 2; gpci.pStages = stages;
        gpci.pVertexInputState = &vin; gpci.pInputAssemblyState = &ia; gpci.pViewportState = &vps;
        gpci.pRasterizationState = &rs; gpci.pMultisampleState = &ms; gpci.pColorBlendState = &cb; gpci.pDynamicState = &dsci;
        // resolve the target's attachment signature to a (cached) render pass.
        // pipeline<->pass compatibility is by attachment format, so this is the
        // same pass the matching swapchain / render target renders into.
        VkFormat pcolor = (desc.target.color == Format::Swapchain) ? m_swapchain_format : vk_format(desc.target.color);
        VkFormat pdepth = (desc.target.depth == Format::None) ? VK_FORMAT_UNDEFINED : vk_format(desc.target.depth);
        bool present = desc.target.color == Format::Swapchain;
        if (pdepth != VK_FORMAT_UNDEFINED) gpci.pDepthStencilState = &ds;
        gpci.layout = p->m_layout;
        gpci.renderPass = get_render_pass(pcolor, pdepth, present);
        gpci.subpass = 0;
        VkResult pr = vkCreateGraphicsPipelines(m_device, VK_NULL_HANDLE, 1, &gpci, nullptr, &p->m_pipeline);
        vkDestroyShaderModule(m_device, vmod, nullptr); vkDestroyShaderModule(m_device, fmod, nullptr);
        if (pr != VK_SUCCESS) DONUT_ERROR("Vulkan RHI: pipeline '{}' creation failed ({})", desc.shader, (int)pr);
        return p;
    }

}