#include "vulkan_common.h"

#include <imgui.h>
#include <imgui_impl_glfw.h>
#include <imgui_impl_vulkan.h>

#include <fstream>
#include <cstdlib>

namespace Donut::RHI
{
    auto VulkanDevice::find_memory_type(uint32_t filter, VkMemoryPropertyFlags flags) const -> uint32_t
    {
        for (uint32_t i = 0; i < m_mem_props.memoryTypeCount; ++i)
            if ((filter & (1u << i)) && (m_mem_props.memoryTypes[i].propertyFlags & flags) == flags)
                return i;
        return UINT32_MAX;
    }

    auto VulkanDevice::create_buffer_raw(VkDeviceSize size, VkBufferUsageFlags usage, VkMemoryPropertyFlags props,
                                         VkBuffer& buf, VkDeviceMemory& mem) const -> bool
    {
        VkBufferCreateInfo bci{ VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
        bci.size = size; bci.usage = usage; bci.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
        if (vkCreateBuffer(m_device, &bci, nullptr, &buf) != VK_SUCCESS) return false;
        VkMemoryRequirements req{}; vkGetBufferMemoryRequirements(m_device, buf, &req);
        VkMemoryAllocateInfo ai{ VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO };
        ai.allocationSize = req.size; ai.memoryTypeIndex = find_memory_type(req.memoryTypeBits, props);
        if (vkAllocateMemory(m_device, &ai, nullptr, &mem) != VK_SUCCESS) return false;
        vkBindBufferMemory(m_device, buf, mem, 0);
        return true;
    }

    auto VulkanDevice::load_spirv(const std::string& path) const -> std::vector<uint32_t>
    {
        std::ifstream file(path, std::ios::ate | std::ios::binary);
        if (!file.is_open()) return {};
        size_t size = (size_t)file.tellg();
        std::vector<uint32_t> data(size / 4);
        file.seekg(0); file.read(reinterpret_cast<char*>(data.data()), size);
        return data;
    }

    auto VulkanDevice::create_shader_module(const std::string& path, VkShaderModule& out) const -> bool
    {
        auto spv = load_spirv(path);
        if (spv.empty()) { DONUT_ERROR("Vulkan RHI: failed to load SPIR-V {}", path); return false; }
        VkShaderModuleCreateInfo ci{ VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO };
        ci.codeSize = spv.size() * 4; ci.pCode = spv.data();
        return vkCreateShaderModule(m_device, &ci, nullptr, &out) == VK_SUCCESS;
    }


    auto VulkanDevice::init(const NativeWindow& window) -> bool
    {
        m_window = (GLFWwindow*)window.glfw_handle; m_width = window.width; m_height = window.height;
        if (!create_instance())              return false;
        if (!pick_physical_and_device())     return false;
        if (!create_swapchain())             return false;
        if (!create_image_views())           return false;
        // The swapchain is a "present" target: its colour format + a D32 depth.
        m_swapchain_rp = get_render_pass(m_swapchain_format, VK_FORMAT_D32_SFLOAT, true);
        if (!m_swapchain_rp)                 return false;
        if (!create_depth_and_framebuffers())return false;
        if (!create_command_and_sync())      return false;
        DONUT_INFO("Vulkan RHI device ready: {} swapchain images, {}x{}", (int)m_images.size(), m_extent.width, m_extent.height);
        return true;
    }


    auto VulkanDevice::begin_frame(const glm::vec4&) -> CommandList*
    {
        if (m_device == VK_NULL_HANDLE) return nullptr;
        vkWaitForFences(m_device, 1, &m_in_flight[m_current_frame], VK_TRUE, UINT64_MAX);
        VkResult r = vkAcquireNextImageKHR(m_device, m_swapchain, UINT64_MAX, m_image_available[m_current_frame], VK_NULL_HANDLE, &m_image_index);
        if (r == VK_ERROR_OUT_OF_DATE_KHR) { recreate_swapchain(); return nullptr; }
        if (r != VK_SUCCESS && r != VK_SUBOPTIMAL_KHR) { DONUT_ERROR("Vulkan RHI: acquire failed ({})", (int)r); return nullptr; }
        if (m_images_in_flight[m_image_index] != VK_NULL_HANDLE)
            vkWaitForFences(m_device, 1, &m_images_in_flight[m_image_index], VK_TRUE, UINT64_MAX);
        m_images_in_flight[m_image_index] = m_in_flight[m_current_frame];
        if (m_geo_in_use != VK_NULL_HANDLE)
            vkWaitForFences(m_device, 1, &m_geo_in_use, VK_TRUE, UINT64_MAX);

        vkResetDescriptorPool(m_device, m_frame_pools[m_current_frame], 0);
        VkCommandBuffer cmd = m_command_buffers[m_current_frame];
        vkResetCommandBuffer(cmd, 0);
        VkCommandBufferBeginInfo bi{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
        vkBeginCommandBuffer(cmd, &bi);

        m_cmds.m_device = m_device; m_cmds.m_cmd = cmd;
        m_cmds.m_swapchain_rp = m_swapchain_rp;
        m_cmds.m_swapchain_fb = m_framebuffers[m_image_index]; m_cmds.m_extent = m_extent;
        m_cmds.m_frame_pool = m_frame_pools[m_current_frame]; m_cmds.m_pipe = nullptr;
        return &m_cmds;
    }

    auto VulkanDevice::end_frame() -> void
    {
        VkCommandBuffer cmd = m_command_buffers[m_current_frame];
        vkEndCommandBuffer(cmd);
        VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
        VkSubmitInfo submit{ VK_STRUCTURE_TYPE_SUBMIT_INFO };
        submit.waitSemaphoreCount = 1; submit.pWaitSemaphores = &m_image_available[m_current_frame]; submit.pWaitDstStageMask = &wait_stage;
        submit.commandBufferCount = 1; submit.pCommandBuffers = &cmd;
        submit.signalSemaphoreCount = 1; submit.pSignalSemaphores = &m_render_finished[m_image_index];
        vkResetFences(m_device, 1, &m_in_flight[m_current_frame]);
        if (vkQueueSubmit(m_graphics_queue, 1, &submit, m_in_flight[m_current_frame]) != VK_SUCCESS)
        { DONUT_ERROR("Vulkan RHI: queue submit failed"); return; }
        m_geo_in_use = m_in_flight[m_current_frame];

        VkPresentInfoKHR present{ VK_STRUCTURE_TYPE_PRESENT_INFO_KHR };
        present.waitSemaphoreCount = 1; present.pWaitSemaphores = &m_render_finished[m_image_index];
        present.swapchainCount = 1; present.pSwapchains = &m_swapchain; present.pImageIndices = &m_image_index;
        VkResult r = vkQueuePresentKHR(m_present_queue, &present);
        if (r == VK_ERROR_OUT_OF_DATE_KHR || r == VK_SUBOPTIMAL_KHR || m_framebuffer_resized)
        { m_framebuffer_resized = false; recreate_swapchain(); }
        m_current_frame = (m_current_frame + 1) % MAX_FRAMES_IN_FLIGHT;
    }

    auto VulkanDevice::run_offscreen(const std::function<void(CommandList&)>& record) -> void
    {
        vkDeviceWaitIdle(m_device);
        vkResetDescriptorPool(m_device, m_frame_pools[0], 0);   // per-draw sets for this pass

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

        m_cmds.m_device = m_device; m_cmds.m_cmd = cmd;
        m_cmds.m_swapchain_rp = m_swapchain_rp; m_cmds.m_swapchain_fb = VK_NULL_HANDLE; m_cmds.m_extent = m_extent;
        m_cmds.m_frame_pool = m_frame_pools[0]; m_cmds.m_pipe = nullptr;
        record(m_cmds);   // records its own off-screen render pass(es)

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

    auto VulkanDevice::read_render_target(RenderTarget* target, std::vector<uint8_t>& out) -> void
    {
        auto* rt = static_cast<VkRenderTargetR*>(target);
        uint32_t w = (uint32_t)rt->m_w, h = (uint32_t)rt->m_h;
        VkDeviceSize sz = (VkDeviceSize)w * h * 4;
        out.resize(sz);

        VkBuffer buf; VkDeviceMemory mem;
        create_buffer_raw(sz, VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                          VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, buf, 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);
        // The target ended its render pass in SHADER_READ_ONLY; move it to TRANSFER_SRC to copy.
        VkImageMemoryBarrier b{ VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER };
        b.oldLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; b.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        b.image = rt->m_image; b.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
        b.srcAccessMask = VK_ACCESS_SHADER_READ_BIT; b.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_FRAGMENT_SHADER_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 = { w, h, 1 };
        vkCmdCopyImageToBuffer(cmd, rt->m_image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, buf, 1, &copy);
        VkImageMemoryBarrier b2 = b; b2.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL; b2.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        b2.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT; b2.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, &b2);
        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);

        void* mp = nullptr; vkMapMemory(m_device, mem, 0, sz, 0, &mp);
        std::memcpy(out.data(), mp, sz);   // RGBA8_UNORM, top-down (top-left origin)
        vkUnmapMemory(m_device, mem);
        vkFreeCommandBuffers(m_device, m_command_pool, 1, &cmd);
        vkDestroyBuffer(m_device, buf, nullptr); vkFreeMemory(m_device, mem, nullptr);
    }

    auto VulkanDevice::read_render_target_float(RenderTarget* target, std::vector<float>& out) -> void
    {
        auto* rt = static_cast<VkRenderTargetR*>(target);
        uint32_t w = (uint32_t)rt->m_w, h = (uint32_t)rt->m_h;
        out.resize((size_t)w * h * 4);
        VkDeviceSize sz = (VkDeviceSize)out.size() * sizeof(float);

        VkBuffer buf; VkDeviceMemory mem;
        create_buffer_raw(sz, VK_BUFFER_USAGE_TRANSFER_DST_BIT,
                          VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, buf, 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_SHADER_READ_ONLY_OPTIMAL; b.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        b.image = rt->m_image; b.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
        b.srcAccessMask = VK_ACCESS_SHADER_READ_BIT; b.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_FRAGMENT_SHADER_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 = { w, h, 1 };
        vkCmdCopyImageToBuffer(cmd, rt->m_image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, buf, 1, &copy);
        VkImageMemoryBarrier b2 = b; b2.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL; b2.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        b2.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT; b2.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, &b2);
        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);

        void* mp = nullptr; vkMapMemory(m_device, mem, 0, sz, 0, &mp);
        std::memcpy(out.data(), mp, sz);   // R32G32B32A32_SFLOAT, top-down
        vkUnmapMemory(m_device, mem);
        vkFreeCommandBuffers(m_device, m_command_pool, 1, &cmd);
        vkDestroyBuffer(m_device, buf, nullptr); vkFreeMemory(m_device, mem, nullptr);
    }

