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

#include "ibl_helpers.h"

#include "aster/core/device.h"
#include "aster/core/image.h"

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

#include <EASTL/fixed_vector.h>
#include <EASTL/tuple.h>

constexpr cstr EQUIRECT_TO_CUBE_SHADER_FILE = "shader/eqrect_to_cube.cs.hlsl.spv";
constexpr cstr DIFFUSE_IRRADIANCE_SHADER_FILE = "shader/diffuse_irradiance.cs.hlsl.spv";
constexpr cstr PREFILTER_SHADER_FILE = "shader/prefilter.cs.hlsl.spv";
constexpr cstr BRDF_LUT_SHADER_FILE = "shader/brdf_lut.cs.hlsl.spv";

void
Environment::Destroy(GpuResourceManager *resourceManager)
{
    resourceManager->Release(Take(m_Skybox));
    resourceManager->Release(Take(m_Diffuse));
    resourceManager->Release(Take(m_Prefilter));
    resourceManager->Release(Take(m_BrdfLut));
}

Environment
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;

    vk::SamplerCreateInfo brdfLutSamplerCreateInfo = resMan->m_DefaultSamplerCreateInfo;
    brdfLutSamplerCreateInfo.addressModeU = vk::SamplerAddressMode::eClampToEdge;
    brdfLutSamplerCreateInfo.addressModeV = vk::SamplerAddressMode::eClampToEdge;
    brdfLutSamplerCreateInfo.addressModeW = vk::SamplerAddressMode::eClampToEdge;

    StorageTextureCube skybox;
    StorageTextureCube diffuseIrradiance;
    StorageTextureCube prefilterCube;
    StorageTexture brdfLut;
    SamplerHandle brdfLutSampler;

    skybox.Init(pDevice, cubeSide, vk::Format::eR16G16B16A16Sfloat, true, true, "Skybox");
    TextureHandle skyboxHandle = resMan->CommitTexture(&skybox);
    StorageTextureHandle skyboxStorageHandle = resMan->CommitStorageTexture(&skybox);

    diffuseIrradiance.Init(pDevice, 64, vk::Format::eR16G16B16A16Sfloat, true, false, "Diffuse Irradiance");
    TextureHandle diffuseIrradianceHandle = resMan->CommitTexture(&diffuseIrradiance);
    StorageTextureHandle diffuseIrradianceStorageHandle = resMan->CommitStorageTexture(&diffuseIrradiance);

    prefilterCube.Init(pDevice, cubeSide, vk::Format::eR16G16B16A16Sfloat, true, true, "Prefilter");
    TextureHandle prefilterHandle = resMan->CommitTexture(&prefilterCube); // This stores the original view for us.
    constexpr u32 prefilterMipCountMax = 6;
    eastl::array<StorageTextureHandle, prefilterMipCountMax> prefilterStorageHandles;
    // All non-owning copies.
    for (u32 mipLevel = 0; auto &tex : prefilterStorageHandles)
    {
        vk::ImageViewCreateInfo imageViewCreateInfo = {
            .image = prefilterCube.m_Image,
            .viewType = vk::ImageViewType::eCube,
            .format = vk::Format::eR16G16B16A16Sfloat,
            .components = vk::ComponentMapping{},
            .subresourceRange =
                {
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .baseMipLevel = mipLevel++,
                    .levelCount = 1,
                    .baseArrayLayer = 0,
                    .layerCount = 6,
                },
        };
        AbortIfFailed(pDevice->m_Device.createImageView(&imageViewCreateInfo, nullptr, &prefilterCube.m_View));
        tex = resMan->CommitStorageTexture(&prefilterCube);
    }

    brdfLut.Init(pDevice, {512, 512}, vk::Format::eR16G16Sfloat, true, "BRDF LUT");
    brdfLutSampler = resMan->CreateSampler(&brdfLutSamplerCreateInfo);
    TextureHandle brdfLutHandle = resMan->CommitTexture(&brdfLut, brdfLutSampler);
    StorageTextureHandle brdfLutStorageHandle = resMan->CommitStorageTexture(&brdfLut);

