#include "iwa/swapchain.hpp"

#include <mijin/debug/assert.hpp>
#include <SDL_vulkan.h>
#include <vulkan/vulkan_to_string.hpp>
#include "iwa/device.hpp"
#include "iwa/log.hpp"

namespace iwa
{
namespace
{
std::pair<vk::PresentModeKHR, std::uint32_t> detectPresentMode(const Device& device, const Window& window, const vk::SurfaceCapabilitiesKHR& surfaceCapabilities)
{
    std::vector<vk::PresentModeKHR> presentModes = device.getVkPhysicalDevice().getSurfacePresentModesKHR(window.getVkSurface());

    vk::PresentModeKHR presentMode = vk::PresentModeKHR::eImmediate;
    std::uint32_t imageCount = 2;
    if (std::ranges::find(presentModes, vk::PresentModeKHR::eMailbox) != presentModes.end())
    {
        presentMode = vk::PresentModeKHR::eMailbox;
        imageCount = 3;
    }
    imageCount = std::max(imageCount, surfaceCapabilities.minImageCount);
    if (surfaceCapabilities.maxImageCount > 0)
    {
        imageCount = std::min(imageCount, surfaceCapabilities.maxImageCount);
    }
    return std::make_pair(presentMode, imageCount);
}

vk::SurfaceFormatKHR detectImageFormat(const Device& device, const Window& window, const vk::ImageUsageFlags imageUsage)
{
    std::vector<vk::SurfaceFormatKHR> surfaceFormats = device.getVkPhysicalDevice().getSurfaceFormatsKHR(window.getVkSurface());

    int resultIdx = -1;
    for (std::size_t idx = 0; idx < surfaceFormats.size(); ++idx)
    {
        const vk::SurfaceFormatKHR surfaceFormat = surfaceFormats[idx];
        vk::ImageFormatProperties formatProperties;

        try
        {
            formatProperties = device.getVkPhysicalDevice().getImageFormatProperties(
                    /* format = */ surfaceFormat.format,
                    /* type = */ vk::ImageType::e2D,
                    /* tiling = */ vk::ImageTiling::eOptimal,
                    /* usage = */ imageUsage
            );
        }
        catch(vk::FormatNotSupportedError&)
        {
            continue; // not supported
        }

        if (resultIdx < 0)
        {
            resultIdx = static_cast<int>(idx);
        }
        else if (surfaceFormat.format == vk::Format::eB8G8R8A8Srgb
                 && surfaceFormat.colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear)
        {
            resultIdx = static_cast<int>(idx);
            break;
        }
    }

    if (resultIdx < 0)
    {
        logAndDie("Error creating Vulkan swapchain, no compatible image format found.");
    }

    return surfaceFormats[resultIdx];
}

vk::Extent2D detectImageExtent(const Window& window, const vk::SurfaceCapabilitiesKHR& surfaceCapabilities)
{
    // default to current extent (chosen by the driver)
    vk::Extent2D extent = surfaceCapabilities.currentExtent;

    // alternatively ask SDL about it
    if (extent.width == 0xFFFFFFFF && extent.height == 0xFFFFFFFF)
    {
        int width, height; // NOLINT(readability-isolate-declaration, cppcoreguidelines-init-variables)
        SDL_Vulkan_GetDrawableSize(window.getSDLWindow(), &width, &height);

        extent.width = std::clamp(static_cast<std::uint32_t>(width), surfaceCapabilities.minImageExtent.width, surfaceCapabilities.maxImageExtent.width);
        extent.height = std::clamp(static_cast<std::uint32_t>(height), surfaceCapabilities.minImageExtent.height, surfaceCapabilities.maxImageExtent.height);
    }

    return extent;
}
} // namespace

Swapchain::Swapchain(ObjectPtr<Device> owner, SwapchainCreationArgs args)
        : super_t(std::move(owner)), mWindow(std::move(args.window)), mImageUsage(args.imageUsage)
{
    MIJIN_ASSERT(mWindow, "Invalid SwapchainCreationArgs: window cannot be NULL.");

    mImageAvailableSemaphores.resize(args.parallelFrames);
    recreate();
}
Swapchain::~Swapchain() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroySwapchainKHR)
}

