#include "renderer.h"

#include "imgui_impl_glfw.h"
#include "render_target.h"

#include "application/application.h"
#include "texture.h"



#ifdef APP_USE_VULKAN_DEBUG_REPORT
static VKAPI_ATTR VkBool32 VKAPI_CALL on_debug_report(VkDebugReportFlagsEXT flags, VkDebugReportObjectTypeEXT objectType, uint64_t object, size_t location, int32_t messageCode, const char* pLayerPrefix, const char* pMessage, void* pUserData)
{
    (void)flags; (void)object; (void)location; (void)messageCode; (void)pUserData; (void)pLayerPrefix; // Unused arguments
    fprintf(stderr, "[vulkan] Debug report from ObjectType: %i\nMessage: %s\n\n", objectType, pMessage);
    return VK_FALSE;
}
#endif // APP_USE_VULKAN_DEBUG_REPORT

static bool is_extension_available(const std::vector<vk::ExtensionProperties>& properties, const char* extension) {
    return std::ranges::any_of(properties, [extension](const vk::ExtensionProperties& p) {
        return strcmp(p.extensionName, extension) == 0;
    });
}

vk::CommandPool renderer::get_command_pool() const {
    return main_window_data.Frames[main_window_data.FrameIndex].CommandPool;
}

vk::CommandBuffer renderer::create_command_buffer(vk::CommandBufferLevel level, bool begin) const {
    vk::CommandBufferAllocateInfo alloc_info;
    alloc_info.setCommandPool(get_command_pool());
    alloc_info.setLevel(level);
    alloc_info.setCommandBufferCount(1);

    vk::CommandBuffer command_buffer;
    auto err = device.allocateCommandBuffers(&alloc_info, &command_buffer);
    check_vk_result(err);

    // If requested, also start the new command buffer
    if (begin) {
        vk::CommandBufferBeginInfo begin_info;
        begin_info.setFlags(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);
        command_buffer.begin(begin_info);
    }

    return command_buffer;
}

void renderer::end_command_buffer(vk::CommandBuffer command_buffer, bool use_fence) const {
    command_buffer.end();

    if (use_fence) {
        vk::FenceCreateInfo fence_create_info = {};
        vk::Fence fence = device.createFence(fence_create_info);

        vk::SubmitInfo submit_info;
        submit_info.setCommandBuffers(command_buffer);
        queue.submit(submit_info, fence);

        const auto err = device.waitForFences(1, &fence, VK_TRUE, 100000000000);
        check_vk_result(err);

        device.destroyFence(fence);
    }
    else {
        vk::SubmitInfo submit_info;
        submit_info.setCommandBuffers(command_buffer);
        queue.submit(submit_info, nullptr);
    }
}

void renderer::init_vulkan(GLFWwindow* window_handle) {
    std::vector<const char*> extensions;
    uint32_t extensions_count = 0;
    const char** glfw_extensions = glfwGetRequiredInstanceExtensions(&extensions_count);
    for (uint32_t i = 0; i < extensions_count; i++)
        extensions.push_back(glfw_extensions[i]);
    setup_vulkan(extensions);

    // Create Window Surface
    VkSurfaceKHR surface;
    VkResult err = glfwCreateWindowSurface(instance, window_handle, reinterpret_cast<VkAllocationCallbacks*>(allocator), &surface);
    check_vk_result(err);

    // Create Framebuffers
    int w, h;
    glfwGetFramebufferSize(window_handle, &w, &h);
    setup_vulkan_window(surface, w, h);

    ImGui_ImplGlfw_InitForVulkan(window_handle, true);

    ImGui_ImplVulkan_InitInfo init_info = {};
    init_info.Instance = instance;
    init_info.PhysicalDevice = physical_device;
    init_info.Device = device;
    init_info.QueueFamily = queue_family;
    init_info.Queue = queue;
    init_info.PipelineCache = VK_NULL_HANDLE;
    init_info.DescriptorPool = descriptor_pool;
    init_info.RenderPass = main_window_data.RenderPass;
    init_info.Subpass = 0;
    init_info.MinImageCount = min_image_count;
    init_info.ImageCount = main_window_data.ImageCount;
    init_info.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
    init_info.Allocator = reinterpret_cast<VkAllocationCallbacks*>(allocator);
    init_info.CheckVkResultFn = check_vk_result;
    ImGui_ImplVulkan_Init(&init_info);
}

