#include <format>
#include <limits>
#include <stdexcept>

#include <destrum/Graphics/Swapchain.h>
#include <destrum/Graphics/Util.h>
#include <destrum/Graphics/Init.h>

#include <volk.h>
#include <tracy/Tracy.hpp>


void Swapchain::initSync(VkDevice device) {
    constexpr auto fenceCreateInfo = VkFenceCreateInfo{
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT,
    };
    constexpr auto semaphoreCreateInfo = VkSemaphoreCreateInfo{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
    };
    for (std::uint32_t i = 0; i < FRAMES_IN_FLIGHT; ++i) {
        VK_CHECK(vkCreateFence(device, &fenceCreateInfo, nullptr, &frames[i].renderFence));
        VK_CHECK(vkCreateSemaphore(device, &semaphoreCreateInfo, nullptr, &frames[i].swapchainSemaphore));
    }
}

void Swapchain::createSwapchain(
    VkDevice,
    vkb::Device vkbDevice,
    VkSurfaceKHR surf,
    VkFormat format,
    std::uint32_t width,
    std::uint32_t height,
    bool vSync)
{
    ZoneScopedN("Swapchain::createSwapchain");

    surface = surf;

    {
        ZoneScopedN("vkb::SwapchainBuilder::build");

        auto res = vkb::SwapchainBuilder{vkbDevice, surface}
        .set_desired_format(VkSurfaceFormatKHR{
            .format = format,
            .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
        })
        .add_image_usage_flags(
            VK_IMAGE_USAGE_TRANSFER_DST_BIT |
            VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT)
        .set_desired_present_mode(
            vSync ? VK_PRESENT_MODE_FIFO_KHR
                  : VK_PRESENT_MODE_IMMEDIATE_KHR)
        .set_desired_extent(width, height)
        .build();

        if (!res.has_value()) {
            throw std::runtime_error(
                "Failed to create swapchain: " + res.error().message());
        }
        m_swapchain = res.value();
    }

    {
        ZoneScopedN("Get Swapchain Images / Views");

        const auto imageResult = m_swapchain.get_images();
        const auto viewResult = m_swapchain.get_image_views();
        if (!imageResult.has_value() || !viewResult.has_value()) {
            throw std::runtime_error("Failed to retrieve swapchain images or views");
        }
        images = imageResult.value();
        imageViews = viewResult.value();
    }

    imageRenderSemaphores.resize(images.size());

    VkSemaphoreCreateInfo sci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO
    };

    for (auto& sem : imageRenderSemaphores) {
        ZoneScopedN("Create Image Render Semaphore");

        VK_CHECK(vkCreateSemaphore(vkbDevice, &sci, nullptr, &sem));
    }

    extent = m_swapchain.extent;
    dirty = false;
}

void Swapchain::recreateSwapchain(
    VkDevice,
    vkb::Device vkbDevice,
    VkSurfaceKHR surf,
    VkFormat format,
    std::uint32_t width,
    std::uint32_t height,
    bool vSync)
{
    ZoneScopedN("Swapchain::recreateSwapchain");

    if (width == 0 || height == 0) {
        dirty = true;
        return;
    }

    surface = surf;

    {
        ZoneScopedN("vkDeviceWaitIdle");
        vkDeviceWaitIdle(vkbDevice);
    }

    auto oldSwapchain = m_swapchain;

    {
        ZoneScopedN("vkb::SwapchainBuilder::rebuild");

        auto res = vkb::SwapchainBuilder{vkbDevice, surface}
        .set_old_swapchain(oldSwapchain)
        .set_desired_format(VkSurfaceFormatKHR{
            .format = format,
            .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
        })
        .add_image_usage_flags(
            VK_IMAGE_USAGE_TRANSFER_DST_BIT |
            VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT)
        .set_desired_present_mode(
            vSync ? VK_PRESENT_MODE_FIFO_KHR
                  : VK_PRESENT_MODE_IMMEDIATE_KHR)
        .set_desired_extent(width, height)
        .build();

        if (!res.has_value()) {
            throw std::runtime_error(
                "Failed to recreate swapchain: " + res.error().message());
        }

        m_swapchain = res.value();
    }

    {
        ZoneScopedN("Destroy Old Image Render Semaphores");

        for (auto sem : imageRenderSemaphores) {
            vkDestroySemaphore(vkbDevice, sem, nullptr);
        }
        imageRenderSemaphores.clear();
    }

    {
        ZoneScopedN("Destroy Old Image Views");

        for (auto imageView : imageViews) {
            vkDestroyImageView(vkbDevice, imageView, nullptr);
        }
        imageViews.clear();
    }

    {
        ZoneScopedN("Destroy Old Swapchain");
        vkb::destroy_swapchain(oldSwapchain);
    }

    {
        ZoneScopedN("Get New Swapchain Images / Views");

        const auto imageResult = m_swapchain.get_images();
        const auto viewResult = m_swapchain.get_image_views();
        if (!imageResult.has_value() || !viewResult.has_value()) {
            throw std::runtime_error("Failed to retrieve recreated swapchain images or views");
        }
        images = imageResult.value();
        imageViews = viewResult.value();
    }

    VkSemaphoreCreateInfo sci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
    };

    imageRenderSemaphores.resize(images.size());

    for (auto& sem : imageRenderSemaphores) {
        ZoneScopedN("Create New Image Render Semaphore");

        VK_CHECK(vkCreateSemaphore(vkbDevice, &sci, nullptr, &sem));
    }

    extent = m_swapchain.extent;
    dirty = false;
}

