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

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

#include <cstdio>
#include <vector>

namespace XCEngine {
namespace RHI {

bool VulkanScreenshot::Capture(RHIDevice* device, RHISwapChain* swapChain, const char* filename) {
    auto* vkDevice = static_cast<VulkanDevice*>(device);
    auto* texture = static_cast<VulkanTexture*>(swapChain->GetCurrentBackBuffer());
    if (vkDevice == nullptr || texture == nullptr) {
        return false;
    }

    const uint32_t width = texture->GetWidth();
    const uint32_t height = texture->GetHeight();
    const VkDevice logicalDevice = vkDevice->GetDevice();

    VkBuffer stagingBuffer = VK_NULL_HANDLE;
    VkDeviceMemory stagingMemory = VK_NULL_HANDLE;
    const VkDeviceSize bufferSize = static_cast<VkDeviceSize>(width) * static_cast<VkDeviceSize>(height) * 4;

    VkBufferCreateInfo bufferInfo = {};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = bufferSize;
    bufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    if (vkCreateBuffer(logicalDevice, &bufferInfo, nullptr, &stagingBuffer) != VK_SUCCESS) {
        return false;
    }

    VkMemoryRequirements memoryRequirements = {};
    vkGetBufferMemoryRequirements(logicalDevice, stagingBuffer, &memoryRequirements);

    VkMemoryAllocateInfo allocateInfo = {};
    allocateInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocateInfo.allocationSize = memoryRequirements.size;
    allocateInfo.memoryTypeIndex = vkDevice->FindMemoryType(
        memoryRequirements.memoryTypeBits,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    if (allocateInfo.memoryTypeIndex == UINT32_MAX ||
        vkAllocateMemory(logicalDevice, &allocateInfo, nullptr, &stagingMemory) != VK_SUCCESS) {
        vkDestroyBuffer(logicalDevice, stagingBuffer, nullptr);
        return false;
    }

    vkBindBufferMemory(logicalDevice, stagingBuffer, stagingMemory, 0);

    VkCommandPool commandPool = VK_NULL_HANDLE;
    VkCommandBuffer commandBuffer = VK_NULL_HANDLE;

    VkCommandPoolCreateInfo poolInfo = {};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.queueFamilyIndex = vkDevice->GetGraphicsQueueFamilyIndex();
    if (vkCreateCommandPool(logicalDevice, &poolInfo, nullptr, &commandPool) != VK_SUCCESS) {
        vkFreeMemory(logicalDevice, stagingMemory, nullptr);
        vkDestroyBuffer(logicalDevice, stagingBuffer, nullptr);
        return false;
    }

    VkCommandBufferAllocateInfo commandBufferInfo = {};
    commandBufferInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    commandBufferInfo.commandPool = commandPool;
    commandBufferInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    commandBufferInfo.commandBufferCount = 1;
    vkAllocateCommandBuffers(logicalDevice, &commandBufferInfo, &commandBuffer);

    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    vkBeginCommandBuffer(commandBuffer, &beginInfo);

    VkImageMemoryBarrier toCopySource = {};
    toCopySource.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    toCopySource.oldLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
    toCopySource.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
    toCopySource.srcAccessMask = 0;
    toCopySource.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
    toCopySource.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    toCopySource.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    toCopySource.image = texture->GetImage();
    toCopySource.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    toCopySource.subresourceRange.baseMipLevel = 0;
    toCopySource.subresourceRange.levelCount = 1;
    toCopySource.subresourceRange.baseArrayLayer = 0;
    toCopySource.subresourceRange.layerCount = 1;

    vkCmdPipelineBarrier(
        commandBuffer,
        VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        0,
        0, nullptr,
        0, nullptr,
        1, &toCopySource);

    VkBufferImageCopy copyRegion = {};
    copyRegion.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    copyRegion.imageSubresource.mipLevel = 0;
    copyRegion.imageSubresource.baseArrayLayer = 0;
    copyRegion.imageSubresource.layerCount = 1;
    copyRegion.imageExtent.width = width;
    copyRegion.imageExtent.height = height;
    copyRegion.imageExtent.depth = 1;

    vkCmdCopyImageToBuffer(
        commandBuffer,
        texture->GetImage(),
        VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
        stagingBuffer,
        1,
        &copyRegion);

    VkImageMemoryBarrier backToPresent = toCopySource;
    backToPresent.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
    backToPresent.newLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
    backToPresent.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
    backToPresent.dstAccessMask = 0;

    vkCmdPipelineBarrier(
        commandBuffer,
        VK_PIPELINE_STAGE_TRANSFER_BIT,
        VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
        0,
        0, nullptr,
        0, nullptr,
        1, &backToPresent);

    vkEndCommandBuffer(commandBuffer);

    VkSubmitInfo submitInfo = {};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;

    vkQueueSubmit(vkDevice->GetGraphicsQueue(), 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(vkDevice->GetGraphicsQueue());

    void* mappedData = nullptr;
    if (vkMapMemory(logicalDevice, stagingMemory, 0, bufferSize, 0, &mappedData) != VK_SUCCESS) {
        vkDestroyCommandPool(logicalDevice, commandPool, nullptr);
        vkFreeMemory(logicalDevice, stagingMemory, nullptr);
        vkDestroyBuffer(logicalDevice, stagingBuffer, nullptr);
        return false;
    }

    FILE* fp = fopen(filename, "wb");
    if (fp == nullptr) {
        vkUnmapMemory(logicalDevice, stagingMemory);
        vkDestroyCommandPool(logicalDevice, commandPool, nullptr);
        vkFreeMemory(logicalDevice, stagingMemory, nullptr);
        vkDestroyBuffer(logicalDevice, stagingBuffer, nullptr);
        return false;
    }

    fprintf(fp, "P6\n%d %d\n255\n", width, height);

    const uint8_t* bytes = static_cast<const uint8_t*>(mappedData);
    const bool isBGRA = texture->GetVkFormat() == VK_FORMAT_B8G8R8A8_UNORM;
    for (uint32_t y = 0; y < height; ++y) {
        for (uint32_t x = 0; x < width; ++x) {
            const size_t index = (static_cast<size_t>(y) * width + x) * 4;
            const uint8_t r = isBGRA ? bytes[index + 2] : bytes[index + 0];
            const uint8_t g = bytes[index + 1];
            const uint8_t b = isBGRA ? bytes[index + 0] : bytes[index + 2];
            fwrite(&r, 1, 1, fp);
            fwrite(&g, 1, 1, fp);
            fwrite(&b, 1, 1, fp);
        }
    }

    fclose(fp);
    vkUnmapMemory(logicalDevice, stagingMemory);
    vkDestroyCommandPool(logicalDevice, commandPool, nullptr);
    vkFreeMemory(logicalDevice, stagingMemory, nullptr);
    vkDestroyBuffer(logicalDevice, stagingBuffer, nullptr);

    return true;
}

} // namespace RHI
} // namespace XCEngine