vk::PhysicalDevice renderer::setup_vulkan_select_physical_device() const {
    const std::vector<vk::PhysicalDevice> gpus = instance.enumeratePhysicalDevices();
    IM_ASSERT(!gpus.empty());

    // If a number >1 of GPUs got reported, find discrete GPU if present, or use first one available. This covers
    // most common cases (multi-gpu/integrated+dedicated graphics). Handling more complicated setups (multiple
    // dedicated GPUs) is out of scope of this sample.
    for (auto& device: gpus) {
        if (const auto properties = device.getProperties(); properties.deviceType == vk::PhysicalDeviceType::eDiscreteGpu)
            return device;
    }

    // Use first GPU (Integrated) is a Discrete one is not available.
    if (!gpus.empty())
        return gpus[0];
    return VK_NULL_HANDLE;
}

void renderer::setup_vulkan(std::vector<const char*> instance_extensions) {
    // Create Vulkan Instance
    {
        vk::InstanceCreateInfo create_info;

        // VkInstanceCreateInfo create_info = {};
        // create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;

        // Enumerate available extensions
        auto properties = vk::enumerateInstanceExtensionProperties();

        // Enable required extensions
        if (is_extension_available(properties, VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME))
            instance_extensions.push_back(VK_KHR_GET_PHYSICAL_DEVICE_PROPERTIES_2_EXTENSION_NAME);
#ifdef VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME
        if (is_extension_available(properties, VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME)) {
            instance_extensions.push_back(VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);
            create_info.flags |= vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR;
        }
#endif
        // Enabling validation layers
#ifdef APP_USE_VULKAN_DEBUG_REPORT
        const char* layers[] = { "VK_LAYER_KHRONOS_validation" };
        create_info.enabledLayerCount = 1;
        create_info.ppEnabledLayerNames = layers;
        instance_extensions.push_back("VK_EXT_debug_report");
#endif

        // Create Vulkan Instance
        create_info.setPEnabledExtensionNames(instance_extensions);
        instance = vk::createInstance(create_info, allocator);

        // Setup the debug report callback
#ifdef APP_USE_VULKAN_DEBUG_REPORT
        auto vkCreateDebugReportCallbackEXT = (PFN_vkCreateDebugReportCallbackEXT)vkGetInstanceProcAddr(instance, "vkCreateDebugReportCallbackEXT");
        IM_ASSERT(vkCreateDebugReportCallbackEXT != nullptr);
        VkDebugReportCallbackCreateInfoEXT debug_report_ci = {};
        debug_report_ci.sType = VK_STRUCTURE_TYPE_DEBUG_REPORT_CALLBACK_CREATE_INFO_EXT;
        debug_report_ci.flags = VK_DEBUG_REPORT_ERROR_BIT_EXT | VK_DEBUG_REPORT_WARNING_BIT_EXT | VK_DEBUG_REPORT_PERFORMANCE_WARNING_BIT_EXT;
        debug_report_ci.pfnCallback = on_debug_report;
        debug_report_ci.pUserData = nullptr;
        auto err = vkCreateDebugReportCallbackEXT(instance, &debug_report_ci, (VkAllocationCallbacks*)allocator, &debug_report);
        check_vk_result(err);
#endif
    }

    // Select Physical Device (GPU)
    physical_device = setup_vulkan_select_physical_device();

    // Select graphics queue family
    {
        const auto queues = physical_device.getQueueFamilyProperties();
        for (uint32_t i = 0; i < queues.size(); i++) {
            if (queues[i].queueFlags & vk::QueueFlagBits::eGraphics) {
                queue_family = i;
                break;
            }
        }
        IM_ASSERT(queue_family != static_cast<uint32_t>(-1));
    }

    // Create Logical Device (with 1 queue)
    {
        std::vector<const char*> device_extensions;
        device_extensions.emplace_back("VK_KHR_swapchain");

        // Enumerate physical device extension
        auto properties = physical_device.enumerateDeviceExtensionProperties();

#ifdef VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME
        if (is_extension_available(properties, VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME))
            device_extensions.push_back(VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME);
#endif

        std::vector<float> queue_priority = {1.0f};
        vk::DeviceQueueCreateInfo queue_info;
        queue_info.setQueueFamilyIndex(queue_family);
        queue_info.setQueuePriorities(queue_priority);

        vk::DeviceCreateInfo create_info;
        create_info.setQueueCreateInfos(queue_info);
        create_info.setPEnabledExtensionNames(device_extensions);
        device = physical_device.createDevice(create_info, allocator);
        queue = device.getQueue(queue_family, 0);
    }

    // Create Descriptor Pool
    // The example only requires a single combined image sampler descriptor for the font image and only uses one descriptor set (for that)
    // If you wish to load e.g. additional textures you may need to alter pools sizes.
    {
        std::vector<vk::DescriptorPoolSize> pool_sizes;
        pool_sizes.emplace_back(vk::DescriptorType::eCombinedImageSampler, 1);

        auto properties = physical_device.getProperties();
        vk::DescriptorPoolCreateInfo descriptor_pool_create_info;
        descriptor_pool_create_info.setMaxSets(properties.limits.maxDescriptorSetStorageImages);
        descriptor_pool_create_info.setPoolSizes(pool_sizes);

        descriptor_pool = device.createDescriptorPool(descriptor_pool_create_info);
    }
}

