#include "vulkan_common.h"

namespace Donut::RHI
{
    auto VulkanDevice::create_instance() -> bool
    {
        VkApplicationInfo app{ VK_STRUCTURE_TYPE_APPLICATION_INFO };
        app.pApplicationName = "Donut"; app.apiVersion = VK_API_VERSION_1_2;

        uint32_t glfwExtCount = 0;
        const char** glfwExts = glfwGetRequiredInstanceExtensions(&glfwExtCount);
        if (!glfwExts) { DONUT_ERROR("Vulkan RHI: GLFW reports no surface support"); return false; }
        std::vector<const char*> exts(glfwExts, glfwExts + glfwExtCount);
        exts.push_back(VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);
        exts.push_back(VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME);

        std::vector<const char*> layers;
        uint32_t layer_count = 0; vkEnumerateInstanceLayerProperties(&layer_count, nullptr);
        std::vector<VkLayerProperties> avail(layer_count);
        vkEnumerateInstanceLayerProperties(&layer_count, avail.data());
        for (const auto& l : avail)
            if (std::strcmp(l.layerName, "VK_LAYER_KHRONOS_validation") == 0)
                layers.push_back("VK_LAYER_KHRONOS_validation");

        VkInstanceCreateInfo ici{ VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO };
        ici.flags = VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
        ici.pApplicationInfo = &app;
        ici.enabledExtensionCount = (uint32_t)exts.size(); ici.ppEnabledExtensionNames = exts.data();
        ici.enabledLayerCount = (uint32_t)layers.size(); ici.ppEnabledLayerNames = layers.data();
        VkResult r = vkCreateInstance(&ici, nullptr, &m_instance);
        if (r != VK_SUCCESS && !layers.empty())
        {
            DONUT_WARN("Vulkan RHI: validation layer unavailable, continuing without it");
            ici.enabledLayerCount = 0; ici.ppEnabledLayerNames = nullptr;
            r = vkCreateInstance(&ici, nullptr, &m_instance);
        }
        if (r != VK_SUCCESS) { DONUT_ERROR("Vulkan RHI: vkCreateInstance failed ({})", (int)r); return false; }
        VKD_CHECK(glfwCreateWindowSurface(m_instance, m_window, nullptr, &m_surface));
        DONUT_INFO("Vulkan RHI: instance + surface created (validation {})", layers.empty() ? "off" : "on");
        return true;
    }

