#include "XCEngine/RHI/Vulkan/VulkanSwapChain.h"

#include "XCEngine/RHI/Vulkan/VulkanCommandQueue.h"
#include "XCEngine/RHI/Vulkan/VulkanDevice.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

#include <algorithm>
#include <vector>

#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>

namespace XCEngine {
namespace RHI {

VulkanSwapChain::~VulkanSwapChain() {
    Shutdown();
}

bool VulkanSwapChain::Initialize(VulkanDevice* device, VulkanCommandQueue* presentQueue, HWND__* window, uint32_t width, uint32_t height) {
    if (device == nullptr || presentQueue == nullptr || window == nullptr) {
        return false;
    }

    m_device = device;
    m_presentQueue = presentQueue;
    m_window = window;
    m_width = width;
    m_height = height;

    if (!CreateSurface(window)) {
        return false;
    }

    return CreateSwapChainResources();
}

bool VulkanSwapChain::CreateSurface(HWND__* window) {
    VkWin32SurfaceCreateInfoKHR surfaceInfo = {};
    surfaceInfo.sType = VK_STRUCTURE_TYPE_WIN32_SURFACE_CREATE_INFO_KHR;
    surfaceInfo.hinstance = GetModuleHandle(nullptr);
    surfaceInfo.hwnd = window;
    return vkCreateWin32SurfaceKHR(m_device->GetInstance(), &surfaceInfo, nullptr, &m_surface) == VK_SUCCESS;
}

bool VulkanSwapChain::CreateSwapChainResources() {
    VkSurfaceCapabilitiesKHR capabilities = {};
    if (vkGetPhysicalDeviceSurfaceCapabilitiesKHR(m_device->GetPhysicalDevice(), m_surface, &capabilities) != VK_SUCCESS) {
        return false;
    }

    uint32_t formatCount = 0;
    vkGetPhysicalDeviceSurfaceFormatsKHR(m_device->GetPhysicalDevice(), m_surface, &formatCount, nullptr);
    if (formatCount == 0) {
        return false;
    }

    std::vector<VkSurfaceFormatKHR> formats(formatCount);
    vkGetPhysicalDeviceSurfaceFormatsKHR(m_device->GetPhysicalDevice(), m_surface, &formatCount, formats.data());

    VkSurfaceFormatKHR selectedFormat = formats[0];
    for (const VkSurfaceFormatKHR& candidate : formats) {
        if (candidate.format == VK_FORMAT_R8G8B8A8_UNORM) {
            selectedFormat = candidate;
            break;
        }
    }
    if (selectedFormat.format != VK_FORMAT_R8G8B8A8_UNORM) {
        for (const VkSurfaceFormatKHR& candidate : formats) {
            if (candidate.format == VK_FORMAT_B8G8R8A8_UNORM) {
                selectedFormat = candidate;
                break;
            }
        }
    }
    m_surfaceFormat = selectedFormat.format;

    if ((capabilities.supportedUsageFlags & VK_IMAGE_USAGE_TRANSFER_DST_BIT) == 0 ||
        (capabilities.supportedUsageFlags & VK_IMAGE_USAGE_TRANSFER_SRC_BIT) == 0) {
        return false;
    }

    uint32_t presentModeCount = 0;
    vkGetPhysicalDeviceSurfacePresentModesKHR(m_device->GetPhysicalDevice(), m_surface, &presentModeCount, nullptr);
    std::vector<VkPresentModeKHR> presentModes(presentModeCount);
    if (presentModeCount > 0) {
        vkGetPhysicalDeviceSurfacePresentModesKHR(m_device->GetPhysicalDevice(), m_surface, &presentModeCount, presentModes.data());
    }

    VkPresentModeKHR presentMode = VK_PRESENT_MODE_FIFO_KHR;
    for (VkPresentModeKHR candidate : presentModes) {
        if (candidate == VK_PRESENT_MODE_MAILBOX_KHR) {
            presentMode = candidate;
            break;
        }
    }

    const uint32_t requestedImageCount = (std::max<uint32_t>)(2, capabilities.minImageCount);
    const uint32_t imageCount = capabilities.maxImageCount > 0
        ? std::min<uint32_t>(requestedImageCount, capabilities.maxImageCount)
        : requestedImageCount;

    VkExtent2D extent = capabilities.currentExtent;
    if (extent.width == UINT32_MAX) {
        extent.width = m_width;
        extent.height = m_height;
    }
    m_width = extent.width;
    m_height = extent.height;

    VkSwapchainCreateInfoKHR createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    createInfo.surface = m_surface;
    createInfo.minImageCount = imageCount;
    createInfo.imageFormat = selectedFormat.format;
    createInfo.imageColorSpace = selectedFormat.colorSpace;
    createInfo.imageExtent = extent;
    createInfo.imageArrayLayers = 1;
    createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
    createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    createInfo.preTransform = capabilities.currentTransform;
    createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    createInfo.presentMode = presentMode;
    createInfo.clipped = VK_TRUE;
    createInfo.oldSwapchain = VK_NULL_HANDLE;

    if (vkCreateSwapchainKHR(m_device->GetDevice(), &createInfo, nullptr, &m_swapChain) != VK_SUCCESS) {
        return false;
    }

    uint32_t swapChainImageCount = 0;
    vkGetSwapchainImagesKHR(m_device->GetDevice(), m_swapChain, &swapChainImageCount, nullptr);
    std::vector<VkImage> images(swapChainImageCount);
    vkGetSwapchainImagesKHR(m_device->GetDevice(), m_swapChain, &swapChainImageCount, images.data());

    m_backBuffers.clear();
    m_backBuffers.reserve(swapChainImageCount);
    for (VkImage image : images) {
        auto texture = std::make_unique<VulkanTexture>();
        texture->InitializeSwapChainImage(
            m_device->GetDevice(),
            image,
            m_width,
            m_height,
            ToRHIFormat(selectedFormat.format),
            selectedFormat.format);
        m_backBuffers.push_back(std::move(texture));
    }

    return !m_backBuffers.empty();
}

void VulkanSwapChain::DestroySwapChainResources() {
    m_backBuffers.clear();
    if (m_swapChain != VK_NULL_HANDLE && m_device != nullptr) {
        vkDestroySwapchainKHR(m_device->GetDevice(), m_swapChain, nullptr);
        m_swapChain = VK_NULL_HANDLE;
    }
}

bool VulkanSwapChain::AcquireNextImage() {
    if (m_device == nullptr || m_swapChain == VK_NULL_HANDLE) {
        return false;
    }

    VkFenceCreateInfo fenceInfo = {};
    fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;

    VkFence acquireFence = VK_NULL_HANDLE;
    if (vkCreateFence(m_device->GetDevice(), &fenceInfo, nullptr, &acquireFence) != VK_SUCCESS) {
        return false;
    }

    const VkResult result = vkAcquireNextImageKHR(
        m_device->GetDevice(),
        m_swapChain,
        UINT64_MAX,
        VK_NULL_HANDLE,
        acquireFence,
        &m_currentImageIndex);

    if (result == VK_SUCCESS || result == VK_SUBOPTIMAL_KHR) {
        vkWaitForFences(m_device->GetDevice(), 1, &acquireFence, VK_TRUE, UINT64_MAX);
        vkDestroyFence(m_device->GetDevice(), acquireFence, nullptr);
        return true;
    }

    vkDestroyFence(m_device->GetDevice(), acquireFence, nullptr);
    return false;
}

void VulkanSwapChain::Shutdown() {
    DestroySwapChainResources();

    if (m_surface != VK_NULL_HANDLE && m_device != nullptr) {
        vkDestroySurfaceKHR(m_device->GetInstance(), m_surface, nullptr);
        m_surface = VK_NULL_HANDLE;
    }

    m_presentQueue = nullptr;
    m_window = nullptr;
    m_device = nullptr;
    m_width = 0;
    m_height = 0;
    m_currentImageIndex = 0;
}

RHITexture* VulkanSwapChain::GetCurrentBackBuffer() {
    if (m_currentImageIndex >= m_backBuffers.size()) {
        return nullptr;
    }
    return m_backBuffers[m_currentImageIndex].get();
}

void VulkanSwapChain::Present(uint32_t syncInterval, uint32_t flags) {
    (void)syncInterval;
    (void)flags;

    if (m_presentQueue == nullptr || m_swapChain == VK_NULL_HANDLE) {
        return;
    }

    VkPresentInfoKHR presentInfo = {};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = &m_swapChain;
    presentInfo.pImageIndices = &m_currentImageIndex;

    vkQueuePresentKHR(m_presentQueue->GetQueue(), &presentInfo);
}

void VulkanSwapChain::Resize(uint32_t width, uint32_t height) {
    m_width = width;
    m_height = height;
    DestroySwapChainResources();
    CreateSwapChainResources();
}

} // namespace RHI
} // namespace XCEngine