mijin::Task<> Swapchain::c_present(const PresentArgs& args)
{
    while (!mWindow->isVisible()) {
        co_await mijin::c_suspend();
    }

    const vk::PresentInfoKHR presentInfo =
    {
        .waitSemaphoreCount = static_cast<std::uint32_t>(args.waitSemaphores.size()),
        .pWaitSemaphores = args.waitSemaphores.data(),
        .swapchainCount = 1,
        .pSwapchains = &mHandle,
        .pImageIndices = &mCurrentImageIdx
    };

    vk::Result result; // NOLINT(cppcoreguidelines-init-variables)
    try {
        result = args.queue.presentKHR(presentInfo);
    } catch(vk::OutOfDateKHRError&) {
        result = vk::Result::eErrorOutOfDateKHR;
    }
    mCurrentImageIdx = INVALID_IMAGE_INDEX;

    // next image, please
    mCurrentFrameIdx = (mCurrentFrameIdx + 1) % static_cast<int>(getNumParallelFrames());
    if (result == vk::Result::eSuccess) {
        co_await c_acquireImage();
    }
    else {
        recreate();
    }
}

void Swapchain::recreate()
{
    const vk::SurfaceCapabilitiesKHR surfaceCapabilities = getOwner()->getVkPhysicalDevice().getSurfaceCapabilitiesKHR(mWindow->getVkSurface());
    const bool firstCreate = !mHandle;

    // doing this here seems to fix the device loss during recreation
    getOwner()->getVkHandle().waitIdle();

    vk::SwapchainCreateInfoKHR createInfo;

    // surface
    createInfo.surface = mWindow->getVkSurface();

    // old swapchain
    createInfo.oldSwapchain = mHandle;

    // some defaults
    createInfo.imageSharingMode = vk::SharingMode::eExclusive;
    createInfo.clipped = VK_TRUE;
    createInfo.preTransform = surfaceCapabilities.currentTransform;

    // queue families
    createInfo.queueFamilyIndexCount = 1;
    createInfo.pQueueFamilyIndices = &getOwner()->getDeviceInfo().graphicsQueueFamily;

    // present mode and image count
    std::tie(createInfo.presentMode, createInfo.minImageCount) = detectPresentMode(*getOwner(), *mWindow, surfaceCapabilities);

    // image usage
    createInfo.imageUsage = mImageUsage;
    MIJIN_ASSERT((surfaceCapabilities.supportedUsageFlags & createInfo.imageUsage) == createInfo.imageUsage, "Invalid image usage flags when creating surface.");

    // image format and color space
    const vk::SurfaceFormatKHR surfaceFormat = detectImageFormat(*getOwner(), *mWindow, createInfo.imageUsage);
    createInfo.imageFormat = surfaceFormat.format;
    createInfo.imageColorSpace = surfaceFormat.colorSpace;

    // image extent
    createInfo.imageExtent = detectImageExtent(*mWindow, surfaceCapabilities);
    createInfo.imageArrayLayers = 1;

    // composite alpha
    createInfo.compositeAlpha = vk::CompositeAlphaFlagBitsKHR::eOpaque;
    MIJIN_ASSERT((surfaceCapabilities.supportedCompositeAlpha & createInfo.compositeAlpha) == createInfo.compositeAlpha, "Invalid composite alpha when creating surface.");

    mHandle = getOwner()->getVkHandle().createSwapchainKHR(createInfo);
    mFormat = createInfo.imageFormat;
    mExtent = createInfo.imageExtent;

    //for (ObjectPtr<Image>& image : mImages)
    //{
    //    image->unwrap();
    //    unregisterImage(&image);
    //}
    mImages.clear();

    // retrieve swapchain images
    const std::vector<vk::Image> imageHandles = getOwner()->getVkHandle().getSwapchainImagesKHR(mHandle);
    mImages.reserve(imageHandles.size());

    // and create image views
    // destroyImageViews();
    // imageViews.reserve(imageHandles.size());

    for (const vk::Image image : imageHandles)
    {
        mImages.emplace_back(getOwner()->createChild<Image>(ImageWrapArgs{
            .handle = image,
            .type = vk::ImageType::e2D,
            .format = createInfo.imageFormat,
            .usage = createInfo.imageUsage,
            .size = {
                createInfo.imageExtent.width,
                createInfo.imageExtent.height,
                1
            }
        }));


        // const ImageViewCreateInfo ivCreateInfo =
        // {
        //         .nativeImage = image,
        //         .viewType = vk::ImageViewType::e2D,
        //         .format = createInfo.imageFormat,
        //         .components =
        //                 {
        //                         .r = vk::ComponentSwizzle::eIdentity,
        //                         .g = vk::ComponentSwizzle::eIdentity,
        //                         .b = vk::ComponentSwizzle::eIdentity,
        //                         .a = vk::ComponentSwizzle::eIdentity
        //                 },
        //         .subresourceRange =
        //                 {
        //                         .aspectMask = vk::ImageAspectFlagBits::eColor,
        //                         .baseMipLevel = 0,
        //                         .levelCount = 1,
        //                         .baseArrayLayer = 0,
        //                         .layerCount = 1
        //                 }
        // };
        // imageViews.emplace_back().create(ivCreateInfo);

        // images.back().setDebugName(fmt::format("Swapchain Image #{}", images.size() - 1));
        // registerImage(&images.back());
    }

    // create "image available" semaphores
    for (ObjectPtr<Semaphore>& semaphore : mImageAvailableSemaphores)
    {
        semaphore = getOwner()->createChild<Semaphore>();
    }

    if (!firstCreate)
    {
        getOwner()->getVkHandle().destroySwapchainKHR(createInfo.oldSwapchain);
    }

    logMsg("Swapchain (re)created:");
    logMsg("  Extent: {}x{}", mExtent.width, mExtent.height);
    logMsg("  Image Usage: {}", vk::to_string(createInfo.imageUsage));
    logMsg("  Image count: {}", mImages.size());
    logMsg("  Format: {}", vk::to_string(mFormat));
    logMsg("  Color Space: {}", vk::to_string(createInfo.imageColorSpace));

    // update size
    // imageSize.width = createInfo.imageExtent.width;
    // imageSize.height = createInfo.imageExtent.height;
    // imageFormat = createInfo.imageFormat;

    // already request the first frame
    acquireImage();

    recreated.emit();
}