void Swapchain::cleanup(VkDevice device) {
    for (auto& frame : frames) {
        if (frame.renderFence != VK_NULL_HANDLE) {
            vkDestroyFence(device, frame.renderFence, nullptr);
            frame.renderFence = VK_NULL_HANDLE;
        }
        if (frame.swapchainSemaphore != VK_NULL_HANDLE) {
            vkDestroySemaphore(device, frame.swapchainSemaphore, nullptr);
            frame.swapchainSemaphore = VK_NULL_HANDLE;
        }
    }

    for (auto& semaphore : imageRenderSemaphores) {
        if (semaphore != VK_NULL_HANDLE) {
            vkDestroySemaphore(device, semaphore, nullptr);
            semaphore = VK_NULL_HANDLE;
        }
    }
    imageRenderSemaphores.clear();

    for (auto imageView : imageViews) {
        vkDestroyImageView(device, imageView, nullptr);
    }
    imageViews.clear();

    vkb::destroy_swapchain(m_swapchain);
    m_swapchain = {};
    images.clear();
    extent = {};
    dirty = false;
}

void Swapchain::beginFrame(VkDevice device, int index) const {
    ZoneScopedN("Swapchain::beginFrame");

    auto& frame = frames[index];

    {
        ZoneScopedN("vkWaitForFences");
        VK_CHECK(vkWaitForFences(device, 1, &frame.renderFence, true, std::numeric_limits<std::uint64_t>::max()));
    }
}

void Swapchain::resetFences(VkDevice device, int index) const {
    ZoneScopedN("Swapchain::resetFences");

    auto& frame = frames[index];

    {
        ZoneScopedN("vkResetFences");
        VK_CHECK(vkResetFences(device, 1, &frame.renderFence));
    }
}

Swapchain::AcquireResult Swapchain::acquireNextImage(VkDevice device, std::uint32_t index) {
    ZoneScopedN("Swapchain::acquireNextImage");

    std::uint32_t swapchainImageIndex{};

    VkResult result = VK_SUCCESS;

    {
        ZoneScopedN("vkAcquireNextImageKHR");

        result = vkAcquireNextImageKHR(
            device,
            m_swapchain,
            std::numeric_limits<std::uint64_t>::max(),
            frames[index].swapchainSemaphore,
            VK_NULL_HANDLE,
            &swapchainImageIndex);
    }

    if (result == VK_ERROR_OUT_OF_DATE_KHR) {
        dirty = true;
        return {.result = result};
    }
    if (result == VK_SUBOPTIMAL_KHR) {
        dirty = true;
        return {
            .result = result,
            .image = images.at(swapchainImageIndex),
            .imageIndex = swapchainImageIndex,
        };
    }
    if (result != VK_SUCCESS) {
        throw std::runtime_error("failed to acquire swap chain image!");
    }

    return {
        .result = result,
        .image = images.at(swapchainImageIndex),
        .imageIndex = swapchainImageIndex,
    };
}

void Swapchain::submitAndPresent(
    VkDevice device,
    VkCommandBuffer cmd,
    VkQueue graphicsQueue,
    VkQueue presentQueue,
    std::uint32_t imageIndex,
    std::uint32_t frameIndex)
{
    ZoneScopedN("Swapchain::submitAndPresent");

    auto& frame = frames[frameIndex];
    VkSemaphore renderFinished = imageRenderSemaphores[imageIndex];

    VkCommandBufferSubmitInfo cmdInfo{
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_SUBMIT_INFO,
        .commandBuffer = cmd,
    };

    VkSemaphoreSubmitInfo waitInfo =
            vkinit::semaphoreSubmitInfo(
                VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
                frame.swapchainSemaphore);

    VkSemaphoreSubmitInfo signalInfo =
            vkinit::semaphoreSubmitInfo(
                VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT,
                renderFinished);

    VkSubmitInfo2 submit = vkinit::submitInfo(&cmdInfo, &waitInfo, &signalInfo);

    VK_CHECK(vkResetFences(device, 1, &frame.renderFence));

    {
        ZoneScopedN("vkQueueSubmit2");
        const VkResult submitResult = vkQueueSubmit2(graphicsQueue, 1, &submit, frame.renderFence);
        if (submitResult != VK_SUCCESS) {
            vkDestroyFence(device, frame.renderFence, nullptr);
            constexpr VkFenceCreateInfo fenceInfo{
                .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
                .flags = VK_FENCE_CREATE_SIGNALED_BIT,
            };
            VK_CHECK(vkCreateFence(device, &fenceInfo, nullptr, &frame.renderFence));
            checkVkResult(submitResult, "vkQueueSubmit2", __FILE__, __LINE__);
        }
    }

    VkPresentInfoKHR presentInfo{
        .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
        .waitSemaphoreCount = 1,
        .pWaitSemaphores = &renderFinished,
        .swapchainCount = 1,
        .pSwapchains = &m_swapchain.swapchain,
        .pImageIndices = &imageIndex,
    };

    VkResult res = VK_SUCCESS;

    {
        ZoneScopedN("vkQueuePresentKHR");
        res = vkQueuePresentKHR(presentQueue, &presentInfo);
    }

    if (res == VK_ERROR_OUT_OF_DATE_KHR || res == VK_SUBOPTIMAL_KHR) {
        dirty = true;
    } else if (res != VK_SUCCESS) {
        throw std::runtime_error("failed to present swap chain image");
    }
}