#include "V210ComputeDecoder.hpp"
#include "VulkanUtils.hpp"

#include <fstream>
#include <iostream>
#include <stdexcept>

namespace
{
    std::vector<char> readFile(const std::string& path)
    {
        std::ifstream file(path, std::ios::ate | std::ios::binary);

        if (!file.is_open())
        {
            throw std::runtime_error(
                "Failed to open shader file: " + path);
        }

        size_t fileSize = static_cast<size_t>(file.tellg());
        std::vector<char> buffer(fileSize);
        file.seekg(0);
        file.read(buffer.data(), static_cast<std::streamsize>(fileSize));
        file.close();

        return buffer;
    }

    VkShaderModule createShaderModule(
        VkDevice device,
        const std::vector<char>& code)
    {
        VkShaderModuleCreateInfo createInfo{};
        createInfo.sType =
            VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
        createInfo.codeSize = code.size();
        createInfo.pCode = reinterpret_cast<const uint32_t*>(
            code.data());

        VkShaderModule shaderModule;

        if (vkCreateShaderModule(
                device,
                &createInfo,
                nullptr,
                &shaderModule) != VK_SUCCESS)
        {
            throw std::runtime_error(
                "Failed to create shader module");
        }

        return shaderModule;
    }
}

void V210ComputeDecoder::init(
    VkDevice device,
    uint32_t feedCount,
    uint32_t dstWidth,
    uint32_t dstHeight)
{
    m_device = device;
    m_dstWidth = dstWidth;
    m_dstHeight = dstHeight;

    std::vector<VkDescriptorSetLayoutBinding> bindings(2);

    bindings[0].binding = 0;
    bindings[0].descriptorType =
        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    bindings[0].descriptorCount = 1;
    bindings[0].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
    bindings[0].pImmutableSamplers = nullptr;

    bindings[1].binding = 1;
    bindings[1].descriptorType =
        VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
    bindings[1].descriptorCount = 1;
    bindings[1].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
    bindings[1].pImmutableSamplers = nullptr;

    VkDescriptorSetLayoutCreateInfo layoutInfo{};
    layoutInfo.sType =
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    layoutInfo.bindingCount =
        static_cast<uint32_t>(bindings.size());
    layoutInfo.pBindings = bindings.data();

    if (vkCreateDescriptorSetLayout(
            device,
            &layoutInfo,
            nullptr,
            &m_descriptorSetLayout) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 descriptor set layout");
    }

    std::vector<char> shaderCode =
        readFile("../shaders/bin/v210_decode.comp.spv");

    VkShaderModule shaderModule =
        createShaderModule(device, shaderCode);

    VkPushConstantRange pushRange{};
    pushRange.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
    pushRange.offset = 0;
    pushRange.size = 5 * sizeof(uint32_t);

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType =
        VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts =
        &m_descriptorSetLayout;
    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = &pushRange;

    if (vkCreatePipelineLayout(
            device,
            &pipelineLayoutInfo,
            nullptr,
            &m_pipelineLayout) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 pipeline layout");
    }

    VkComputePipelineCreateInfo pipelineInfo{};
    pipelineInfo.sType =
        VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
    pipelineInfo.stage.sType =
        VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    pipelineInfo.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
    pipelineInfo.stage.module = shaderModule;
    pipelineInfo.stage.pName = "main";
    pipelineInfo.layout = m_pipelineLayout;

    if (vkCreateComputePipelines(
            device,
            VK_NULL_HANDLE,
            1,
            &pipelineInfo,
            nullptr,
            &m_pipeline) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 compute pipeline");
    }

    vkDestroyShaderModule(device, shaderModule, nullptr);

    VkDescriptorPoolSize poolSizes[2];
    poolSizes[0].type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    poolSizes[0].descriptorCount = feedCount;
    poolSizes[1].type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
    poolSizes[1].descriptorCount = feedCount;

    VkDescriptorPoolCreateInfo poolInfo{};
    poolInfo.sType =
        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    poolInfo.poolSizeCount = 2;
    poolInfo.pPoolSizes = poolSizes;
    poolInfo.maxSets = feedCount;

    if (vkCreateDescriptorPool(
            device,
            &poolInfo,
            nullptr,
            &m_descriptorPool) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 descriptor pool");
    }

    std::vector<VkDescriptorSetLayout> layouts(
        feedCount, m_descriptorSetLayout);

    VkDescriptorSetAllocateInfo allocInfo{};
    allocInfo.sType =
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = m_descriptorPool;
    allocInfo.descriptorSetCount = feedCount;
    allocInfo.pSetLayouts = layouts.data();

    m_descSets.resize(feedCount);

    if (vkAllocateDescriptorSets(
            device,
            &allocInfo,
            m_descSets.data()) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to allocate v210 descriptor sets");
    }
}

