// =============================================
//  Aster: gui.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "gui.h"

#include "context.h"
#include "device.h"
#include "helpers.h"
#include "image.h"
#include "window.h"

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

namespace ImGui
{
vk::Format g_AttachmentFormat;
vk::DescriptorPool g_DescriptorPool;

void
VulkanAssert(VkResult result)
{
    AbortIfFailed(result);
}

void
Init(const Context *context, const Device *device, const Window *window, vk::Format attachmentFormat,
     const u32 imageCount, const u32 queueFamily, const vk::Queue queue)
{
    g_AttachmentFormat = attachmentFormat;

    eastl::vector<vk::DescriptorPoolSize> poolSizes = {
        {vk::DescriptorType::eSampler, 1000},
        {vk::DescriptorType::eCombinedImageSampler, 1000},
        {vk::DescriptorType::eSampledImage, 1000},
        {vk::DescriptorType::eStorageImage, 1000},
        {vk::DescriptorType::eUniformTexelBuffer, 1000},
        {vk::DescriptorType::eStorageTexelBuffer, 1000},
        {vk::DescriptorType::eUniformBuffer, 1000},
        {vk::DescriptorType::eStorageBuffer, 1000},
        {vk::DescriptorType::eUniformBufferDynamic, 1000},
        {vk::DescriptorType::eStorageBufferDynamic, 1000},
        {vk::DescriptorType::eInputAttachment, 1000},
    };

    const vk::DescriptorPoolCreateInfo descriptorPoolCreateInfo = {
        .flags = vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet,
        .maxSets = 1000,
        .poolSizeCount = Cast<u32>(poolSizes.size()),
        .pPoolSizes = poolSizes.data(),
    };

    AbortIfFailed(device->m_Device.createDescriptorPool(&descriptorPoolCreateInfo, nullptr, &g_DescriptorPool));

    IMGUI_CHECKVERSION();
    CreateContext();
    ImGuiIO &io = GetIO();
    (void)io;
    io.ConfigFlags |= ImGuiConfigFlags_DockingEnable;
    // io.ConfigFlags |= ImGuiConfigFlags_ViewportsEnable; // Viewports bad

    StyleColorsDark();

    ImGui_ImplGlfw_InitForVulkan(window->m_Window, true);

    vk::PipelineRenderingCreateInfo renderingCreateInfo = {
        .colorAttachmentCount = 1,
        .pColorAttachmentFormats = &g_AttachmentFormat,
    };

    ImGui_ImplVulkan_InitInfo imguiVulkanInitInfo = {
        .Instance = context->m_Instance,
        .PhysicalDevice = device->m_PhysicalDevice,
        .Device = device->m_Device,
        .QueueFamily = queueFamily,
        .Queue = queue,
        .DescriptorPool = g_DescriptorPool,
        .MinImageCount = imageCount,
        .ImageCount = imageCount,
        .PipelineCache = nullptr,
        .UseDynamicRendering = true,
        .PipelineRenderingCreateInfo = renderingCreateInfo,
        .Allocator = nullptr,
        .CheckVkResultFn = VulkanAssert,
    };
    ImGui_ImplVulkan_Init(&imguiVulkanInitInfo);

    ImGui_ImplVulkan_CreateFontsTexture();
}

void
Destroy(const Device *device)
{

    ImGui_ImplVulkan_Shutdown();
    ImGui_ImplGlfw_Shutdown();
    DestroyContext();

    device->m_Device.destroy(Take(g_DescriptorPool), nullptr);
}

void
StartBuild()
{
    ImGui_ImplVulkan_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    NewFrame();

    static ImGuiDockNodeFlags dockspaceFlags = ImGuiDockNodeFlags_None | ImGuiDockNodeFlags_PassthruCentralNode;

    // We are using the ImGuiWindowFlags_NoDocking flag to make the parent window not dockable into,
    // because it would be confusing to have two docking targets within each others.
    ImGuiWindowFlags windowFlags = ImGuiWindowFlags_NoDocking;

    const ImGuiViewport *viewport = GetMainViewport();
    SetNextWindowPos(viewport->WorkPos);
    SetNextWindowSize(viewport->WorkSize);
    SetNextWindowViewport(viewport->ID);
    PushStyleVar(ImGuiStyleVar_WindowRounding, 0.0f);
    PushStyleVar(ImGuiStyleVar_WindowBorderSize, 0.0f);
    windowFlags |=
        ImGuiWindowFlags_NoTitleBar | ImGuiWindowFlags_NoCollapse | ImGuiWindowFlags_NoResize | ImGuiWindowFlags_NoMove;
    windowFlags |= ImGuiWindowFlags_NoBringToFrontOnFocus | ImGuiWindowFlags_NoNavFocus;
    windowFlags |= ImGuiWindowFlags_NoBackground;

    // Important: note that we proceed even if Begin() returns false (aka window is collapsed).
    // This is because we want to keep our DockSpace() active. If a DockSpace() is inactive,
    // all active windows docked into it will lose their parent and become undocked.
    // We cannot preserve the docking relationship between an active window and an inactive docking, otherwise
    // any change of dockspace/settings would lead to windows being stuck in limbo and never being visible.
    PushStyleVar(ImGuiStyleVar_WindowPadding, ImVec2(0.0f, 0.0f));
    Begin("DockSpace Demo", nullptr, windowFlags);
    PopStyleVar();

    PopStyleVar(2);

    // DockSpace
    if (GetIO().ConfigFlags & ImGuiConfigFlags_DockingEnable)
    {
        const ImGuiID dockspaceId = GetID("MyDockSpace");
        DockSpace(dockspaceId, ImVec2(0.0f, 0.0f), dockspaceFlags);
    }
}

void
EndBuild()
{
    End();
    Render();

    EndFrame();
    if (GetIO().ConfigFlags & ImGuiConfigFlags_ViewportsEnable)
    {
        GLFWwindow *backupCurrentContext = glfwGetCurrentContext();
        UpdatePlatformWindows();
        RenderPlatformWindowsDefault();
        glfwMakeContextCurrent(backupCurrentContext);
    }
}

void
Draw(const vk::CommandBuffer commandBuffer, const vk::Extent2D extent, const vk::ImageView view)
{
    // OPTICK_EVENT();

    constexpr vk::DebugUtilsLabelEXT label = {
        .pLabelName = "UI pass",
        .color = std::array{0.0f, 0.0f, 1.0f, 1.0f},
    };
    commandBuffer.beginDebugUtilsLabelEXT(&label);
    vk::RenderingAttachmentInfo attachmentInfo = {
        .imageView = view,
        .imageLayout = vk::ImageLayout::eColorAttachmentOptimal,
        .resolveMode = vk::ResolveModeFlagBits::eNone,
        .loadOp = vk::AttachmentLoadOp::eLoad,
        .storeOp = vk::AttachmentStoreOp::eStore,
        .clearValue = vk::ClearColorValue{0.0f, 0.0f, 0.0f, 1.0f},
    };

    const vk::RenderingInfo renderingInfo = {
        .renderArea = {.extent = extent},
        .layerCount = 1,
        .colorAttachmentCount = 1,
        .pColorAttachments = &attachmentInfo,
        .pDepthAttachment = nullptr,
    };

    commandBuffer.beginRendering(&renderingInfo);

    ImGui_ImplVulkan_RenderDrawData(GetDrawData(), commandBuffer);

    commandBuffer.endRendering();

    commandBuffer.endDebugUtilsLabelEXT();
}

void
PushDisable()
{
    PushItemFlag(ImGuiItemFlags_Disabled, true);
    PushStyleVar(ImGuiStyleVar_Alpha, GetStyle().Alpha * 0.5f);
}

void
PopDisable()
{
    PopStyleVar();
    PopItemFlag();
}
} // namespace ImGui