// =============================================
//  Aster: ibl_helpers.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "ibl_helpers.h"

#include "device.h"
#include "gpu_resource_manager.h"
#include "helpers.h"
#include "image.h"
#include "asset_loader.h"
#include "pipeline_utils.h"

constexpr cstr EQUIRECT_TO_CUBE_SHADER_FILE = "shader/eqrectToCube.cs.hlsl.spv";

TextureHandle
CreateCubeFromHdrEnv(AssetLoader *assetLoader, vk::Queue computeQueue, const u32 cubeSide, TextureHandle hdrEnv,
                     const cstr name)
{
    GpuResourceManager *resMan = assetLoader->m_ResourceManager;
    const Device *pDevice = resMan->m_Device;

    TextureCube cubeMap;
    cubeMap.Init(pDevice, cubeSide, vk::Format::eR16G16B16A16Sfloat, false, name ? name : "Env CubeMap");

    StorageTexture stagingTexture;
    stagingTexture.Init(pDevice, {cubeSide * 3, cubeSide * 2}, vk::Format::eR16G16B16A16Sfloat, false, "EnvStaging");
    auto envStagingHandle = resMan->Commit(&stagingTexture);

    vk::ImageSubresourceRange stagingSubresRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
    };
    vk::ImageSubresourceRange cubeSubresRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 6,
    };

    vk::ImageMemoryBarrier2 readyToWriteBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eTopOfPipe,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eComputeShader,
        .dstAccessMask = vk::AccessFlagBits2::eShaderStorageWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eGeneral,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = stagingTexture.m_Image,
        .subresourceRange = stagingSubresRange,
    };
    vk::DependencyInfo readyToWriteDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &readyToWriteBarrier,
    };
    vk::ImageMemoryBarrier2 stagingReadyForTransfer = {
        .srcStageMask = vk::PipelineStageFlagBits2::eComputeShader,
        .srcAccessMask = vk::AccessFlagBits2::eShaderStorageWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferRead,
        .oldLayout = vk::ImageLayout::eGeneral,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = stagingTexture.m_Image,
        .subresourceRange = stagingSubresRange,
    };
    vk::ImageMemoryBarrier2 cubeReadyForTransfer = {
        .srcStageMask = vk::PipelineStageFlagBits2::eTopOfPipe,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eTransferDstOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = cubeMap.m_Image,
        .subresourceRange = cubeSubresRange,
    };
    eastl::array transferBarriers = {stagingReadyForTransfer, cubeReadyForTransfer};
    vk::DependencyInfo preTransferDependency = {
        .imageMemoryBarrierCount = transferBarriers.size(),
        .pImageMemoryBarriers = transferBarriers.data(),
    };
    vk::ImageMemoryBarrier2 cubemapToRead = {
        .srcStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eFragmentShader,
        .dstAccessMask = vk::AccessFlagBits2::eShaderSampledRead,
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = cubeMap.m_Image,
        .subresourceRange = cubeSubresRange,
    };
    vk::DependencyInfo cubemapToReadDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &cubemapToRead,
    };

    eastl::array<vk::ImageCopy2, 6> imageCopies;
    for (i32 i = 0; i < 6; ++i)
    {
        imageCopies[i] = vk::ImageCopy2{
            .srcSubresource =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = 0,
                    .layerCount = 1,
                },
            .srcOffset =
                {
                    (i % 3) * Cast<i32>(cubeSide),
                    (i / 3) * Cast<i32>(cubeSide),
                    0,
                },
            .dstSubresource =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = Cast<u32>(i),
                    .layerCount = 1,
                },
            .dstOffset = {0, 0, 0},
            .extent = {cubeSide, cubeSide, 1},
        };
    }

    vk::CopyImageInfo2 copyFlatToCube = {
        .srcImage = stagingTexture.m_Image,
        .srcImageLayout = vk::ImageLayout::eTransferSrcOptimal,
        .dstImage = cubeMap.m_Image,
        .dstImageLayout = vk::ImageLayout::eTransferDstOptimal,
        .regionCount = imageCopies.size(),
        .pRegions = imageCopies.data(),
    };

    struct WorkloadPushConstants
    {
        TextureHandle m_HdrEnvHandle;
        StorageTextureHandle m_OutputTexture;
        u32 m_CubeSide;
    };

