#include "FeedTexture.hpp"

#include "VulkanUtils.hpp"

#include <cstring>
#include <stdexcept>

void createFeedTexture(
    VkDevice device,
    VkPhysicalDevice physicalDevice,
    VkCommandPool commandPool,
    VkQueue graphicsQueue,
    const VideoFrame& initialFrame,
    VkDeviceSize imageSize,
    FeedTexture& texture)
{
    createBuffer(
        device,
        physicalDevice,
        imageSize,
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
        VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
        texture.stagingBuffer,
        texture.stagingBufferMemory
    );

    if (vkMapMemory(
            device,
            texture.stagingBufferMemory,
            0,
            imageSize,
            0,
            &texture.mappedData) != VK_SUCCESS)
    {
        throw std::runtime_error("Failed to map staging buffer memory");
    }

    std::memcpy(
        texture.mappedData,
        initialFrame.pixels.data(),
        static_cast<size_t>(imageSize)
    );

    createImage(
        device,
        physicalDevice,
        initialFrame.width,
        initialFrame.height,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_TILING_OPTIMAL,
        VK_IMAGE_USAGE_TRANSFER_DST_BIT |
        VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
        texture.image,
        texture.imageMemory
    );

    transitionImageLayout(
        device,
        commandPool,
        graphicsQueue,
        texture.image,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_LAYOUT_UNDEFINED,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    );

    copyBufferToImage(
        device,
        commandPool,
        graphicsQueue,
        texture.stagingBuffer,
        texture.image,
        initialFrame.width,
        initialFrame.height
    );

    transitionImageLayout(
        device,
        commandPool,
        graphicsQueue,
        texture.image,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
    );

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

    if (vkCreateImageView(
            device,
            &textureViewInfo,
            nullptr,
            &texture.imageView) != VK_SUCCESS)
    {
        throw std::runtime_error("Failed to create texture 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(
            device,
            &samplerInfo,
            nullptr,
            &texture.sampler) != VK_SUCCESS)
    {
        throw std::runtime_error("Failed to create texture sampler");
    }
}

void updateFeedDescriptorSet(
    VkDevice device,
    const FeedTexture& texture)
{
    updateSampledImageDescriptorSet(
        device,
        texture.descriptorSet,
        texture.imageView,
        texture.sampler,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
    );
}

void updateSampledImageDescriptorSet(
    VkDevice device,
    VkDescriptorSet descriptorSet,
    VkImageView imageView,
    VkSampler sampler,
    VkImageLayout imageLayout)
{
    VkDescriptorImageInfo imageInfo{};
    imageInfo.imageLayout = imageLayout;
    imageInfo.imageView = imageView;
    imageInfo.sampler = sampler;

    VkWriteDescriptorSet descriptorWrite{};
    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = descriptorSet;
    descriptorWrite.dstBinding = 0;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType =
        VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pImageInfo = &imageInfo;

    vkUpdateDescriptorSets(
        device,
        1,
        &descriptorWrite,
        0,
        nullptr
    );
}

void copyFrameToFeedTextureStaging(
    const VideoFrame& frame,
    VkDeviceSize imageSize,
    FeedTexture& texture)
{
    std::memcpy(
        texture.mappedData,
        frame.pixels.data(),
        static_cast<size_t>(imageSize)
    );
}

void recordFeedTextureUpload(
    VkCommandBuffer commandBuffer,
    const FeedTexture& texture,
    uint32_t width,
    uint32_t height)
{
    recordImageLayoutTransition(
        commandBuffer,
        texture.image,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
    );

    recordCopyBufferToImage(
        commandBuffer,
        texture.stagingBuffer,
        texture.image,
        width,
        height
    );

    recordImageLayoutTransition(
        commandBuffer,
        texture.image,
        VK_FORMAT_R8G8B8A8_UNORM,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
    );
}

void destroyFeedTexture(
    VkDevice device,
    FeedTexture& texture)
{
    vkDestroySampler(device, texture.sampler, nullptr);

    vkDestroyImageView(device, texture.imageView, nullptr);

    vkDestroyImage(device, texture.image, nullptr);
    vkFreeMemory(device, texture.imageMemory, nullptr);

    vkUnmapMemory(device, texture.stagingBufferMemory);

    vkDestroyBuffer(device, texture.stagingBuffer, nullptr);
    vkFreeMemory(device, texture.stagingBufferMemory, nullptr);
}