#pragma region Dependencies and Copies
    vk::ImageSubresourceRange cubeSubresRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 6,
    };
    vk::ImageSubresourceRange lutSubresRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 0,
        .levelCount = 1,
        .baseArrayLayer = 0,
        .layerCount = 1,
    };

    vk::ImageMemoryBarrier2 readyToWriteBarrierTemplate = {
        .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,
        .subresourceRange = cubeSubresRange,
    };
    eastl::fixed_vector<vk::ImageMemoryBarrier2, 4> readyToWriteBarriers(4, readyToWriteBarrierTemplate);
    readyToWriteBarriers[0].image = skybox.m_Image;
    readyToWriteBarriers[1].image = diffuseIrradiance.m_Image;
    readyToWriteBarriers[2].image = prefilterCube.m_Image;
    readyToWriteBarriers[3].image = brdfLut.m_Image;
    readyToWriteBarriers[3].subresourceRange = lutSubresRange;

    vk::DependencyInfo readyToWriteDependency = {
        .imageMemoryBarrierCount = Cast<u32>(readyToWriteBarriers.size()),
        .pImageMemoryBarriers = readyToWriteBarriers.data(),
    };

    vk::ImageMemoryBarrier2 readyToSampleBarrierTemplate = {
        .srcStageMask = vk::PipelineStageFlagBits2::eComputeShader,
        .srcAccessMask = vk::AccessFlagBits2::eShaderStorageWrite | vk::AccessFlagBits2::eShaderStorageRead,
        .dstStageMask = vk::PipelineStageFlagBits2::eBottomOfPipe,
        .oldLayout = vk::ImageLayout::eGeneral,
        .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .subresourceRange = cubeSubresRange,
    };
    auto skyboxToSampleBarrier = readyToSampleBarrierTemplate;
    skyboxToSampleBarrier.image = skybox.m_Image;

    auto diffIrrToSampleBarrier = readyToSampleBarrierTemplate;
    diffIrrToSampleBarrier.image = diffuseIrradiance.m_Image;

    auto prefilterToSampleBarrier = readyToSampleBarrierTemplate;
    prefilterToSampleBarrier.image = prefilterCube.m_Image;

    auto brdfToSampleBarrier = readyToSampleBarrierTemplate;
    prefilterToSampleBarrier.image = brdfLut.m_Image;
    prefilterToSampleBarrier.subresourceRange = lutSubresRange;

    vk::DependencyInfo skyboxToSampleDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &skyboxToSampleBarrier,
    };
    vk::DependencyInfo diffIrrToSampleDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &diffIrrToSampleBarrier,
    };
    vk::DependencyInfo prefilterToSampleDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &prefilterToSampleBarrier,
    };
    vk::DependencyInfo brdfToSampleDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &brdfToSampleBarrier,
    };

#pragma endregion

    struct SkyboxPushConstants
    {
        TextureHandle m_HdrEnvHandle;
        StorageTextureHandle m_OutputTexture;
        u32 m_CubeSide;
    };
    struct DiffuseIrradiancePushConstants
    {
        TextureHandle m_SkyboxHandle;
        StorageTextureHandle m_OutputTexture;
        u32 m_CubeSide;
    };
    struct PrefilterPushConstants
    {
        TextureHandle m_SkyboxHandle;
        StorageTextureHandle m_OutputTexture;
        u32 m_CubeSide;
        f32 m_Roughness;
        u32 m_EnvSide;
    };
    struct BrdfLutPushConstants
    {
        StorageTextureHandle m_OutputTexture;
    };

#pragma region Pipeline Creation etc

    vk::PushConstantRange pcr = {
        .stageFlags = vk::ShaderStageFlagBits::eCompute,
        .offset = 0,
        .size = Cast<u32>(eastl::max(eastl::max(sizeof(SkyboxPushConstants), sizeof(BrdfLutPushConstants)),
                           eastl::max(sizeof(DiffuseIrradiancePushConstants), sizeof(PrefilterPushConstants)))),
    };

    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 auto eqRectToCubeShader = CreateShader(pDevice, EQUIRECT_TO_CUBE_SHADER_FILE);
    const auto diffuseRadianceShader = CreateShader(pDevice, DIFFUSE_IRRADIANCE_SHADER_FILE);
    const auto prefilterShader = CreateShader(pDevice, PREFILTER_SHADER_FILE);
    const auto brdfLutShader = CreateShader(pDevice, BRDF_LUT_SHADER_FILE);
    eastl::array computePipelineCreateInfo = {
        vk::ComputePipelineCreateInfo{
            .stage =
                {
                    .stage = vk::ShaderStageFlagBits::eCompute,
                    .module = eqRectToCubeShader,
                    .pName = "main",
                },
            .layout = pipelineLayout,
        },
        vk::ComputePipelineCreateInfo{
            .stage =
                {
                    .stage = vk::ShaderStageFlagBits::eCompute,
                    .module = diffuseRadianceShader,
                    .pName = "main",
                },
            .layout = pipelineLayout,
        },
        vk::ComputePipelineCreateInfo{
            .stage =
                {
                    .stage = vk::ShaderStageFlagBits::eCompute,
                    .module = prefilterShader,
                    .pName = "main",
                },
            .layout = pipelineLayout,
        },
        vk::ComputePipelineCreateInfo{
            .stage =
                {
                    .stage = vk::ShaderStageFlagBits::eCompute,
                    .module = brdfLutShader,
                    .pName = "main",
                },
            .layout = pipelineLayout,
        },
    };

    eastl::array<vk::Pipeline, computePipelineCreateInfo.size()> pipelines;
    AbortIfFailed(pDevice->m_Device.createComputePipelines(pDevice->m_PipelineCache, Cast<u32>(computePipelineCreateInfo.size()),
                                                           computePipelineCreateInfo.data(), nullptr,
                                                           pipelines.data()));

    vk::Pipeline eqRectToCubePipeline = pipelines[0];
    vk::Pipeline diffuseIrradiancePipeline = pipelines[1];
    vk::Pipeline prefilterPipeline = pipelines[2];
    vk::Pipeline brdfLutPipeline = pipelines[3];

    for (auto &createInfos : computePipelineCreateInfo)
    {
        pDevice->m_Device.destroy(createInfos.stage.module, nullptr);
    }