// All the ImGui_ImplVulkanH_XXX structures/functions are optional helpers used by the demo.
// Your real engine/app may not use them.
void renderer::setup_vulkan_window(VkSurfaceKHR surface, int width,
                                          int height) {
    main_window_data.Surface = surface;

    // Check for WSI support
    vk::Bool32 res;
    const auto err = physical_device.getSurfaceSupportKHR(queue_family, main_window_data.Surface, &res);
    check_vk_result(err);
    if (res != VK_TRUE) {
        fprintf(stderr, "Error no WSI support on physical device 0\n");
        exit(-1);
    }

    // Select Surface Format
    constexpr VkFormat requestSurfaceImageFormat[] = {
        VK_FORMAT_B8G8R8A8_UNORM, VK_FORMAT_R8G8B8A8_UNORM, VK_FORMAT_B8G8R8_UNORM, VK_FORMAT_R8G8B8_UNORM
    };
    constexpr VkColorSpaceKHR requestSurfaceColorSpace = VK_COLORSPACE_SRGB_NONLINEAR_KHR;
    main_window_data.SurfaceFormat = ImGui_ImplVulkanH_SelectSurfaceFormat(physical_device, main_window_data.Surface, requestSurfaceImageFormat,
                                                              (size_t) IM_ARRAYSIZE(requestSurfaceImageFormat),
                                                              requestSurfaceColorSpace);

    // Select Present Mode
#ifdef APP_USE_UNLIMITED_FRAME_RATE
    VkPresentModeKHR present_modes[] = { VK_PRESENT_MODE_MAILBOX_KHR, VK_PRESENT_MODE_IMMEDIATE_KHR, VK_PRESENT_MODE_FIFO_KHR };
#else
    VkPresentModeKHR present_modes[] = {VK_PRESENT_MODE_FIFO_KHR};
#endif
    main_window_data.PresentMode = ImGui_ImplVulkanH_SelectPresentMode(physical_device, main_window_data.Surface, &present_modes[0],
                                                          IM_ARRAYSIZE(present_modes));
    //printf("[vulkan] Selected PresentMode = %d\n", wd->PresentMode);

    // Create SwapChain, RenderPass, Framebuffer, etc.
    IM_ASSERT(min_image_count >= 2);

    ImGui_ImplVulkanH_CreateOrResizeWindow(instance, physical_device, device, &main_window_data, queue_family,
                                           reinterpret_cast<VkAllocationCallbacks*>(allocator), width,
                                           height, min_image_count);
}

