#include <format>

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

#include <vulkan/vulkan.h>
#include <vulkan/vk_enum_string_helper.h>
#include "volk.h"


void Swapchain::initSync(VkDevice device)
{
    const auto fenceCreateInfo = VkFenceCreateInfo{
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT,
    };
    const 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(GfxDevice* gfxDevice, VkFormat format, std::uint32_t width, std::uint32_t height, bool vSync) {
    m_gfxDevice = gfxDevice;
    assert(format == VK_FORMAT_B8G8R8A8_SRGB && "TODO: test other formats");
    vSync = true;

    auto res = vkb::SwapchainBuilder{gfxDevice->getDevice()}
            .set_desired_format(VkSurfaceFormatKHR{
                .format = format,
                .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
            })
            .add_image_usage_flags(VK_IMAGE_USAGE_TRANSFER_DST_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(std::format(
            "failed to create swapchain: error = {}, vk result = {}",
            res.full_error().type.message(),
            string_VkResult(res.full_error().vk_result)));
    }
    m_swapchain = res.value();

    images = m_swapchain.get_images().value();
    imageViews = m_swapchain.get_image_views().value();

    imageRenderSemaphores.resize(images.size());

    VkSemaphoreCreateInfo sci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO
    };

    for (auto& sem : imageRenderSemaphores) {
        VK_CHECK(vkCreateSemaphore(m_gfxDevice->getDevice(), &sci, nullptr, &sem));
    }

    // TODO: if re-creation of swapchain is supported, don't forget to call
    // vkutil::initSwapchainViews here.

    extent = m_swapchain.extent;
}

void Swapchain::recreateSwapchain(const GfxDevice& gfxDevice, VkFormat format, std::uint32_t width, std::uint32_t height, bool vSync) {
    assert(m_swapchain);

    assert(format == VK_FORMAT_B8G8R8A8_SRGB && "TODO: test other formats");
    auto res = vkb::SwapchainBuilder{gfxDevice.getDevice()}
            .set_old_swapchain(m_swapchain)
            .set_desired_format(VkSurfaceFormatKHR{
                .format = format,
                .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
            })
            .add_image_usage_flags(VK_IMAGE_USAGE_TRANSFER_DST_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(std::format(
            "failed to create swapchain: error = {}, vk result = {}",
            res.full_error().type.message(),
            string_VkResult(res.full_error().vk_result)));
    }
    vkb::destroy_swapchain(m_swapchain);

    for (auto imageView: imageViews) {
        vkDestroyImageView(m_gfxDevice->getDevice(), imageView, nullptr);
    }

    m_swapchain = res.value();

    images = m_swapchain.get_images().value();
    imageViews = m_swapchain.get_image_views().value();

    dirty = false;
    extent = m_swapchain.extent;
}

void Swapchain::cleanup() {
    for (auto& frame: frames) {
        vkDestroyFence(m_gfxDevice->getDevice(), frame.renderFence, nullptr);
        vkDestroySemaphore(m_gfxDevice->getDevice(), frame.swapchainSemaphore, nullptr);
    } {
        // destroy swapchain and its views
        for (auto imageView: imageViews) {
            vkDestroyImageView(m_gfxDevice->getDevice(), imageView, nullptr);
        }
        imageViews.clear();

        vkb::destroy_swapchain(m_swapchain);
    }
}

void Swapchain::beginFrame(int index) const {
    auto& frame = frames[index];
    VK_CHECK(vkWaitForFences(m_gfxDevice->getDevice(), 1, &frame.renderFence, true, std::numeric_limits<std::uint64_t>::max()));
}

void Swapchain::resetFences(int index) const {
    auto& frame = frames[index];
    VK_CHECK(vkResetFences(m_gfxDevice->getDevice(), 1, &frame.renderFence));
}

std::pair<VkImage, int> Swapchain::acquireNextImage(int index) {
    std::uint32_t swapchainImageIndex{};
    const auto result = vkAcquireNextImageKHR(
        m_gfxDevice->getDevice(),
        m_swapchain,
        std::numeric_limits<std::uint64_t>::max(),
        frames[index].swapchainSemaphore,
        VK_NULL_HANDLE,
        &swapchainImageIndex);
    if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR) {
        dirty = true;
        return {images[swapchainImageIndex], swapchainImageIndex};
    } else if (result != VK_SUCCESS) {
        throw std::runtime_error("failed to acquire swap chain image!");
    }

    return {images[swapchainImageIndex], swapchainImageIndex};
}

void Swapchain::submitAndPresent(
    VkCommandBuffer cmd,
    VkQueue graphicsQueue,
    uint32_t imageIndex,   // from vkAcquireNextImageKHR
    uint32_t frameIndex)   // 0..FRAMES_IN_FLIGHT-1
{
    auto& frame = frames[frameIndex]; // ✅ per-frame

    VkSemaphore renderFinished = imageRenderSemaphores[imageIndex]; // ✅ per-image

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

    VkSemaphoreSubmitInfo waitInfo =
        vkinit::semaphoreSubmitInfo(
            VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT_KHR,
            frame.swapchainSemaphore); // ✅ acquire semaphore (per-frame)

    VkSemaphoreSubmitInfo signalInfo =
        vkinit::semaphoreSubmitInfo(
            VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT,
            renderFinished); // ✅ signal semaphore (per-image)

    VkSubmitInfo2 submit = vkinit::submitInfo(&cmdInfo, &waitInfo, &signalInfo);
    VK_CHECK(vkQueueSubmit2(graphicsQueue, 1, &submit, frame.renderFence)); // ✅ fence (per-frame)

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

    VkResult res = vkQueuePresentKHR(graphicsQueue, &presentInfo);
    if (res == VK_ERROR_OUT_OF_DATE_KHR || res == VK_SUBOPTIMAL_KHR) dirty = true;
    else if (res != VK_SUCCESS) dirty = true;
}