#include "UI.hpp"

#include "../../Window.hpp"
#include "Renderer.hpp"
#include "Allocator.hpp"
#include <imgui_impl_glfw.h>
#include <imgui_impl_vulkan.h>
#include "../../../core/utils/logger.hpp"

using namespace render::vk;

static void imgui_check_vk_result(VkResult err) {
    if (err == 0)
        return;
    if (err < 0) {
        FATAL("VK ImGui error " << err);
    } else {
        LOG_W("VK ImGui error " << err);
    }
}

UI::UI(GLFWwindow *window): render::UI() {
    auto ctx = Renderer::Get()->getUICtx();
    device = ctx.device;
    { // Descriptor pool
        std::array<VkDescriptorPoolSize, 11> poolSizes{};
        poolSizes[0].type = VK_DESCRIPTOR_TYPE_SAMPLER;
        poolSizes[0].descriptorCount = 1000;
        poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        poolSizes[1].descriptorCount = 1000;
        poolSizes[2].type = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
        poolSizes[2].descriptorCount = 1000;
        poolSizes[3].type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
        poolSizes[3].descriptorCount = 1000;
        poolSizes[4].type = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
        poolSizes[4].descriptorCount = 1000;
        poolSizes[5].type = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
        poolSizes[5].descriptorCount = 1000;
        poolSizes[6].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[6].descriptorCount = 1000;
        poolSizes[7].type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
        poolSizes[7].descriptorCount = 1000;
        poolSizes[8].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
        poolSizes[8].descriptorCount = 1000;
        poolSizes[9].type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
        poolSizes[9].descriptorCount = 1000;
        poolSizes[10].type = VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT;
        poolSizes[10].descriptorCount = 1000;

        VkDescriptorPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.poolSizeCount = poolSizes.size();
        poolInfo.pPoolSizes = poolSizes.data();
        poolInfo.maxSets = 1000 * poolSizes.size();
        poolInfo.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;

        if (vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {
            FATAL("Failed to create UI descriptor pool!");
        }
    }

    // Setup Platform/Renderer bindings
    ImGui_ImplGlfw_InitForVulkan(window, true);
    ImGui_ImplVulkan_InitInfo init_info = {};
    init_info.Instance = ctx.instance;
    init_info.PhysicalDevice = ctx.physicalDevice;
    init_info.Device = ctx.device;
    init_info.QueueFamily = ctx.queueFamily;
    init_info.Queue = ctx.queue;
    init_info.PipelineCache = VK_NULL_HANDLE;
    init_info.DescriptorPool = descriptorPool;
    init_info.Allocator = nullptr;
    init_info.MinImageCount = ctx.imageCount;
    init_info.ImageCount = ctx.imageCount;
    init_info.CheckVkResultFn = imgui_check_vk_result;
    ImGui_ImplVulkan_Init(&init_info, ctx.renderPass);
    Allocator::GetDefault()->transfer(ImGui_ImplVulkan_CreateFontsTexture);
    aimTexture = Texture::LoadFromFile("content/textures/ui/Aim.dds", {false, false, Texture::Wrap::MIRRORED_REPEAT, 0, false});
    aimHandle = (intptr_t)ImGui_ImplVulkan_AddTexture(aimTexture->getDescriptor());
}
UI::~UI() {
    ImGui_ImplVulkan_Shutdown();
    ImGui_ImplGlfw_Shutdown();

    vkDestroyDescriptorPool(device, descriptorPool, nullptr);
}

void UI::setImageCount(uint32_t count) {
    ImGui_ImplVulkan_SetMinImageCount(count);
}

UI::Actions UI::draw(config::client::options &o, state::state &s, const state::reports &r) {
    ImGui_ImplVulkan_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    return render::UI::draw(o, s, r, aimHandle);
}
void UI::render() {
    Renderer::Get()->recordUI([](VkCommandBuffer buffer){
        ImGui_ImplVulkan_RenderDrawData(ImGui::GetDrawData(), buffer);
    });
}

void UI::Load(Window& w) {
    UI::sInstance = new UI(w.getPtr());
}