    auto VulkanDevice::init_imgui() -> void
    {
        VkDescriptorPoolSize pool_size{ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1000 };
        VkDescriptorPoolCreateInfo dpci{ VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO };
        dpci.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT; dpci.maxSets = 1000;
        dpci.poolSizeCount = 1; dpci.pPoolSizes = &pool_size;
        vkCreateDescriptorPool(m_device, &dpci, nullptr, &m_imgui_pool);

        IMGUI_CHECKVERSION(); ImGui::CreateContext();
        ImGui::GetIO().ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;
        ImGui::GetIO().ConfigFlags |= ImGuiConfigFlags_DockingEnable;
        ImGui::StyleColorsDark();
        ImGui_ImplGlfw_InitForVulkan(m_window, true);
        ImGui_ImplVulkan_InitInfo info{};
        info.ApiVersion = VK_API_VERSION_1_2; info.Instance = m_instance; info.PhysicalDevice = m_physical;
        info.Device = m_device; info.QueueFamily = m_graphics_family; info.Queue = m_graphics_queue;
        info.DescriptorPool = m_imgui_pool; info.RenderPass = m_swapchain_rp;
        info.MinImageCount = 2; info.ImageCount = (uint32_t)m_images.size(); info.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
        if (!ImGui_ImplVulkan_Init(&info)) { DONUT_ERROR("Vulkan RHI: ImGui_ImplVulkan_Init failed"); return; }
        m_imgui = true;
        DONUT_INFO("Vulkan RHI: ImGui backend initialized");
    }