    auto VulkanDevice::pick_physical_and_device() -> bool
    {
        uint32_t count = 0; vkEnumeratePhysicalDevices(m_instance, &count, nullptr);
        if (count == 0) { DONUT_ERROR("Vulkan RHI: no physical devices"); return false; }
        std::vector<VkPhysicalDevice> devices(count);
        vkEnumeratePhysicalDevices(m_instance, &count, devices.data());
        m_physical = devices[0];

        uint32_t q = 0; vkGetPhysicalDeviceQueueFamilyProperties(m_physical, &q, nullptr);
        std::vector<VkQueueFamilyProperties> qfams(q);
        vkGetPhysicalDeviceQueueFamilyProperties(m_physical, &q, qfams.data());
        bool fg = false, fp = false;
        for (uint32_t i = 0; i < q; ++i)
        {
            if (!fg && (qfams[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)) { m_graphics_family = i; fg = true; }
            VkBool32 present = VK_FALSE; vkGetPhysicalDeviceSurfaceSupportKHR(m_physical, i, m_surface, &present);
            if (!fp && present) { m_present_family = i; fp = true; }
        }
        if (!fg || !fp) { DONUT_ERROR("Vulkan RHI: no graphics/present queue"); return false; }

        std::vector<const char*> dev_exts = { VK_KHR_SWAPCHAIN_EXTENSION_NAME };
        uint32_t dec = 0; vkEnumerateDeviceExtensionProperties(m_physical, nullptr, &dec, nullptr);
        std::vector<VkExtensionProperties> dep(dec);
        vkEnumerateDeviceExtensionProperties(m_physical, nullptr, &dec, dep.data());
        for (const auto& e : dep)
            if (std::strcmp(e.extensionName, "VK_KHR_portability_subset") == 0)
                dev_exts.push_back("VK_KHR_portability_subset");

        float priority = 1.0f;
        std::vector<VkDeviceQueueCreateInfo> qcis;
        uint32_t families[2] = { m_graphics_family, m_present_family };
        for (uint32_t i = 0; i < (m_graphics_family == m_present_family ? 1u : 2u); ++i)
        {
            VkDeviceQueueCreateInfo qci{ VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO };
            qci.queueFamilyIndex = families[i]; qci.queueCount = 1; qci.pQueuePriorities = &priority;
            qcis.push_back(qci);
        }
        VkDeviceCreateInfo dci{ VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO };
        dci.queueCreateInfoCount = (uint32_t)qcis.size(); dci.pQueueCreateInfos = qcis.data();
        dci.enabledExtensionCount = (uint32_t)dev_exts.size(); dci.ppEnabledExtensionNames = dev_exts.data();
        VKD_CHECK(vkCreateDevice(m_physical, &dci, nullptr, &m_device));
        vkGetDeviceQueue(m_device, m_graphics_family, 0, &m_graphics_queue);
        vkGetDeviceQueue(m_device, m_present_family, 0, &m_present_queue);

        VkPhysicalDeviceProperties props{}; vkGetPhysicalDeviceProperties(m_physical, &props);
        vkGetPhysicalDeviceMemoryProperties(m_physical, &m_mem_props);
        m_gpu_name = props.deviceName;
        DONUT_INFO("Vulkan RHI device: {}", m_gpu_name);
        return true;
    }

    auto VulkanDevice::create_swapchain() -> bool
    {
        VkSurfaceCapabilitiesKHR caps{};
        vkGetPhysicalDeviceSurfaceCapabilitiesKHR(m_physical, m_surface, &caps);
        uint32_t fc = 0; vkGetPhysicalDeviceSurfaceFormatsKHR(m_physical, m_surface, &fc, nullptr);
        std::vector<VkSurfaceFormatKHR> formats(fc);
        vkGetPhysicalDeviceSurfaceFormatsKHR(m_physical, m_surface, &fc, formats.data());
        VkSurfaceFormatKHR chosen = formats[0];
        for (const auto& f : formats)
            if (f.format == VK_FORMAT_B8G8R8A8_UNORM && f.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) chosen = f;
        m_swapchain_format = chosen.format;

        if (caps.currentExtent.width != UINT32_MAX) m_extent = caps.currentExtent;
        else {
            m_extent.width  = std::clamp((uint32_t)m_width,  caps.minImageExtent.width,  caps.maxImageExtent.width);
            m_extent.height = std::clamp((uint32_t)m_height, caps.minImageExtent.height, caps.maxImageExtent.height);
        }
        uint32_t image_count = caps.minImageCount + 1;
        if (caps.maxImageCount > 0 && image_count > caps.maxImageCount) image_count = caps.maxImageCount;

        VkSwapchainCreateInfoKHR sci{ VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR };
        sci.surface = m_surface; sci.minImageCount = image_count;
        sci.imageFormat = chosen.format; sci.imageColorSpace = chosen.colorSpace;
        sci.imageExtent = m_extent; sci.imageArrayLayers = 1;
        sci.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
        sci.preTransform = caps.currentTransform; sci.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
        // FIFO is always available (vsync). With vsync off, prefer IMMEDIATE
        // (uncapped, may tear) then MAILBOX, if the surface supports them.
        VkPresentModeKHR present_mode = VK_PRESENT_MODE_FIFO_KHR;
        if (!m_vsync)
        {
            uint32_t pmc = 0; vkGetPhysicalDeviceSurfacePresentModesKHR(m_physical, m_surface, &pmc, nullptr);
            std::vector<VkPresentModeKHR> modes(pmc);
            vkGetPhysicalDeviceSurfacePresentModesKHR(m_physical, m_surface, &pmc, modes.data());
            auto has = [&](VkPresentModeKHR m){ return std::find(modes.begin(), modes.end(), m) != modes.end(); };
            if (has(VK_PRESENT_MODE_IMMEDIATE_KHR))    present_mode = VK_PRESENT_MODE_IMMEDIATE_KHR;
            else if (has(VK_PRESENT_MODE_MAILBOX_KHR)) present_mode = VK_PRESENT_MODE_MAILBOX_KHR;
        }
        sci.presentMode = present_mode; sci.clipped = VK_TRUE;
        uint32_t fam[2] = { m_graphics_family, m_present_family };
        if (m_graphics_family != m_present_family)
        { sci.imageSharingMode = VK_SHARING_MODE_CONCURRENT; sci.queueFamilyIndexCount = 2; sci.pQueueFamilyIndices = fam; }
        else sci.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
        VKD_CHECK(vkCreateSwapchainKHR(m_device, &sci, nullptr, &m_swapchain));
        uint32_t n = 0; vkGetSwapchainImagesKHR(m_device, m_swapchain, &n, nullptr);
        m_images.resize(n); vkGetSwapchainImagesKHR(m_device, m_swapchain, &n, m_images.data());
        return true;
    }

    auto VulkanDevice::create_image_views() -> bool
    {
        m_image_views.resize(m_images.size());
        for (size_t i = 0; i < m_images.size(); ++i)
        {
            VkImageViewCreateInfo vci{ VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO };
            vci.image = m_images[i]; vci.viewType = VK_IMAGE_VIEW_TYPE_2D; vci.format = m_swapchain_format;
            vci.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };
            VKD_CHECK(vkCreateImageView(m_device, &vci, nullptr, &m_image_views[i]));
        }
        return true;
    }