void renderer::cleanup_vulkan() const {
    device.destroyDescriptorPool(descriptor_pool);
#ifdef APP_USE_VULKAN_DEBUG_REPORT
    // Remove the debug report callback
    auto vkDestroyDebugReportCallbackEXT = (PFN_vkDestroyDebugReportCallbackEXT)vkGetInstanceProcAddr(instance, "vkDestroyDebugReportCallbackEXT");
    vkDestroyDebugReportCallbackEXT(instance, debug_report, (VkAllocationCallbacks*)allocator);
#endif // APP_USE_VULKAN_DEBUG_REPORT

    device.destroy();
    instance.destroy();
}

void renderer::cleanup_vulkan_window() {
    ImGui_ImplVulkanH_DestroyWindow(instance, device, &main_window_data,
                                    reinterpret_cast<VkAllocationCallbacks*>(allocator));
}

void renderer::frame_render(ImDrawData* draw_data) {
    vk::Semaphore image_acquired_semaphore = main_window_data.FrameSemaphores[main_window_data.SemaphoreIndex].ImageAcquiredSemaphore;
    vk::Semaphore render_complete_semaphore = main_window_data.FrameSemaphores[main_window_data.SemaphoreIndex].RenderCompleteSemaphore;

    vk::Result err = device.acquireNextImageKHR(main_window_data.Swapchain, UINT64_MAX, image_acquired_semaphore, VK_NULL_HANDLE,
                                                &main_window_data.FrameIndex);
    if (err == vk::Result::eErrorOutOfDateKHR || err == vk::Result::eSuboptimalKHR) {
        swap_chain_rebuild_ = true;
        return;
    }

    check_vk_result(err);

    ImGui_ImplVulkanH_Frame* fd = &main_window_data.Frames[main_window_data.FrameIndex];
    const vk::CommandBuffer cmd_buf = fd->CommandBuffer;
    const vk::Fence fence = fd->Fence; {
        err = device.waitForFences(1, &fence, VK_TRUE, UINT64_MAX);

        // wait indefinitely instead of periodically checking
        check_vk_result(err);

        err = device.resetFences(1, &fence);
        check_vk_result(err);
    } {
        const vk::CommandPool command_pool = fd->CommandPool;
        device.resetCommandPool(command_pool);

        vk::CommandBufferBeginInfo info = {};
        info.setFlags(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);

        cmd_buf.begin(info);
    } {
        const vk::Framebuffer framebuffer = fd->Framebuffer;
        const vk::RenderPass render_pass = main_window_data.RenderPass;

        const auto clear_color = main_window_data.ClearValue.color.float32;
        const auto clear_depth = main_window_data.ClearValue.depthStencil.depth;
        const auto clear_stencil = main_window_data.ClearValue.depthStencil.stencil;

        vk::ClearValue clear_value;
        clear_value.color = vk::ClearColorValue(std::array<float, 4>{
            clear_color[0], clear_color[1], clear_color[2], clear_color[3]
        });
        clear_value.depthStencil = vk::ClearDepthStencilValue(clear_depth, clear_stencil);

        std::vector<vk::ClearValue> clear_values;
        clear_values.emplace_back(clear_value);

        vk::RenderPassBeginInfo info;
        info.setRenderPass(render_pass);
        info.setFramebuffer(framebuffer);
        info.renderArea.extent.width = main_window_data.Width;
        info.renderArea.extent.height = main_window_data.Height;
        info.setClearValues(clear_values);

        cmd_buf.beginRenderPass(info, vk::SubpassContents::eInline);
    }

    // Record dear imgui primitives into command buffer
    ImGui_ImplVulkan_RenderDrawData(draw_data, fd->CommandBuffer);

    // Submit command buffer
    vkCmdEndRenderPass(fd->CommandBuffer); {
        vk::PipelineStageFlags wait_stage = vk::PipelineStageFlagBits::eColorAttachmentOutput;

        vk::SubmitInfo info;
        info.setWaitSemaphores(image_acquired_semaphore);
        info.setWaitDstStageMask(wait_stage);
        info.setCommandBuffers(cmd_buf);
        info.setSignalSemaphores(render_complete_semaphore);

        cmd_buf.end();
        err = queue.submit(1, &info, fence);
        check_vk_result(err);
    }
}