#pragma endregion

    SkyboxPushConstants skyboxPushConstant = {
        .m_HdrEnvHandle = hdrEnv,
        .m_OutputTexture = skyboxStorageHandle,
        .m_CubeSide = cubeSide,
    };
    DiffuseIrradiancePushConstants diffuseIrradiancePushConstants = {
        .m_SkyboxHandle = skyboxHandle,
        .m_OutputTexture = diffuseIrradianceStorageHandle,
        .m_CubeSide = diffuseIrradiance.m_Extent.width,
    };
    PrefilterPushConstants prefilterPushConstants = {
        .m_SkyboxHandle = skyboxHandle,
        .m_EnvSide = cubeSide,
    };
    BrdfLutPushConstants brdfLutPushConstants = {
        .m_OutputTexture = brdfLutStorageHandle,
    };

    resMan->Update();

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

#if !defined(ASTER_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, eqRectToCubePipeline);
    cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof skyboxPushConstant,
                      &skyboxPushConstant);
    assert(skybox.m_Extent.width % 16 == 0 && skybox.m_Extent.height % 16 == 0);
    cmd.dispatch(skybox.m_Extent.width / 16, skybox.m_Extent.height / 16, 6);

    GenerateMipMaps(cmd, &skybox, vk::ImageLayout::eGeneral, vk::ImageLayout::eGeneral,
                    vk::PipelineStageFlagBits2::eComputeShader, vk::PipelineStageFlagBits2::eComputeShader);

    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, diffuseIrradiancePipeline);
    cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof skyboxPushConstant,
                      &diffuseIrradiancePushConstants);
    assert(diffuseIrradiance.m_Extent.width % 16 == 0 && diffuseIrradiance.m_Extent.height % 16 == 0);
    cmd.dispatch(diffuseIrradiance.m_Extent.width / 16, diffuseIrradiance.m_Extent.width / 16, 6);

    cmd.pipelineBarrier2(&diffIrrToSampleDependency);

    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, prefilterPipeline);
    u32 mipSize = prefilterCube.m_Extent.width;
    assert(mipSize % 16 == 0);
    for (u32 mipCount = 0; auto &tex : prefilterStorageHandles)
    {
        prefilterPushConstants.m_OutputTexture = tex;
        prefilterPushConstants.m_CubeSide = mipSize;
        prefilterPushConstants.m_Roughness = Cast<f32>(mipCount) / Cast<f32>(prefilterMipCountMax);
        cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof prefilterPushConstants,
                          &prefilterPushConstants);
        u32 groupCount = eastl::max(mipSize / 16u, 1u);
        cmd.dispatch(groupCount, groupCount, 6);

        ++mipCount;
        mipSize = mipSize >> 1;
    }

    cmd.pipelineBarrier2(&skyboxToSampleDependency);
    cmd.pipelineBarrier2(&prefilterToSampleDependency);

    cmd.bindPipeline(vk::PipelineBindPoint::eCompute, brdfLutPipeline);
    cmd.pushConstants(pipelineLayout, vk::ShaderStageFlagBits::eCompute, 0, sizeof brdfLutPushConstants,
                      &brdfLutPushConstants);
    assert(brdfLut.m_Extent.width % 16 == 0 && brdfLut.m_Extent.height % 16 == 0);
    cmd.dispatch(brdfLut.m_Extent.width / 16, brdfLut.m_Extent.height / 16, 1);

#if !defined(ASTER_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, {}));

    skybox = {};
    resMan->Release(skyboxStorageHandle);
    resMan->Release(diffuseIrradianceStorageHandle);
    resMan->Release(brdfLutStorageHandle);
    for (auto &texHandles : prefilterStorageHandles)
    {
        StorageTextureCube st;
        resMan->Release(&st, texHandles);
        pDevice->m_Device.destroy(st.m_View, nullptr);
    }
    for (auto &pipeline : pipelines)
    {
        pDevice->m_Device.destroy(pipeline, nullptr);
    }
    pDevice->m_Device.destroy(pipelineLayout, nullptr);

    return {
        .m_Skybox = skyboxHandle,
        .m_Diffuse = diffuseIrradianceHandle,
        .m_Prefilter = prefilterHandle,
        .m_BrdfLut = brdfLutHandle,
    };
}