#include "backend.hpp"
#include "lib/vkcheck.hpp"
#include "vulkan-context.hpp"

#include <cstdlib>
#include <print>
#include <vector>

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

namespace {
constexpr uint32_t min_image_count = 2;

inline void check_vk_result(VkResult err) {
    if (err == 0) {
        return;
    }
    std::println(stderr, "[vulkan] Error: VkResult = {}", static_cast<int>(err));
    if (err < 0) {
        std::abort();
    }
}

void init_vulkan_window(VulkanContext& context, ImGui_ImplVulkanH_Window& window, VkSurfaceKHR surface, int width, int height) {
    VkBool32 res = VK_FALSE;
    vkGetPhysicalDeviceSurfaceSupportKHR(context.physical_device, context.queue_family, surface, &res);
    if (res != VK_TRUE) {
        std::println(stderr, "Physical device does not support WSI");
        std::exit(1);
    }

    const std::vector<VkFormat> request_surface_image_format = {
        VK_FORMAT_R8G8B8_UNORM, VK_FORMAT_R8G8B8A8_UNORM, VK_FORMAT_B8G8R8A8_UNORM, VK_FORMAT_B8G8R8_UNORM};

    const VkColorSpaceKHR request_surface_color_space = VK_COLORSPACE_SRGB_NONLINEAR_KHR;
    window.Surface = surface;
    window.SurfaceFormat = ImGui_ImplVulkanH_SelectSurfaceFormat(context.physical_device, window.Surface,
        request_surface_image_format.data(), request_surface_image_format.size(), request_surface_color_space);

    const std::vector<VkPresentModeKHR> present_modes = {VK_PRESENT_MODE_FIFO_KHR};
    window.PresentMode = ImGui_ImplVulkanH_SelectPresentMode(
        context.physical_device, window.Surface, present_modes.data(), present_modes.size());

    static_assert(min_image_count >= 2);
    ImGui_ImplVulkanH_CreateOrResizeWindow(context.instance, context.physical_device, context.logical_device, &window,
        context.queue_family, context.allocator, width, height, min_image_count, 0);
}

void render_frame(
    VulkanContext& context, ImGui_ImplVulkanH_Window& window, ImDrawData* draw_data, bool& should_rebuild_swapchain) {
    auto* image_acquired = window.FrameSemaphores[window.SemaphoreIndex].ImageAcquiredSemaphore;
    auto* render_complete = window.FrameSemaphores[window.SemaphoreIndex].RenderCompleteSemaphore;

    VkResult err = vkAcquireNextImageKHR(
        context.logical_device, window.Swapchain, UINT64_MAX, image_acquired, VK_NULL_HANDLE, &window.FrameIndex);

    if (err == VK_ERROR_OUT_OF_DATE_KHR) {
        should_rebuild_swapchain = true;
        return;
    }

    if (err == VK_SUBOPTIMAL_KHR) {
        should_rebuild_swapchain = true;
    } else {
        check_vk_result(err);
    }

    auto* frame = &window.Frames[window.FrameIndex];

    check_vk_result(vkWaitForFences(context.logical_device, 1, &frame->Fence, VK_TRUE, UINT64_MAX));
    check_vk_result(vkResetFences(context.logical_device, 1, &frame->Fence));

    check_vk_result(vkResetCommandPool(context.logical_device, frame->CommandPool, 0));
    VkCommandBufferBeginInfo command_buffer_begin_info = {};
    command_buffer_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    command_buffer_begin_info.flags |= VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    check_vk_result(vkBeginCommandBuffer(frame->CommandBuffer, &command_buffer_begin_info));

    VkRenderPassBeginInfo render_pass_begin_info = {};
    render_pass_begin_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    render_pass_begin_info.renderPass = window.RenderPass;
    render_pass_begin_info.framebuffer = frame->Framebuffer;
    render_pass_begin_info.renderArea.extent.width = window.Width;
    render_pass_begin_info.renderArea.extent.height = window.Height;
    render_pass_begin_info.clearValueCount = 1;
    render_pass_begin_info.pClearValues = &window.ClearValue;

    vkCmdBeginRenderPass(frame->CommandBuffer, &render_pass_begin_info, VK_SUBPASS_CONTENTS_INLINE);
    ImGui_ImplVulkan_RenderDrawData(draw_data, frame->CommandBuffer);
    vkCmdEndRenderPass(frame->CommandBuffer);

    VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    VkSubmitInfo submit_info = {};
    submit_info.waitSemaphoreCount = 1;
    submit_info.pWaitSemaphores = &image_acquired;
    submit_info.pWaitDstStageMask = &wait_stage;
    submit_info.commandBufferCount = 1;
    submit_info.pCommandBuffers = &frame->CommandBuffer;
    submit_info.signalSemaphoreCount = 1;
    submit_info.pSignalSemaphores = &render_complete;

    check_vk_result(vkEndCommandBuffer(frame->CommandBuffer));
    check_vk_result(vkQueueSubmit(context.queue, 1, &submit_info, frame->Fence));
}

void present_frame(VulkanContext& context, ImGui_ImplVulkanH_Window& window, bool& should_rebuild_swapchain) {
    if (should_rebuild_swapchain) {
        return;
    }

    auto* render_complete = window.FrameSemaphores[window.SemaphoreIndex].RenderCompleteSemaphore;

    VkPresentInfoKHR present_info = {};
    present_info.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    present_info.waitSemaphoreCount = 1;
    present_info.pWaitSemaphores = &render_complete;
    present_info.swapchainCount = 1;
    present_info.pSwapchains = &window.Swapchain;
    present_info.pImageIndices = &window.FrameIndex;

    auto err = vkQueuePresentKHR(context.queue, &present_info);
    if (err == VK_ERROR_OUT_OF_DATE_KHR) {
        should_rebuild_swapchain = true;
        return;
    }

    if (err == VK_SUBOPTIMAL_KHR) {
        should_rebuild_swapchain = true;
    } else {
        check_vk_result(err);
    }

    window.SemaphoreIndex = (window.SemaphoreIndex + 1) % window.SemaphoreCount;
}
} // namespace