void renderer::frame_present() {
    if (swap_chain_rebuild_)
        return;
    vk::Semaphore render_complete_semaphore = main_window_data.FrameSemaphores[main_window_data.SemaphoreIndex].RenderCompleteSemaphore;
    vk::SwapchainKHR swapchain = main_window_data.Swapchain;
    uint32_t frame_index = main_window_data.FrameIndex;

    vk::PresentInfoKHR info;
    info.setWaitSemaphores(render_complete_semaphore);
    info.setSwapchains(swapchain);
    info.setImageIndices(frame_index);

    try {
        (void)queue.presentKHR(info);
    } catch (const vk::OutOfDateKHRError& e) {
        swap_chain_rebuild_ = true;
        return;
    }
    main_window_data.SemaphoreIndex = (main_window_data.SemaphoreIndex + 1) % main_window_data.SemaphoreCount; // Now we can use the next set of semaphores
}

void renderer::pre_init() {
    glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
}

bool renderer::init(GLFWwindow* window_handle) {
    if (has_initialized_)
        return true;

    if (!glfwVulkanSupported()) {
        throw std::runtime_error("Vulkan not supported");
    }
    init_vulkan(window_handle);

    has_initialized_ = true;
    return true;
}

void renderer::shutdown() {
    ImGui_ImplGlfw_Shutdown();
    ImGui_ImplVulkan_Shutdown();

    cleanup_vulkan_window();
    cleanup_vulkan();
}

void renderer::new_frame(GLFWwindow* window_handle) {
    // Resize swap chain?
    if (swap_chain_rebuild_)
    {
        int width, height;
        glfwGetFramebufferSize(window_handle, &width, &height);
        if (width > 0 && height > 0)
        {
            ImGui_ImplVulkan_SetMinImageCount(min_image_count);
            ImGui_ImplVulkanH_CreateOrResizeWindow(instance, physical_device, device, &main_window_data, queue_family, reinterpret_cast<VkAllocationCallbacks*>(allocator), width, height, min_image_count);
            main_window_data.FrameIndex = 0;
            swap_chain_rebuild_ = false;
        }
    }

    // Start the Dear ImGui frame
    ImGui_ImplVulkan_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();
}

void renderer::end_frame(GLFWwindow* window_handle) {
    ImGuiIO& io = ImGui::GetIO();

    // Rendering
    ImGui::Render();
    ImDrawData* main_draw_data = ImGui::GetDrawData();
    const bool main_is_minimized = (main_draw_data->DisplaySize.x <= 0.0f || main_draw_data->DisplaySize.y <= 0.0f);
    main_window_data.ClearValue.color.float32[0] = clear_color.x * clear_color.w;
    main_window_data.ClearValue.color.float32[1] = clear_color.y * clear_color.w;
    main_window_data.ClearValue.color.float32[2] = clear_color.z * clear_color.w;
    main_window_data.ClearValue.color.float32[3] = clear_color.w;
    if (!main_is_minimized)
        frame_render(main_draw_data);

    // Update and Render additional Platform Windows
    if (io.ConfigFlags & ImGuiConfigFlags_ViewportsEnable)
    {
        ImGui::UpdatePlatformWindows();
        ImGui::RenderPlatformWindowsDefault();
    }

    // Present Main Platform Window
    if (!main_is_minimized)
        frame_present();
}

std::shared_ptr<texture> renderer::create_texture(const uint8_t* data, int width, int height, vk::Format format) {
    auto out = std::make_shared<texture>();
    out->init(width, height, format);
    if (data)
        out->upload(data, width * height * 4);
    return out;
}

std::shared_ptr<render_target> renderer::create_render_target(int width, int height, vk::Format format) {
    auto out = std::make_shared<render_target>();
    out->init(width, height, format);
    return out;
}