#pragma once

#include <XCEngine/Rendering/Builtin/BuiltinPassMetadataUtils.h>

#include <algorithm>

namespace XCEngine {
namespace Rendering {

inline bool TryBuildBuiltinPassResourceBindingPlan(
    const Containers::Array<Resources::ShaderResourceBindingDesc>& bindings,
    BuiltinPassResourceBindingPlan& outPlan,
    Containers::String* outError = nullptr) {
    outPlan = {};

    auto fail = [&outError](const char* message) {
        if (outError != nullptr) {
            *outError = message;
        }
        return false;
    };

    if (bindings.Empty()) {
        return true;
    }

    outPlan.bindings.Reserve(bindings.Size());
    Core::uint32 minBoundSet = UINT32_MAX;
    Core::uint32 maxBoundSet = 0;

    for (const Resources::ShaderResourceBindingDesc& binding : bindings) {
        const BuiltinPassResourceSemantic semantic = ResolveBuiltinPassResourceSemantic(binding);
        if (semantic == BuiltinPassResourceSemantic::Unknown) {
            return fail("Unsupported builtin pass resource semantic");
        }
        if (!IsBuiltinPassResourceTypeCompatible(semantic, binding.type)) {
            return fail("Builtin pass resource semantic/type combination is invalid");
        }

        PassResourceBindingLocation* location = nullptr;
        switch (semantic) {
        case BuiltinPassResourceSemantic::PerObject:
            location = &outPlan.perObject;
            break;
        case BuiltinPassResourceSemantic::Material:
            location = &outPlan.material;
            break;
        case BuiltinPassResourceSemantic::MaterialBuffer:
            break;
        case BuiltinPassResourceSemantic::Lighting:
            location = &outPlan.lighting;
            break;
        case BuiltinPassResourceSemantic::ShadowReceiver:
            location = &outPlan.shadowReceiver;
            break;
        case BuiltinPassResourceSemantic::Environment:
            location = &outPlan.environment;
            break;
        case BuiltinPassResourceSemantic::PassConstants:
            location = &outPlan.passConstants;
            break;
        case BuiltinPassResourceSemantic::VolumeField:
            location = &outPlan.volumeField;
            break;
        case BuiltinPassResourceSemantic::BaseColorTexture:
            location = &outPlan.baseColorTexture;
            break;
        case BuiltinPassResourceSemantic::SourceColorTexture:
            location = &outPlan.sourceColorTexture;
            break;
        case BuiltinPassResourceSemantic::SkyboxPanoramicTexture:
            location = &outPlan.skyboxPanoramicTexture;
            break;
        case BuiltinPassResourceSemantic::SkyboxTexture:
            location = &outPlan.skyboxTexture;
            break;
        case BuiltinPassResourceSemantic::LinearClampSampler:
            location = &outPlan.linearClampSampler;
            break;
        case BuiltinPassResourceSemantic::ShadowMapTexture:
            location = &outPlan.shadowMapTexture;
            break;
        case BuiltinPassResourceSemantic::ShadowMapSampler:
            location = &outPlan.shadowMapSampler;
            break;
        case BuiltinPassResourceSemantic::Unknown:
        default:
            break;
        }

        if (semantic != BuiltinPassResourceSemantic::MaterialBuffer && location == nullptr) {
            return fail("Builtin pass resource semantic could not be mapped");
        }
        if (location != nullptr && location->IsValid()) {
            return fail("Builtin pass resource semantic appears more than once");
        }

        for (const BuiltinPassResourceBindingDesc& existingBinding : outPlan.bindings) {
            if (existingBinding.location.set == binding.set &&
                existingBinding.location.binding == binding.binding) {
                return fail("Builtin pass resource set/binding pair appears more than once");
            }
        }

        const PassResourceBindingLocation resolvedLocation = { binding.set, binding.binding };
        if (location != nullptr) {
            *location = resolvedLocation;
        }

        BuiltinPassResourceBindingDesc resolvedBinding = {};
        resolvedBinding.name = binding.name;
        resolvedBinding.semantic = semantic;
        resolvedBinding.resourceType = binding.type;
        resolvedBinding.location = resolvedLocation;
        outPlan.bindings.PushBack(resolvedBinding);
        if (semantic == BuiltinPassResourceSemantic::MaterialBuffer) {
            outPlan.materialBufferBindings.PushBack(resolvedBinding);
            outPlan.usesMaterialBuffers = true;
        }

        outPlan.maxSetIndex = std::max(outPlan.maxSetIndex, binding.set);
        minBoundSet = std::min(minBoundSet, binding.set);
        maxBoundSet = std::max(maxBoundSet, binding.set);

        switch (binding.type) {
        case Resources::ShaderResourceType::ConstantBuffer:
            outPlan.usesConstantBuffers = true;
            break;
        case Resources::ShaderResourceType::Texture2D:
        case Resources::ShaderResourceType::TextureCube:
            outPlan.usesTextures = true;
            break;
        case Resources::ShaderResourceType::Sampler:
            outPlan.usesSamplers = true;
            break;
        default:
            break;
        }
    }

    outPlan.firstDescriptorSet = minBoundSet;
    outPlan.descriptorSetCount = maxBoundSet - minBoundSet + 1u;
    return true;
}

inline void AppendBuiltinPassResourceBinding(
    Containers::Array<Resources::ShaderResourceBindingDesc>& bindings,
    const char* name,
    Resources::ShaderResourceType type,
    Core::uint32 set,
    Core::uint32 binding,
    const char* semantic) {
    Resources::ShaderResourceBindingDesc desc = {};
    desc.name = name;
    desc.type = type;
    desc.set = set;
    desc.binding = binding;
    desc.semantic = semantic;
    bindings.PushBack(desc);
}

inline bool TryBuildBuiltinPassDefaultResourceBindings(
    const Resources::ShaderPass& shaderPass,
    Containers::Array<Resources::ShaderResourceBindingDesc>& outBindings) {
    outBindings.Clear();

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::ForwardLit)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "PerObjectConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PerObject");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LightingConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            1u,
            0u,
            "Lighting");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "MaterialConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            2u,
            0u,
            "Material");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "ShadowReceiverConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            3u,
            0u,
            "ShadowReceiver");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "BaseColorTexture",
            Resources::ShaderResourceType::Texture2D,
            4u,
            0u,
            "BaseColorTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LinearClampSampler",
            Resources::ShaderResourceType::Sampler,
            5u,
            0u,
            "LinearClampSampler");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "ShadowMapTexture",
            Resources::ShaderResourceType::Texture2D,
            6u,
            0u,
            "ShadowMapTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "ShadowMapSampler",
            Resources::ShaderResourceType::Sampler,
            7u,
            0u,
            "ShadowMapSampler");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::Unlit) ||
        ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::DepthOnly) ||
        ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::ShadowCaster)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "PerObjectConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PerObject");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "MaterialConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            1u,
            0u,
            "Material");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "BaseColorTexture",
            Resources::ShaderResourceType::Texture2D,
            2u,
            0u,
            "BaseColorTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LinearClampSampler",
            Resources::ShaderResourceType::Sampler,
            3u,
            0u,
            "LinearClampSampler");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::ObjectId)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "PerObjectConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PerObject");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::Skybox)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "EnvironmentConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "Environment");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "MaterialConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            1u,
            0u,
            "Material");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "SkyboxPanoramicTexture",
            Resources::ShaderResourceType::Texture2D,
            2u,
            0u,
            "SkyboxPanoramicTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "SkyboxTexture",
            Resources::ShaderResourceType::TextureCube,
            3u,
            0u,
            "SkyboxTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LinearClampSampler",
            Resources::ShaderResourceType::Sampler,
            4u,
            0u,
            "LinearClampSampler");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::Volumetric)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "PerObjectConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PerObject");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "MaterialConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            1u,
            0u,
            "Material");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::PostProcess)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "PostProcessConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PassConstants");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "SourceColorTexture",
            Resources::ShaderResourceType::Texture2D,
            1u,
            0u,
            "SourceColorTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LinearClampSampler",
            Resources::ShaderResourceType::Sampler,
            2u,
            0u,
            "LinearClampSampler");
        return true;
    }

    if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::FinalColor)) {
        AppendBuiltinPassResourceBinding(
            outBindings,
            "FinalColorConstants",
            Resources::ShaderResourceType::ConstantBuffer,
            0u,
            0u,
            "PassConstants");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "SourceColorTexture",
            Resources::ShaderResourceType::Texture2D,
            1u,
            0u,
            "SourceColorTexture");
        AppendBuiltinPassResourceBinding(
            outBindings,
            "LinearClampSampler",
            Resources::ShaderResourceType::Sampler,
            2u,
            0u,
            "LinearClampSampler");
        return true;
    }

    return false;
}