    auto VulkanDevice::imgui_new_frame() -> void
    {
        if (!m_imgui) return;
        ImGui_ImplVulkan_NewFrame(); ImGui_ImplGlfw_NewFrame(); ImGui::NewFrame();
    }

    auto VulkanDevice::imgui_render(CommandList& cmds) -> void
    {
        if (!m_imgui) return;
        ImGui::Render();
        ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), static_cast<VkCommandListR&>(cmds).m_cmd);
    }

    auto VulkanDevice::shutdown() -> void
    {
        if (m_device == VK_NULL_HANDLE)
        {
            if (m_instance && m_surface) { vkDestroySurfaceKHR(m_instance, m_surface, nullptr); m_surface = VK_NULL_HANDLE; }
            if (m_instance) { vkDestroyInstance(m_instance, nullptr); m_instance = VK_NULL_HANDLE; }
            return;
        }
        vkDeviceWaitIdle(m_device);
        if (m_imgui) { ImGui_ImplVulkan_Shutdown(); ImGui_ImplGlfw_Shutdown(); ImGui::DestroyContext(); m_imgui = false; }
        if (m_imgui_pool) vkDestroyDescriptorPool(m_device, m_imgui_pool, nullptr);
        for (auto p : m_frame_pools) vkDestroyDescriptorPool(m_device, p, nullptr);
        m_frame_pools.clear();
        for (auto s : m_render_finished) vkDestroySemaphore(m_device, s, nullptr);
        for (auto s : m_image_available) vkDestroySemaphore(m_device, s, nullptr);
        for (auto f : m_in_flight)       vkDestroyFence(m_device, f, nullptr);
        m_render_finished.clear(); m_image_available.clear(); m_in_flight.clear();
        if (m_command_pool) vkDestroyCommandPool(m_device, m_command_pool, nullptr);
        for (auto& [key, rp] : m_pass_cache) vkDestroyRenderPass(m_device, rp, nullptr);
        m_pass_cache.clear(); m_swapchain_rp = VK_NULL_HANDLE;
        cleanup_swapchain();
        vkDestroyDevice(m_device, nullptr); m_device = VK_NULL_HANDLE;
        if (m_surface)  vkDestroySurfaceKHR(m_instance, m_surface, nullptr);
        if (m_instance) vkDestroyInstance(m_instance, nullptr);
        m_surface = VK_NULL_HANDLE; m_instance = VK_NULL_HANDLE;
    }

    auto create_vulkan_device() -> Scope<Device> { return create_scope<VulkanDevice>(); }

    auto vulkan_prepare_glfw() -> void
    {
#ifdef __APPLE__
        if (!getenv("VK_ICD_FILENAMES"))
            setenv("VK_ICD_FILENAMES", "/opt/homebrew/etc/vulkan/icd.d/MoltenVK_icd.json", 0);
        if (!getenv("VK_LAYER_PATH"))
            setenv("VK_LAYER_PATH", "/opt/homebrew/share/vulkan/explicit_layer.d", 0);
        if (!getenv("DYLD_LIBRARY_PATH"))
            setenv("DYLD_LIBRARY_PATH", "/opt/homebrew/lib", 0);
#endif
        glfwInitVulkanLoader(vkGetInstanceProcAddr);
    }
}