    // A render pass for one attachment signature, created once and cached.
    // `present` targets (the swapchain) finish PRESENT_SRC and sync on the
    // colour-output stage; `sampled` targets (off-screen) finish
    // SHADER_READ_ONLY and round-trip through the fragment shader so the next
    // pass can sample them. Depth (VK_FORMAT_UNDEFINED = none) is optional.
    auto VulkanDevice::get_render_pass(VkFormat color, VkFormat depth, bool present) -> VkRenderPass
    {
        uint64_t key = (uint64_t)(uint32_t)color
                     | ((uint64_t)(uint32_t)depth << 24)
                     | ((uint64_t)(present ? 1 : 0) << 48);
        auto it = m_pass_cache.find(key);
        if (it != m_pass_cache.end()) return it->second;

        const bool has_depth = depth != VK_FORMAT_UNDEFINED;
        VkAttachmentDescription atts[2]{};
        atts[0].format = color; atts[0].samples = VK_SAMPLE_COUNT_1_BIT;
        atts[0].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; atts[0].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
        atts[0].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; atts[0].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        atts[0].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        atts[0].finalLayout = present ? VK_IMAGE_LAYOUT_PRESENT_SRC_KHR : VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        atts[1].format = depth; atts[1].samples = VK_SAMPLE_COUNT_1_BIT;
        atts[1].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; atts[1].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        atts[1].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; atts[1].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
        atts[1].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; atts[1].finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
        VkAttachmentReference color_ref{ 0, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL };
        VkAttachmentReference depth_ref{ 1, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL };
        VkSubpassDescription subpass{}; subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
        subpass.colorAttachmentCount = 1; subpass.pColorAttachments = &color_ref;
        if (has_depth) subpass.pDepthStencilAttachment = &depth_ref;

        VkSubpassDependency deps[2]{};
        uint32_t dep_count;
        if (present)
        {
            deps[0].srcSubpass = VK_SUBPASS_EXTERNAL; deps[0].dstSubpass = 0;
            deps[0].srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
            deps[0].dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
            deps[0].dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
            dep_count = 1;
        }
        else
        {
            deps[0].srcSubpass = VK_SUBPASS_EXTERNAL; deps[0].dstSubpass = 0;
            deps[0].srcStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT; deps[0].srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
            deps[0].dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT; deps[0].dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
            deps[1].srcSubpass = 0; deps[1].dstSubpass = VK_SUBPASS_EXTERNAL;
            deps[1].srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT; deps[1].srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
            deps[1].dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT; deps[1].dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
            dep_count = 2;
        }

        VkRenderPassCreateInfo rpci{ VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO };
        rpci.attachmentCount = has_depth ? 2 : 1; rpci.pAttachments = atts;
        rpci.subpassCount = 1; rpci.pSubpasses = &subpass;
        rpci.dependencyCount = dep_count; rpci.pDependencies = deps;
        VkRenderPass rp = VK_NULL_HANDLE;
        if (vkCreateRenderPass(m_device, &rpci, nullptr, &rp) != VK_SUCCESS)
        { DONUT_ERROR("Vulkan RHI: render pass creation failed"); return VK_NULL_HANDLE; }
        m_pass_cache[key] = rp;
        return rp;
    }