inline bool TryBuildBuiltinPassResourceBindingPlan(
    const Resources::ShaderPass& shaderPass,
    BuiltinPassResourceBindingPlan& outPlan,
    Containers::String* outError = nullptr) {
    if (shaderPass.resources.Empty()) {
        if (outError != nullptr) {
            *outError = "Builtin pass must declare explicit resource bindings";
        }
        outPlan = {};
        return false;
    }

    return TryBuildBuiltinPassResourceBindingPlan(shaderPass.resources, outPlan, outError);
}

inline RHI::DescriptorType ToBuiltinPassDescriptorType(Resources::ShaderResourceType type) {
    switch (type) {
    case Resources::ShaderResourceType::ConstantBuffer:
        return RHI::DescriptorType::CBV;
    case Resources::ShaderResourceType::Texture2D:
    case Resources::ShaderResourceType::TextureCube:
    case Resources::ShaderResourceType::StructuredBuffer:
    case Resources::ShaderResourceType::RawBuffer:
        return RHI::DescriptorType::SRV;
    case Resources::ShaderResourceType::RWStructuredBuffer:
    case Resources::ShaderResourceType::RWRawBuffer:
        return RHI::DescriptorType::UAV;
    case Resources::ShaderResourceType::Sampler:
        return RHI::DescriptorType::Sampler;
    default:
        return RHI::DescriptorType::CBV;
    }
}