void V210ComputeDecoder::destroy()
{
    if (m_device == VK_NULL_HANDLE) return;

    vkDestroyPipeline(m_device, m_pipeline, nullptr);
    vkDestroyPipelineLayout(m_device, m_pipelineLayout, nullptr);
    vkDestroyDescriptorSetLayout(
        m_device, m_descriptorSetLayout, nullptr);
    vkDestroyDescriptorPool(m_device, m_descriptorPool, nullptr);
}

void V210ComputeDecoder::createFeedResources(
    VkPhysicalDevice physicalDevice,
    VkCommandPool commandPool,
    VkQueue graphicsQueue,
    uint32_t feedIndex,
    uint32_t srcWidth,
    uint32_t srcHeight,
    uint32_t srcStride,
    uint32_t dstWidth,
    uint32_t dstHeight,
    V210ComputeFeed& feed)
{
    feed.srcWidth = srcWidth;
    feed.srcHeight = srcHeight;
    feed.srcStride = srcStride;

    VkDeviceSize v210Size = static_cast<VkDeviceSize>(srcStride) * srcHeight;

    createBuffer(
        m_device,
        physicalDevice,
        v210Size,
        VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
        VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
        feed.v210Buffer,
        feed.v210BufferMemory);

    if (vkMapMemory(
            m_device,
            feed.v210BufferMemory,
            0,
            v210Size,
            0,
            &feed.v210MappedData) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to map v210 buffer memory");
    }

    createImage(
        m_device,
        physicalDevice,
        dstWidth,
        dstHeight,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_TILING_OPTIMAL,
        VK_IMAGE_USAGE_STORAGE_BIT |
            VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        feed.image,
        feed.imageMemory);

    transitionImageLayout(
        m_device,
        commandPool,
        graphicsQueue,
        feed.image,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_LAYOUT_UNDEFINED,
        VK_IMAGE_LAYOUT_GENERAL);

    VkImageViewCreateInfo viewInfo{};
    viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    viewInfo.image = feed.image;
    viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    viewInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    viewInfo.subresourceRange.aspectMask =
        VK_IMAGE_ASPECT_COLOR_BIT;
    viewInfo.subresourceRange.baseMipLevel = 0;
    viewInfo.subresourceRange.levelCount = 1;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount = 1;

    if (vkCreateImageView(
            m_device,
            &viewInfo,
            nullptr,
            &feed.imageView) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 output image view");
    }

    VkSamplerCreateInfo samplerInfo{};
    samplerInfo.sType =
        VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
    samplerInfo.magFilter = VK_FILTER_LINEAR;
    samplerInfo.minFilter = VK_FILTER_LINEAR;
    samplerInfo.addressModeU =
        VK_SAMPLER_ADDRESS_MODE_REPEAT;
    samplerInfo.addressModeV =
        VK_SAMPLER_ADDRESS_MODE_REPEAT;
    samplerInfo.addressModeW =
        VK_SAMPLER_ADDRESS_MODE_REPEAT;
    samplerInfo.anisotropyEnable = VK_FALSE;
    samplerInfo.borderColor =
        VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    samplerInfo.unnormalizedCoordinates = VK_FALSE;
    samplerInfo.compareEnable = VK_FALSE;
    samplerInfo.mipmapMode =
        VK_SAMPLER_MIPMAP_MODE_LINEAR;

    if (vkCreateSampler(
            m_device,
            &samplerInfo,
            nullptr,
            &feed.sampler) != VK_SUCCESS)
    {
        throw std::runtime_error(
            "Failed to create v210 output sampler");
    }

    feed.descriptorSet = m_descSets[feedIndex];

    VkDescriptorBufferInfo bufferInfo{};
    bufferInfo.buffer = feed.v210Buffer;
    bufferInfo.offset = 0;
    bufferInfo.range = v210Size;

    VkDescriptorImageInfo imageInfo{};
    imageInfo.imageLayout = VK_IMAGE_LAYOUT_GENERAL;
    imageInfo.imageView = feed.imageView;

    VkWriteDescriptorSet writes[2];

    writes[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[0].dstSet = feed.descriptorSet;
    writes[0].dstBinding = 0;
    writes[0].dstArrayElement = 0;
    writes[0].descriptorCount = 1;
    writes[0].descriptorType =
        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    writes[0].pBufferInfo = &bufferInfo;
    writes[0].pImageInfo = nullptr;
    writes[0].pTexelBufferView = nullptr;

    writes[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[1].dstSet = feed.descriptorSet;
    writes[1].dstBinding = 1;
    writes[1].dstArrayElement = 0;
    writes[1].descriptorCount = 1;
    writes[1].descriptorType =
        VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
    writes[1].pBufferInfo = nullptr;
    writes[1].pImageInfo = &imageInfo;
    writes[1].pTexelBufferView = nullptr;

    vkUpdateDescriptorSets(
        m_device, 2, writes, 0, nullptr);
}

void V210ComputeDecoder::destroyFeedResources(
    V210ComputeFeed& feed)
{
    if (m_device == VK_NULL_HANDLE) return;

    vkDestroySampler(m_device, feed.sampler, nullptr);
    vkDestroyImageView(m_device, feed.imageView, nullptr);
    vkDestroyImage(m_device, feed.image, nullptr);
    vkFreeMemory(m_device, feed.imageMemory, nullptr);

    vkUnmapMemory(m_device, feed.v210BufferMemory);
    vkDestroyBuffer(m_device, feed.v210Buffer, nullptr);
    vkFreeMemory(m_device, feed.v210BufferMemory, nullptr);
}

void V210ComputeDecoder::recordDecode(
    VkCommandBuffer cmdBuf,
    uint32_t feedIndex,
    uint32_t srcWidth,
    uint32_t srcHeight,
    uint32_t srcStrideBytes,
    uint32_t dstWidth,
    uint32_t dstHeight) const
{
    vkCmdBindPipeline(
        cmdBuf,
        VK_PIPELINE_BIND_POINT_COMPUTE,
        m_pipeline);

    vkCmdBindDescriptorSets(
        cmdBuf,
        VK_PIPELINE_BIND_POINT_COMPUTE,
        m_pipelineLayout,
        0,
        1,
        &m_descSets[feedIndex],
        0,
        nullptr);

    uint32_t pushConstants[5];
    pushConstants[0] = srcWidth;
    pushConstants[1] = srcHeight;
    pushConstants[2] = srcStrideBytes;
    pushConstants[3] = dstWidth;
    pushConstants[4] = dstHeight;

    vkCmdPushConstants(
        cmdBuf,
        m_pipelineLayout,
        VK_SHADER_STAGE_COMPUTE_BIT,
        0,
        sizeof(pushConstants),
        pushConstants);

    uint32_t groupCountX = (dstWidth + 15) / 16;
    uint32_t groupCountY = (dstHeight + 15) / 16;

    vkCmdDispatch(cmdBuf, groupCountX, groupCountY, 1);
}

void V210ComputeDecoder::recordOutputReadyForSampling(
    VkCommandBuffer cmdBuf,
    const V210ComputeFeed& feed) const
{
    VkImageMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    barrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL;
    barrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.image = feed.image;
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = 1;

    vkCmdPipelineBarrier(
        cmdBuf,
        VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
        0,
        0,
        nullptr,
        0,
        nullptr,
        1,
        &barrier);
}