    auto VulkanDevice::create_depth_and_framebuffers() -> bool
    {
        VkImageCreateInfo dici{ VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO };
        dici.imageType = VK_IMAGE_TYPE_2D; dici.format = VK_FORMAT_D32_SFLOAT;
        dici.extent = { m_extent.width, m_extent.height, 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;
        VKD_CHECK(vkCreateImage(m_device, &dici, nullptr, &m_depth_image));
        VkMemoryRequirements dreq{}; vkGetImageMemoryRequirements(m_device, 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);
        VKD_CHECK(vkAllocateMemory(m_device, &dai, nullptr, &m_depth_mem));
        VKD_CHECK(vkBindImageMemory(m_device, m_depth_image, m_depth_mem, 0));
        VkImageViewCreateInfo dvci{ VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO };
        dvci.image = m_depth_image; dvci.viewType = VK_IMAGE_VIEW_TYPE_2D; dvci.format = VK_FORMAT_D32_SFLOAT;
        dvci.subresourceRange = { VK_IMAGE_ASPECT_DEPTH_BIT, 0, 1, 0, 1 };
        VKD_CHECK(vkCreateImageView(m_device, &dvci, nullptr, &m_depth_view));

        m_framebuffers.resize(m_image_views.size());
        for (size_t i = 0; i < m_image_views.size(); ++i)
        {
            VkImageView att[2] = { m_image_views[i], m_depth_view };
            VkFramebufferCreateInfo fbci{ VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO };
            fbci.renderPass = m_swapchain_rp; fbci.attachmentCount = 2; fbci.pAttachments = att;
            fbci.width = m_extent.width; fbci.height = m_extent.height; fbci.layers = 1;
            VKD_CHECK(vkCreateFramebuffer(m_device, &fbci, nullptr, &m_framebuffers[i]));
        }
        return true;
    }

    auto VulkanDevice::create_command_and_sync() -> bool
    {
        VkCommandPoolCreateInfo pci{ VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO };
        pci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; pci.queueFamilyIndex = m_graphics_family;
        VKD_CHECK(vkCreateCommandPool(m_device, &pci, nullptr, &m_command_pool));
        m_command_buffers.resize(MAX_FRAMES_IN_FLIGHT);
        VkCommandBufferAllocateInfo cbai{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
        cbai.commandPool = m_command_pool; cbai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cbai.commandBufferCount = MAX_FRAMES_IN_FLIGHT;
        VKD_CHECK(vkAllocateCommandBuffers(m_device, &cbai, m_command_buffers.data()));

        m_image_available.resize(MAX_FRAMES_IN_FLIGHT);
        m_in_flight.resize(MAX_FRAMES_IN_FLIGHT);
        m_render_finished.resize(m_images.size());
        m_images_in_flight.assign(m_images.size(), VK_NULL_HANDLE);
        VkSemaphoreCreateInfo sci{ VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO };
        VkFenceCreateInfo fci{ VK_STRUCTURE_TYPE_FENCE_CREATE_INFO }; fci.flags = VK_FENCE_CREATE_SIGNALED_BIT;
        for (int i = 0; i < MAX_FRAMES_IN_FLIGHT; ++i)
        { VKD_CHECK(vkCreateSemaphore(m_device, &sci, nullptr, &m_image_available[i])); VKD_CHECK(vkCreateFence(m_device, &fci, nullptr, &m_in_flight[i])); }
        for (size_t i = 0; i < m_images.size(); ++i)
            VKD_CHECK(vkCreateSemaphore(m_device, &sci, nullptr, &m_render_finished[i]));

        m_frame_pools.resize(MAX_FRAMES_IN_FLIGHT);
        VkDescriptorPoolSize sizes[2] = {
            { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 512 },
            { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 256 },
        };
        for (int i = 0; i < MAX_FRAMES_IN_FLIGHT; ++i)
        {
            VkDescriptorPoolCreateInfo dpci{ VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO };
            dpci.maxSets = 256; dpci.poolSizeCount = 2; dpci.pPoolSizes = sizes;
            VKD_CHECK(vkCreateDescriptorPool(m_device, &dpci, nullptr, &m_frame_pools[i]));
        }
        return true;
    }


    auto VulkanDevice::cleanup_swapchain() -> void
    {
        for (auto fb : m_framebuffers) vkDestroyFramebuffer(m_device, fb, nullptr);
        m_framebuffers.clear();
        if (m_depth_view)  { vkDestroyImageView(m_device, m_depth_view, nullptr); m_depth_view = VK_NULL_HANDLE; }
        if (m_depth_image) { vkDestroyImage(m_device, m_depth_image, nullptr); m_depth_image = VK_NULL_HANDLE; }
        if (m_depth_mem)   { vkFreeMemory(m_device, m_depth_mem, nullptr); m_depth_mem = VK_NULL_HANDLE; }
        for (auto iv : m_image_views) vkDestroyImageView(m_device, iv, nullptr);
        m_image_views.clear();
        if (m_swapchain) { vkDestroySwapchainKHR(m_device, m_swapchain, nullptr); m_swapchain = VK_NULL_HANDLE; }
    }

    auto VulkanDevice::recreate_swapchain() -> bool
    {
        int w = 0, h = 0; glfwGetFramebufferSize(m_window, &w, &h);
        while (w == 0 || h == 0) { glfwGetFramebufferSize(m_window, &w, &h); glfwWaitEvents(); }
        m_width = w; m_height = h;
        vkDeviceWaitIdle(m_device);
        cleanup_swapchain();
        if (!create_swapchain())             return false;
        if (!create_image_views())           return false;
        if (!create_depth_and_framebuffers())return false;
        m_images_in_flight.assign(m_images.size(), VK_NULL_HANDLE);
        return true;
    }

}