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

#include "ibl_helpers.h"

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

#include "asset_loader.h"
#include "helpers.h"

#include "aster/systems/commit_manager.h"
#include "aster/systems/rendering_device.h"

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

constexpr auto EQUIRECT_TO_CUBE_SHADER_FILE = "eqrect_to_cube";
constexpr auto ENVIRONMENT_SHADER_FILE = "environment";
constexpr auto DIFFUSE_IRRADIANCE_ENTRY = "diffuseIrradiance";
constexpr auto PREFILTER_ENTRY = "prefilter";
constexpr auto BRDF_LUT_ENTRY = "brdfLut";

Environment
CreateCubeFromHdrEnv(AssetLoader &assetLoader, u32 const cubeSide, systems::ResId<TextureView> hdrEnv)
{
    systems::RenderingDevice &device = *assetLoader.m_Device;
    auto *commitManager = device.m_CommitManager.get();

    auto skybox = device.CreateTextureCubeWithView<StorageTextureCubeView>({
        .m_Format = vk::Format::eR16G16B16A16Sfloat,
        .m_Side = cubeSide,
        .m_Name = "Skybox",
        .m_IsSampled = true,
        .m_IsMipMapped = true,
        .m_IsStorage = true,
    });

    auto skyboxHandle = commitManager->CommitTexture(skybox);
    auto skyboxStorageHandle = commitManager->CommitStorageImage(skybox);

    auto diffuseIrradiance = device.CreateTextureCubeWithView<StorageTextureCubeView>({
        .m_Format = vk::Format::eR16G16B16A16Sfloat,
        .m_Side = 64,
        .m_Name = "Diffuse Irradiance",
        .m_IsSampled = true,
        .m_IsMipMapped = false,
        .m_IsStorage = true,
    });
    auto diffuseIrradianceHandle = commitManager->CommitTexture(diffuseIrradiance);
    auto diffuseIrradianceStorageHandle = commitManager->CommitStorageImage(diffuseIrradiance);

    auto prefilterCube = device.CreateTextureCubeWithView<StorageTextureCubeView>({
        .m_Format = vk::Format::eR16G16B16A16Sfloat,
        .m_Side = cubeSide,
        .m_Name = "Prefilter",
        .m_IsSampled = true,
        .m_IsMipMapped = true,
        .m_IsStorage = true,
    });
    auto prefilterHandle = commitManager->CommitTexture(prefilterCube); // This stores the original view for us.
    constexpr u32 prefilterMipCountMax = 6;
    eastl::fixed_vector<systems::ResId<StorageImageView>, prefilterMipCountMax> prefilterStorageHandles;
    // All non-owning copies.
    for (u8 mipLevel = 0; mipLevel < prefilterMipCountMax; ++mipLevel)
    {
        auto view = device.CreateView<StorageTextureCubeView>({
            .m_Image = systems::CastImage<StorageTextureCube>(prefilterCube->m_Image),
            .m_ViewType = vk::ImageViewType::eCube,
            .m_AspectMask = vk::ImageAspectFlagBits::eColor,
            .m_MipLevelCount = 1,
            .m_LayerCount = 6,
            .m_BaseMipLevel = mipLevel,
            .m_BaseLayer = 0,
        });
        prefilterStorageHandles.push_back(commitManager->CommitStorageImage(view));
    }

    auto brdfLut = device.CreateTexture2DWithView<StorageTextureView>({
        .m_Format = vk::Format::eR16G16Sfloat,
        .m_Extent = {512, 512},
        .m_Name = "BRDF LUT",
        .m_IsSampled = true,
        .m_IsMipMapped = false,
        .m_IsStorage = true,
    });

    auto brdfLutSampler = device.CreateSampler({
        .m_AddressModeU = vk::SamplerAddressMode::eClampToEdge,
        .m_AddressModeV = vk::SamplerAddressMode::eClampToEdge,
        .m_AddressModeW = vk::SamplerAddressMode::eClampToEdge,
    });
    auto brdfLutHandle = commitManager->CommitTexture(brdfLut, brdfLutSampler);
    auto brdfLutStorageHandle = commitManager->CommitStorageImage(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->GetImage();
    readyToWriteBarriers[1].image = diffuseIrradiance->GetImage();
    readyToWriteBarriers[2].image = prefilterCube->GetImage();
    readyToWriteBarriers[3].image = brdfLut->GetImage();
    readyToWriteBarriers[3].subresourceRange = lutSubresRange;

    vk::DependencyInfo readyToWriteDependency = {
        .imageMemoryBarrierCount = static_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->GetImage();

    auto diffIrrToSampleBarrier = readyToSampleBarrierTemplate;
    diffIrrToSampleBarrier.image = diffuseIrradiance->GetImage();

    auto prefilterToSampleBarrier = readyToSampleBarrierTemplate;
    prefilterToSampleBarrier.image = prefilterCube->GetImage();

    auto brdfToSampleBarrier = readyToSampleBarrierTemplate;
    prefilterToSampleBarrier.image = brdfLut->GetImage();
    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
    {
        systems::ResId<TextureView> m_HdrEnvHandle;
        systems::ResId<StorageImageView> m_OutputTexture;
        u32 m_CubeSide;
    };
    struct DiffuseIrradiancePushConstants
    {
        systems::ResId<TextureView> m_SkyboxHandle;
        systems::ResId<StorageImageView> m_OutputTexture;
        u32 m_CubeSide;
    };
    struct PrefilterPushConstants
    {
        systems::ResId<TextureView> m_SkyboxHandle;
        systems::ResId<StorageImageView> m_OutputTexture;
        u32 m_CubeSide;
        f32 m_Roughness;
        u32 m_EnvSide;
    };
    struct BrdfLutPushConstants
    {
        systems::ResId<StorageImageView> m_OutputTexture;
    };

#pragma region Pipeline Creation etc

    // vk::PushConstantRange pcr = {
    //     .stageFlags = vk::ShaderStageFlagBits::eCompute,
    //     .offset = 0,
    //     .size = static_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 = &commitManager->GetDescriptorSetLayout(),
    //     .pushConstantRangeCount = 1,
    //     .pPushConstantRanges = &pcr,
    // };
    // AbortIfFailed(device.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, static_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];

    Pipeline eqRectToCubePipeline;

    if (auto result =
            device.CreateComputePipeline(eqRectToCubePipeline, {
                                                                   .m_Shader =
                                                                       {
                                                                           .m_ShaderFile = EQUIRECT_TO_CUBE_SHADER_FILE,
                                                                           .m_EntryPoints = {"main"},
                                                                       },
                                                                   .m_Name = "EqRect -> Cubemap",
                                                               }))
    {
        ERROR("EqRect -> Cubemap Pipeline Creation failed. Cause: {}", result.What()) THEN_ABORT(result.Value());
    }

    Pipeline diffuseIrradiancePipeline;

    if (auto result = device.CreateComputePipeline(
            diffuseIrradiancePipeline,
            {{.m_ShaderFile = ENVIRONMENT_SHADER_FILE, .m_EntryPoints = {DIFFUSE_IRRADIANCE_ENTRY}},
             "DiffuseIrradiance"}))
    {
        ERROR("Diffuse Irradiance compute pipeline creation failed. Cause: {}", result.What())
        THEN_ABORT(result.Value());
    }

    Pipeline prefilterPipeline;

    if (auto result = device.CreateComputePipeline(
            prefilterPipeline,
            {{.m_ShaderFile = ENVIRONMENT_SHADER_FILE, .m_EntryPoints = {PREFILTER_ENTRY}}, "Prefilter"}))
    {
        ERROR("Prefilter compute pipeline creation failed. Cause: {}", result.What())
        THEN_ABORT(result.Value());
    }

    Pipeline brdfLutPipeline;

    if (auto result = device.CreateComputePipeline(
            brdfLutPipeline, {{.m_ShaderFile = ENVIRONMENT_SHADER_FILE, .m_EntryPoints = {BRDF_LUT_ENTRY}}, "BRDF"}))
    {
        ERROR("BRDF LUT compute pipeline creation failed. Cause: {}", result.What())
        THEN_ABORT(result.Value());
    }

#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_OutputTexture = systems::NullId{},
        .m_EnvSide = cubeSide,
    };
    BrdfLutPushConstants brdfLutPushConstants = {
        .m_OutputTexture = brdfLutStorageHandle,
    };

    commitManager->Update();

    auto context = assetLoader.m_Device->CreateComputeContext();

    context.Begin();

    context.BeginDebugRegion("Eqrect -> Cubemap");

    context.Dependency(readyToWriteDependency);

    assert(skybox->m_Extent.width % 16 == 0 && skybox->m_Extent.height % 16 == 0);
    context.Dispatch(eqRectToCubePipeline, skybox->m_Extent.width / 16, skybox->m_Extent.height / 16, 6,
                     skyboxPushConstant);

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

    assert(diffuseIrradiance->m_Extent.width % 16 == 0 && diffuseIrradiance->m_Extent.height % 16 == 0);
    context.Dispatch(diffuseIrradiancePipeline, diffuseIrradiance->m_Extent.width / 16,
                     diffuseIrradiance->m_Extent.width / 16, 6, diffuseIrradiancePushConstants);

    context.Dependency(diffIrrToSampleDependency);

    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 = static_cast<f32>(mipCount) / static_cast<f32>(prefilterMipCountMax - 1);

        u32 groupCount = eastl::max(mipSize / 16u, 1u);
        context.Dispatch(prefilterPipeline, groupCount, groupCount, 6, prefilterPushConstants);

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

    context.Dependency(skyboxToSampleDependency);
    context.Dependency(prefilterToSampleDependency);

    assert(brdfLut->m_Extent.width % 16 == 0 && brdfLut->m_Extent.height % 16 == 0);
    context.Dispatch(brdfLutPipeline, brdfLut->m_Extent.width / 16, brdfLut->m_Extent.height / 16, 1,
                     brdfLutPushConstants);

    context.EndDebugRegion();

    context.End();

    auto receipt = device.Submit(context);
    device.WaitOn(receipt);

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