#include "SwapChain.hpp"

#include "shared.hpp"

using namespace render::vk;

SwapChain::SwapChain(VkDevice device, const PhysicalDeviceInfo& info): device(device) {
    { // Swapchain
        VkPresentModeKHR presentMode = [&]() {
            // MAYBE: add prefer no triple buffering options
            for (const auto& availablePresentMode: info.swapDetails.presentModes) {
                if (availablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR) {
                    return availablePresentMode;
                }
            }

            return VK_PRESENT_MODE_FIFO_KHR;
        }();

        uint32_t imageCount = info.swapDetails.capabilities.minImageCount + 1;
        if (info.swapDetails.capabilities.maxImageCount > 0 && imageCount > info.swapDetails.capabilities.maxImageCount) {
            imageCount = info.swapDetails.capabilities.maxImageCount;
        }

        VkSwapchainCreateInfoKHR createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
        createInfo.surface = info.surface;
        createInfo.minImageCount = imageCount;
        createInfo.imageFormat = info.surfaceFormat.format;
        createInfo.imageColorSpace = info.surfaceFormat.colorSpace;
        createInfo.imageExtent = info.swapDetails.capabilities.currentExtent;
        createInfo.imageArrayLayers = 1;
        createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT; //VK_IMAGE_USAGE_TRANSFER_DST_BIT

        if (info.queueIndices.graphicsFamily != info.queueIndices.presentFamily) {
            uint32_t queueFamilyIndices[] = {info.queueIndices.graphicsFamily.value(), info.queueIndices.presentFamily.value()};
            createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
            createInfo.queueFamilyIndexCount = 2;
            createInfo.pQueueFamilyIndices = queueFamilyIndices;
        } else {
            createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
            createInfo.queueFamilyIndexCount = 0;
            createInfo.pQueueFamilyIndices = nullptr;
        }
        createInfo.preTransform = info.swapDetails.capabilities.currentTransform;
        createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
        createInfo.presentMode = presentMode;
        createInfo.clipped = VK_TRUE;
        createInfo.oldSwapchain = chain;

        if (vkCreateSwapchainKHR(device, &createInfo, ALLOC, &chain) != VK_SUCCESS) {
            FATAL("Failed to create swap chain!");
        }
    }

    { // Images
        uint32_t imageCount;
        vkGetSwapchainImagesKHR(device, chain, &imageCount, nullptr);
        images.resize(imageCount);
        vkGetSwapchainImagesKHR(device, chain, &imageCount, images.data());

        imageViews.resize(images.size());
        for (size_t i = 0; i < images.size(); i++) {
            VkImageViewCreateInfo createInfo{};
            createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            createInfo.image = images[i];

            createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            createInfo.format = info.surfaceFormat.format;
            createInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;

            createInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            createInfo.subresourceRange.baseMipLevel = 0;
            createInfo.subresourceRange.levelCount = 1;
            createInfo.subresourceRange.baseArrayLayer = 0;
            createInfo.subresourceRange.layerCount = 1;

            if (vkCreateImageView(device, &createInfo, ALLOC, &imageViews[i]) != VK_SUCCESS) {
                FATAL("Failed to create image views!");
            }
        }
    }

    imagesInFlight.resize(imageViews.size(), VK_NULL_HANDLE);
}
SwapChain::~SwapChain() {
    for (auto imageView: imageViews) {
        vkDestroyImageView(device, imageView, ALLOC);
    }
    vkDestroySwapchainKHR(device, chain, ALLOC);
}

std::optional<uint32_t> SwapChain::acquireNextImage(VkSemaphore semaphore, VkFence fence) {
    uint32_t imageIndex;
    auto result = vkAcquireNextImageKHR(device, chain, UINT64_MAX, semaphore, VK_NULL_HANDLE, &imageIndex);
    if (result == VK_ERROR_OUT_OF_DATE_KHR)
        return {};

    if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
        FATAL("Failed to acquire swap chain image!");
    }

    // Check if a previous frame is using this image (i.e. there is its fence to wait on)
    if (imagesInFlight[imageIndex] != VK_NULL_HANDLE) {
        vkWaitForFences(device, 1, &imagesInFlight[imageIndex], VK_TRUE, UINT64_MAX);
    }
    // Mark the image as now being in use by this frame
    imagesInFlight[imageIndex] = fence;
    return imageIndex;
}
bool SwapChain::presentImage(uint32_t idx, VkQueue queue, VkSemaphore signalSemaphore) {
    VkSemaphore signalSemaphores[] = {signalSemaphore};

    VkPresentInfoKHR presentInfo{};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    presentInfo.waitSemaphoreCount = 1;
    presentInfo.pWaitSemaphores = signalSemaphores;

    VkSwapchainKHR swapChains[] = {chain};
    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = swapChains;
    presentInfo.pImageIndices = &idx;
    presentInfo.pResults = nullptr;

    auto result = vkQueuePresentKHR(queue, &presentInfo);

    if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR) {
        return false;
    } else if (result != VK_SUCCESS) {
        FATAL("Failed to present swap chain image!");
    }
    return true;
}