void Swapchain::acquireImage()
{
    MIJIN_ASSERT(mCurrentImageIdx == INVALID_IMAGE_INDEX, "Attempting to acquire a new image before presenting.");

    try
    {
        const vk::ResultValue<std::uint32_t> result = getOwner()->getVkHandle().acquireNextImageKHR(mHandle, std::numeric_limits<std::uint64_t>::max(), *mImageAvailableSemaphores[mCurrentFrameIdx]);
        if (result.result != vk::Result::eSuccess)
        {
            assert(result.result == vk::Result::eSuboptimalKHR);
            recreate();
            return;
        }
        mCurrentImageIdx = result.value;
        mImages[mCurrentImageIdx]->resetUsage(ResetLayout::YES);
    }
    catch(vk::OutOfDateKHRError&) {
        recreate();
    }
}

mijin::Task<> Swapchain::c_acquireImage()
{
    try
    {
        vk::ResultValue<std::uint32_t> result = vk::ResultValue<std::uint32_t>(vk::Result::eTimeout, 0);
        while(true)
        {
            result = getOwner()->getVkHandle().acquireNextImageKHR(mHandle, 0, *mImageAvailableSemaphores[mCurrentFrameIdx]);
            if (result.result != vk::Result::eTimeout && result.result != vk::Result::eNotReady) {
                break;
            }
            co_await mijin::c_suspend();
        }
        if (result.result != vk::Result::eSuccess)
        {
            MIJIN_ASSERT(result.result == vk::Result::eSuboptimalKHR, "Error acquiring swapchain image.");
            recreate();
            co_return;
        }
        mCurrentImageIdx = result.value;
        mImages[mCurrentImageIdx]->resetUsage(ResetLayout::YES);
    }
    catch(vk::OutOfDateKHRError&) {
        recreate();
    }
}
}