#pragma region Pipeline Creation etc
    const auto shaderModule = CreateShader(pDevice, EQUIRECT_TO_CUBE_SHADER_FILE);

    const vk::PipelineShaderStageCreateInfo shaderStageCreateInfo = {
        .stage = vk::ShaderStageFlagBits::eCompute,
        .module = shaderModule,
        .pName = "main",
    };

    vk::PushConstantRange pcr = {
        .stageFlags = vk::ShaderStageFlagBits::eCompute,
        .offset = 0,
        .size = sizeof(WorkloadPushConstants),
    };

    vk::PipelineLayout pipelineLayout;
    const vk::PipelineLayoutCreateInfo layoutCreateInfo = {
        .setLayoutCount = 1,
        .pSetLayouts = &resMan->m_SetLayout,
        .pushConstantRangeCount = 1,
        .pPushConstantRanges = &pcr,
    };
    AbortIfFailed(pDevice->m_Device.createPipelineLayout(&layoutCreateInfo, nullptr, &pipelineLayout));

    const vk::ComputePipelineCreateInfo computePipelineCreateInfo = {
        .stage = shaderStageCreateInfo,
        .layout = pipelineLayout,
    };

    vk::Pipeline pipeline;
    AbortIfFailed(pDevice->m_Device.createComputePipelines(pDevice->m_PipelineCache, 1, &computePipelineCreateInfo,
                                                           nullptr, &pipeline));

    pDevice->m_Device.destroy(shaderModule, nullptr);

#pragma endregion

    WorkloadPushConstants pushConstants = {
        .m_HdrEnvHandle = hdrEnv, .m_OutputTexture = envStagingHandle, .m_CubeSide = cubeSide};

    resMan->Update();

    auto cmd = assetLoader->m_CommandBuffer;
    constexpr vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
    AbortIfFailed(cmd.begin(&beginInfo));

#if !defined(NDEBUG)
    StackString<128> labelName = "Eqrect -> Cubemap: ";
    labelName += name ? name : "<unknown env>";
    vk::DebugUtilsLabelEXT label = {
        .pLabelName = labelName.c_str(),
        .color = std::array{1.0f, 1.0f, 1.0f, 1.0f},
    };
    cmd.beginDebugUtilsLabelEXT(&label);
#endif

    cmd.pipelineBarrier2(&readyToWriteDependency);

    cmd.bindDescriptorSets(vk::PipelineBindPoint::eCompute, pipelineLayout, 0, 1, &resMan->m_DescriptorSet, 0, nullptr);
    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, pipeline);
    cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof pushConstants, &pushConstants);
    cmd.dispatch(cubeSide / 16, cubeSide / 16, 6);

    cmd.pipelineBarrier2(&preTransferDependency);

    cmd.copyImage2(&copyFlatToCube);

    cmd.pipelineBarrier2(&cubemapToReadDependency);

#if !defined(NDEBUG)
    cmd.endDebugUtilsLabelEXT();
#endif

    AbortIfFailed(cmd.end());

    vk::SubmitInfo submitInfo = {
        .waitSemaphoreCount = 0,
        .pWaitDstStageMask = nullptr,
        .commandBufferCount = 1,
        .pCommandBuffers = &cmd,
    };

    vk::Fence fence;
    vk::FenceCreateInfo fenceCreateInfo = {};
    AbortIfFailed(pDevice->m_Device.createFence(&fenceCreateInfo, nullptr, &fence));
    AbortIfFailed(computeQueue.submit(1, &submitInfo, fence));
    AbortIfFailed(pDevice->m_Device.waitForFences(1, &fence, true, MaxValue<u32>));
    pDevice->m_Device.destroy(fence, nullptr);

    AbortIfFailed(pDevice->m_Device.resetCommandPool(assetLoader->m_CommandPool, {}));

    resMan->Release(envStagingHandle);
    pDevice->m_Device.destroy(pipeline, nullptr);
    pDevice->m_Device.destroy(pipelineLayout, nullptr);

    return resMan->Commit(&cubeMap);
}