int32_t main([[maybe_unused]] int32_t argc, [[maybe_unused]] char** argv) {
    if (!backend_init()) {
        return 1;
    }

    float main_scale = backend_get_main_scale();

    if (!backend_supports_vulkan()) {
        std::println("Vulkan not supported");
        return 1;
    }

    BackendWindow* window = nullptr;
    if (!backend_create_window((int) (800 * main_scale), (int) (800 * main_scale), "vkcheck", &window)) {
        std::println("Failed to create backend window: Vulkan not supported?");
        return 1;
    }

    std::vector<const char*> extensions = backend_get_instance_extensions();

    auto context = VulkanContext(extensions);

    VkSurfaceKHR surface = nullptr;
    backend_create_window_surface(context, window, &surface);

    int32_t width = 0;
    int32_t height = 0;
    backend_get_framebuffer_size(window, &width, &height);

    ImGui_ImplVulkanH_Window main_window;
    init_vulkan_window(context, main_window, surface, width, height);
    backend_show_window(window);

    IMGUI_CHECKVERSION();
    ImGui::CreateContext();

    auto& io = ImGui::GetIO();
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard; // Enable Keyboard Controls
    io.IniFilename = nullptr;                             // Disable .ini file

    ImGui::StyleColorsDark();
    ImGuiStyle& style = ImGui::GetStyle();
    style.ScaleAllSizes(main_scale);
    style.FontScaleDpi = main_scale;

    backend_init_for_vulkan(window, true);
    ImGui_ImplVulkan_InitInfo init_info = {};
    init_info.Instance = context.instance;
    init_info.PhysicalDevice = context.physical_device;
    init_info.Device = context.logical_device;
    init_info.QueueFamily = context.queue_family;
    init_info.Queue = context.queue;
    init_info.DescriptorPool = context.descriptor_pool;
    init_info.MinImageCount = min_image_count;
    init_info.ImageCount = main_window.ImageCount;
    init_info.Allocator = context.allocator;
    init_info.PipelineInfoMain.RenderPass = main_window.RenderPass;
    init_info.PipelineInfoMain.Subpass = 0;
    init_info.PipelineInfoMain.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
    init_info.CheckVkResultFn = check_vk_result;
    ImGui_ImplVulkan_Init(&init_info);

    auto gui = vkcheck::Gui(context.instance, main_scale, true);

    bool swapchain_rebuild = false;
    while (!backend_poll_events(window)) {
        int32_t fb_width = 0;
        int32_t fb_height = 0;
        backend_get_framebuffer_size(window, &fb_width, &fb_height);

        if (fb_width > 0 && fb_height > 0 &&
            (swapchain_rebuild || main_window.Width != fb_width || main_window.Height != fb_height)) {
            ImGui_ImplVulkan_SetMinImageCount(min_image_count);
            ImGui_ImplVulkanH_CreateOrResizeWindow(context.instance, context.physical_device, context.logical_device,
                &main_window, context.queue_family, context.allocator, fb_width, fb_height, min_image_count, 0);
            main_window.FrameIndex = 0;
            swapchain_rebuild = false;
        }

        if (backend_window_is_minimized(window)) {
            backend_sleep(10);
            continue;
        }

        ImGui_ImplVulkan_NewFrame();
        backend_new_frame();
        ImGui::NewFrame();

        gui.render();

        ImGui::Render();
        auto* draw_data = ImGui::GetDrawData();
        const bool is_minimized = draw_data->DisplaySize.x <= 0 || draw_data->DisplaySize.y <= 0;
        if (!is_minimized) {
            render_frame(context, main_window, draw_data, swapchain_rebuild);
            present_frame(context, main_window, swapchain_rebuild);
        }
    }

    check_vk_result(vkDeviceWaitIdle(context.logical_device));
    ImGui_ImplVulkan_Shutdown();
    backend_shutdown();
    ImGui::DestroyContext();

    ImGui_ImplVulkanH_DestroyWindow(context.instance, context.logical_device, &main_window, context.allocator);
    vkDestroySurfaceKHR(context.instance, main_window.Surface, context.allocator);

    vkDestroyDescriptorPool(context.logical_device, context.descriptor_pool, context.allocator);
    vkDestroyDevice(context.logical_device, context.allocator);
    vkDestroyInstance(context.instance, context.allocator);

    backend_destroy_window(window);
    backend_terminate();
}