inline RHI::ResourceViewDimension ToBuiltinPassResourceDimension(Resources::ShaderResourceType type) {
    switch (type) {
    case Resources::ShaderResourceType::Texture2D:
        return RHI::ResourceViewDimension::Texture2D;
    case Resources::ShaderResourceType::TextureCube:
        return RHI::ResourceViewDimension::TextureCube;
    case Resources::ShaderResourceType::StructuredBuffer:
    case Resources::ShaderResourceType::RWStructuredBuffer:
        return RHI::ResourceViewDimension::StructuredBuffer;
    case Resources::ShaderResourceType::RawBuffer:
    case Resources::ShaderResourceType::RWRawBuffer:
        return RHI::ResourceViewDimension::RawBuffer;
    default:
        return RHI::ResourceViewDimension::Unknown;
    }
}

inline RHI::DescriptorHeapType ResolveBuiltinPassDescriptorHeapType(Resources::ShaderResourceType type) {
    return type == Resources::ShaderResourceType::Sampler
        ? RHI::DescriptorHeapType::Sampler
        : RHI::DescriptorHeapType::CBV_SRV_UAV;
}

inline bool IsBuiltinPassShaderVisibleSet(const std::vector<RHI::DescriptorSetLayoutBinding>& bindings) {
    for (const RHI::DescriptorSetLayoutBinding& binding : bindings) {
        const RHI::DescriptorType descriptorType = static_cast<RHI::DescriptorType>(binding.type);
        if (descriptorType == RHI::DescriptorType::SRV ||
            descriptorType == RHI::DescriptorType::UAV ||
            descriptorType == RHI::DescriptorType::Sampler) {
            return true;
        }
    }

    return false;
}

inline Core::uint32 CountBuiltinPassHeapDescriptors(
    RHI::DescriptorHeapType heapType,
    const std::vector<RHI::DescriptorSetLayoutBinding>& bindings) {
    Core::uint32 descriptorCount = 0;
    for (const RHI::DescriptorSetLayoutBinding& binding : bindings) {
        const RHI::DescriptorType descriptorType = static_cast<RHI::DescriptorType>(binding.type);
        if (heapType == RHI::DescriptorHeapType::Sampler) {
            if (descriptorType == RHI::DescriptorType::Sampler) {
                descriptorCount += binding.count > 0 ? binding.count : 1u;
            }
            continue;
        }

        if (descriptorType != RHI::DescriptorType::Sampler) {
            descriptorCount += binding.count > 0 ? binding.count : 1u;
        }
    }

    return descriptorCount > 0 ? descriptorCount : 1u;
}

inline void RefreshBuiltinPassSetLayoutMetadata(BuiltinPassSetLayoutMetadata& setLayout) {
    std::sort(
        setLayout.bindings.begin(),
        setLayout.bindings.end(),
        [](const RHI::DescriptorSetLayoutBinding& left, const RHI::DescriptorSetLayoutBinding& right) {
            return left.binding < right.binding;
        });
    setLayout.shaderVisible = IsBuiltinPassShaderVisibleSet(setLayout.bindings);
    setLayout.layout.bindings = setLayout.bindings.empty() ? nullptr : setLayout.bindings.data();
    setLayout.layout.bindingCount = static_cast<Core::uint32>(setLayout.bindings.size());
}

inline void RefreshBuiltinPassSetLayouts(std::vector<BuiltinPassSetLayoutMetadata>& setLayouts) {
    for (BuiltinPassSetLayoutMetadata& setLayout : setLayouts) {
        RefreshBuiltinPassSetLayoutMetadata(setLayout);
    }
}

inline bool TryBuildBuiltinPassSetLayouts(
    const BuiltinPassResourceBindingPlan& bindingPlan,
    std::vector<BuiltinPassSetLayoutMetadata>& outSetLayouts,
    Containers::String* outError = nullptr) {
    outSetLayouts.clear();

    auto fail = [&outError](const char* message) {
        if (outError != nullptr) {
            *outError = message;
        }
        return false;
    };

    if (bindingPlan.bindings.Empty()) {
        return true;
    }

    outSetLayouts.resize(static_cast<size_t>(bindingPlan.maxSetIndex) + 1u);
    for (const BuiltinPassResourceBindingDesc& binding : bindingPlan.bindings) {
        if (binding.location.set >= outSetLayouts.size()) {
            return fail("Builtin pass encountered an invalid descriptor set index");
        }

        const RHI::DescriptorType descriptorType = ToBuiltinPassDescriptorType(binding.resourceType);
        const RHI::DescriptorHeapType heapType = ResolveBuiltinPassDescriptorHeapType(binding.resourceType);
        BuiltinPassSetLayoutMetadata& setLayout = outSetLayouts[binding.location.set];

        if (!setLayout.bindings.empty() && setLayout.heapType != heapType) {
            return fail("Builtin pass does not support mixing sampler and non-sampler bindings in one set");
        }

        for (const RHI::DescriptorSetLayoutBinding& existingBinding : setLayout.bindings) {
            if (existingBinding.binding == binding.location.binding) {
                return fail("Builtin pass encountered duplicate bindings inside one descriptor set");
            }
        }

        if (setLayout.bindings.empty()) {
            setLayout.heapType = heapType;
        }

        RHI::DescriptorSetLayoutBinding layoutBinding = {};
        layoutBinding.binding = binding.location.binding;
        layoutBinding.type = static_cast<Core::uint32>(descriptorType);
        layoutBinding.count = 1;
        layoutBinding.resourceDimension = ToBuiltinPassResourceDimension(binding.resourceType);
        setLayout.bindings.push_back(layoutBinding);

        switch (binding.semantic) {
        case BuiltinPassResourceSemantic::PerObject:
            setLayout.usesPerObject = true;
            break;
        case BuiltinPassResourceSemantic::Material:
            setLayout.usesMaterial = true;
            break;
        case BuiltinPassResourceSemantic::MaterialBuffer:
            setLayout.usesMaterialBuffers = true;
            setLayout.materialBufferBindings.push_back(binding);
            break;
        case BuiltinPassResourceSemantic::Lighting:
            setLayout.usesLighting = true;
            break;
        case BuiltinPassResourceSemantic::ShadowReceiver:
            setLayout.usesShadowReceiver = true;
            break;
        case BuiltinPassResourceSemantic::Environment:
            setLayout.usesEnvironment = true;
            break;
        case BuiltinPassResourceSemantic::PassConstants:
            setLayout.usesPassConstants = true;
            break;
        case BuiltinPassResourceSemantic::VolumeField:
            setLayout.usesVolumeField = true;
            break;
        case BuiltinPassResourceSemantic::BaseColorTexture:
            setLayout.usesTexture = true;
            setLayout.usesBaseColorTexture = true;
            break;
        case BuiltinPassResourceSemantic::SourceColorTexture:
            setLayout.usesTexture = true;
            setLayout.usesSourceColorTexture = true;
            break;
        case BuiltinPassResourceSemantic::SkyboxPanoramicTexture:
            setLayout.usesTexture = true;
            setLayout.usesSkyboxPanoramicTexture = true;
            break;
        case BuiltinPassResourceSemantic::SkyboxTexture:
            setLayout.usesTexture = true;
            setLayout.usesSkyboxTexture = true;
            break;
        case BuiltinPassResourceSemantic::ShadowMapTexture:
            setLayout.usesTexture = true;
            setLayout.usesShadowMapTexture = true;
            break;
        case BuiltinPassResourceSemantic::LinearClampSampler:
            setLayout.usesSampler = true;
            setLayout.usesLinearClampSampler = true;
            break;
        case BuiltinPassResourceSemantic::ShadowMapSampler:
            setLayout.usesSampler = true;
            setLayout.usesShadowMapSampler = true;
            break;
        case BuiltinPassResourceSemantic::Unknown:
        default:
            return fail("Builtin pass encountered an unsupported resource semantic");
        }
    }

    RefreshBuiltinPassSetLayouts(outSetLayouts);
    return true;
}

} // namespace Rendering
} // namespace XCEngine