#include "Rendering/Passes/BuiltinDepthStylePassBase.h"

#include "Components/GameObject.h"
#include "Core/Asset/ResourceManager.h"
#include "Debug/Logger.h"
#include "RHI/RHICommandList.h"
#include "Rendering/BuiltinPassLayoutUtils.h"
#include "Rendering/Detail/ShaderVariantUtils.h"
#include "Rendering/RenderMaterialResolve.h"
#include "Rendering/Extraction/RenderSceneExtractor.h"
#include "Rendering/RenderSurface.h"
#include "Resources/Material/Material.h"
#include "Resources/Mesh/Mesh.h"

#include <vector>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

bool IsSupportedPerObjectOnlyBindingPlan(const BuiltinPassResourceBindingPlan& bindingPlan) {
    return bindingPlan.perObject.IsValid() &&
           bindingPlan.bindings.Size() == 1u &&
           bindingPlan.descriptorSetCount == 1u &&
           bindingPlan.usesConstantBuffers &&
           !bindingPlan.usesTextures &&
           !bindingPlan.usesSamplers;
}

uint32_t ResolveSurfaceColorAttachmentCount(const RenderSurface& surface) {
    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface.GetColorAttachments();
    return (!colorAttachments.empty() && colorAttachments[0] != nullptr) ? 1u : 0u;
}

RHI::Format ResolveSurfaceColorFormat(const RenderSurface& surface) {
    const std::vector<RHI::RHIResourceView*>& colorAttachments = surface.GetColorAttachments();
    if (colorAttachments.empty() || colorAttachments[0] == nullptr) {
        return RHI::Format::Unknown;
    }

    return colorAttachments[0]->GetFormat();
}

RHI::Format ResolveSurfaceDepthFormat(const RenderSurface& surface) {
    if (RHI::RHIResourceView* depthAttachment = surface.GetDepthAttachment();
        depthAttachment != nullptr) {
        return depthAttachment->GetFormat();
    }

    return RHI::Format::Unknown;
}

RHI::GraphicsPipelineDesc CreatePipelineDesc(
    RHI::RHIType backendType,
    RHI::RHIPipelineLayout* pipelineLayout,
    const Resources::Shader& shader,
    const Containers::String& passName,
    const Resources::Material* material,
    const RenderSurface& surface,
    const RHI::InputLayoutDesc& inputLayout) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    pipelineDesc.renderTargetCount = ResolveSurfaceColorAttachmentCount(surface);
    if (pipelineDesc.renderTargetCount > 0) {
        pipelineDesc.renderTargetFormats[0] =
            static_cast<uint32_t>(ResolveSurfaceColorFormat(surface));
    }
    pipelineDesc.depthStencilFormat =
        static_cast<uint32_t>(ResolveSurfaceDepthFormat(surface));
    pipelineDesc.sampleCount = 1;
    pipelineDesc.inputLayout = inputLayout;
    ApplyMaterialRenderState(material, pipelineDesc);

    pipelineDesc.blendState.blendEnable = false;
    pipelineDesc.blendState.colorWriteMask = pipelineDesc.renderTargetCount > 0 ? 0xF : 0;
    pipelineDesc.depthStencilState.depthTestEnable = true;
    pipelineDesc.depthStencilState.depthWriteEnable = true;
    pipelineDesc.depthStencilState.depthFunc = static_cast<uint32_t>(RHI::ComparisonFunc::LessEqual);

    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Detail::ToShaderBackend(backendType);
    if (const Resources::ShaderStageVariant* vertexVariant =
            shader.FindVariant(passName, Resources::ShaderType::Vertex, backend)) {
        ::XCEngine::Rendering::Detail::ApplyShaderStageVariant(*vertexVariant, pipelineDesc.vertexShader);
    }
    if (const Resources::ShaderStageVariant* fragmentVariant =
            shader.FindVariant(passName, Resources::ShaderType::Fragment, backend)) {
        ::XCEngine::Rendering::Detail::ApplyShaderStageVariant(*fragmentVariant, pipelineDesc.fragmentShader);
    }

    return pipelineDesc;
}

} // namespace

bool BuiltinDepthStylePassBase::EnsureInitialized(const RenderContext& context) {
    if (!context.IsValid()) {
        return false;
    }

    if (m_device == context.device &&
        m_backendType == context.backendType &&
        m_builtinShader.IsValid()) {
        return true;
    }

    DestroyResources();
    return CreateResources(context);
}

bool BuiltinDepthStylePassBase::CreateResources(const RenderContext& context) {
    m_device = context.device;
    m_backendType = context.backendType;
    m_builtinShader = Resources::ResourceManager::Get().Load<Resources::Shader>(m_builtinShaderPath);
    if (!m_builtinShader.IsValid()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase failed to load builtin shader resource: ") + m_builtinShaderPath).CStr());
        DestroyResources();
        return false;
    }

    return true;
}

void BuiltinDepthStylePassBase::DestroyResources() {
    m_resourceCache.Shutdown();

    for (auto& descriptorSetEntry : m_perObjectSets) {
        DestroyOwnedDescriptorSet(descriptorSetEntry.second);
    }
    m_perObjectSets.clear();

    for (auto& pipelineEntry : m_pipelineStates) {
        if (pipelineEntry.second != nullptr) {
            pipelineEntry.second->Shutdown();
            delete pipelineEntry.second;
        }
    }
    m_pipelineStates.clear();

    for (auto& layoutEntry : m_passResourceLayouts) {
        DestroyPassResourceLayout(layoutEntry.second);
    }
    m_passResourceLayouts.clear();

    m_device = nullptr;
    m_backendType = RHI::RHIType::D3D12;
    m_builtinShader.Reset();
}

BuiltinDepthStylePassBase::ResolvedShaderPass BuiltinDepthStylePassBase::ResolveSurfaceShaderPass(
    const Resources::Material* material) const {
    ResolvedShaderPass resolved = {};
    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Detail::ToShaderBackend(m_backendType);

    auto tryResolveFromShader =
        [this, backend, &resolved](
            const Resources::Shader* shader,
            const Resources::Material* ownerMaterial) -> bool {
            if (shader == nullptr) {
                return false;
            }

            auto tryAcceptPass =
                [this, shader, &resolved](const Resources::ShaderPass& shaderPass) -> bool {
                    BuiltinPassResourceBindingPlan bindingPlan = {};
                    Containers::String error;
                    if (!TryBuildSupportedBindingPlan(shaderPass, bindingPlan, &error)) {
                        return false;
                    }

                    resolved.shader = shader;
                    resolved.pass = &shaderPass;
                    resolved.passName = shaderPass.name;
                    return true;
                };

            if (ownerMaterial != nullptr && !ownerMaterial->GetShaderPass().Empty()) {
                const Resources::ShaderPass* explicitPass = shader->FindPass(ownerMaterial->GetShaderPass());
                if (explicitPass != nullptr &&
                    ShaderPassMatchesBuiltinPass(*explicitPass, m_passType) &&
                    ::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(
                        *shader,
                        explicitPass->name,
                        backend) &&
                    tryAcceptPass(*explicitPass)) {
                    return true;
                }
            }

            for (const Resources::ShaderPass& shaderPass : shader->GetPasses()) {
                if (!ShaderPassMatchesBuiltinPass(shaderPass, m_passType) ||
                    !::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(*shader, shaderPass.name, backend)) {
                    continue;
                }

                if (tryAcceptPass(shaderPass)) {
                    return true;
                }
            }

            return false;
        };

    if (material != nullptr &&
        material->GetShader() != nullptr &&
        tryResolveFromShader(material->GetShader(), material)) {
        return resolved;
    }

    if (material != nullptr && IsTransparentRenderQueue(ResolveMaterialRenderQueue(material))) {
        return {};
    }

    if (m_builtinShader.IsValid() &&
        tryResolveFromShader(m_builtinShader.Get(), nullptr)) {
        return resolved;
    }

    return {};
}

bool BuiltinDepthStylePassBase::TryBuildSupportedBindingPlan(
    const Resources::ShaderPass& shaderPass,
    BuiltinPassResourceBindingPlan& outPlan,
    Containers::String* outError) const {
    if (shaderPass.resources.Empty()) {
        if (outError != nullptr) {
            *outError =
                Containers::String("Builtin depth-style pass requires explicit resource bindings on shader pass: ") +
                shaderPass.name;
        }
        return false;
    }

    if (!TryBuildBuiltinPassResourceBindingPlan(shaderPass.resources, outPlan, outError)) {
        return false;
    }

    if (!IsSupportedPerObjectOnlyBindingPlan(outPlan)) {
        if (outError != nullptr) {
            *outError = "Builtin depth-style pass currently requires exactly one PerObject constant-buffer binding";
        }
        return false;
    }

    return true;
}

BuiltinDepthStylePassBase::PassResourceLayout* BuiltinDepthStylePassBase::GetOrCreatePassResourceLayout(
    const RenderContext& context,
    const ResolvedShaderPass& resolvedShaderPass) {
    if (resolvedShaderPass.shader == nullptr || resolvedShaderPass.pass == nullptr) {
        return nullptr;
    }

    PassLayoutKey passLayoutKey = {};
    passLayoutKey.shader = resolvedShaderPass.shader;
    passLayoutKey.passName = resolvedShaderPass.passName;

    const auto existing = m_passResourceLayouts.find(passLayoutKey);
    if (existing != m_passResourceLayouts.end()) {
        return &existing->second;
    }

    PassResourceLayout passLayout = {};
    auto failLayout = [this, &passLayout](const char* message) -> PassResourceLayout* {
        Debug::Logger::Get().Error(Debug::LogCategory::Rendering, message);
        DestroyPassResourceLayout(passLayout);
        return nullptr;
    };

    BuiltinPassResourceBindingPlan bindingPlan = {};
    Containers::String bindingPlanError;
    if (!TryBuildSupportedBindingPlan(*resolvedShaderPass.pass, bindingPlan, &bindingPlanError)) {
        const Containers::String contextualError =
            Containers::String("BuiltinDepthStylePassBase failed to resolve pass resource bindings for shader='") +
            resolvedShaderPass.shader->GetPath() +
            "', pass='" + resolvedShaderPass.passName +
            "': " + bindingPlanError +
            ". Bindings: " + DescribeShaderResourceBindings(resolvedShaderPass.pass->resources);
        return failLayout(contextualError.CStr());
    }

    std::vector<BuiltinPassSetLayoutMetadata> setLayouts;
    Containers::String setLayoutError;
    if (!TryBuildBuiltinPassSetLayouts(bindingPlan, setLayouts, &setLayoutError)) {
        return failLayout(setLayoutError.CStr());
    }

    if (bindingPlan.perObject.set >= setLayouts.size()) {
        return failLayout("Builtin depth-style pass produced an invalid PerObject descriptor set index");
    }

    passLayout.perObject = bindingPlan.perObject;
    passLayout.firstDescriptorSet = bindingPlan.firstDescriptorSet;
    passLayout.perObjectSetLayout = setLayouts[bindingPlan.perObject.set];
    RefreshBuiltinPassSetLayoutMetadata(passLayout.perObjectSetLayout);

    std::vector<RHI::DescriptorSetLayoutDesc> nativeSetLayouts(setLayouts.size());
    for (size_t setIndex = 0; setIndex < setLayouts.size(); ++setIndex) {
        nativeSetLayouts[setIndex] = setLayouts[setIndex].layout;
    }

    RHI::RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = nativeSetLayouts.empty() ? nullptr : nativeSetLayouts.data();
    pipelineLayoutDesc.setLayoutCount = static_cast<uint32_t>(nativeSetLayouts.size());
    passLayout.pipelineLayout = context.device->CreatePipelineLayout(pipelineLayoutDesc);
    if (passLayout.pipelineLayout == nullptr) {
        return failLayout("Builtin depth-style pass failed to create pipeline layout");
    }

    const auto result = m_passResourceLayouts.emplace(passLayoutKey, std::move(passLayout));
    RefreshBuiltinPassSetLayoutMetadata(result.first->second.perObjectSetLayout);
    return &result.first->second;
}

RHI::RHIPipelineState* BuiltinDepthStylePassBase::GetOrCreatePipelineState(
    const RenderContext& context,
    const RenderSurface& surface,
    const Resources::Material* material) {
    const ResolvedShaderPass resolvedShaderPass = ResolveSurfaceShaderPass(material);
    if (resolvedShaderPass.shader == nullptr || resolvedShaderPass.pass == nullptr) {
        return nullptr;
    }

    PassResourceLayout* passLayout = GetOrCreatePassResourceLayout(context, resolvedShaderPass);
    if (passLayout == nullptr || passLayout->pipelineLayout == nullptr) {
        return nullptr;
    }

    PipelineStateKey pipelineKey = {};
    pipelineKey.renderState =
        material != nullptr ? material->GetRenderState() : Resources::MaterialRenderState();
    pipelineKey.shader = resolvedShaderPass.shader;
    pipelineKey.passName = resolvedShaderPass.passName;
    pipelineKey.renderTargetCount = ResolveSurfaceColorAttachmentCount(surface);
    pipelineKey.renderTargetFormat = static_cast<uint32_t>(ResolveSurfaceColorFormat(surface));
    pipelineKey.depthStencilFormat = static_cast<uint32_t>(ResolveSurfaceDepthFormat(surface));

    const auto existing = m_pipelineStates.find(pipelineKey);
    if (existing != m_pipelineStates.end()) {
        return existing->second;
    }

    const RHI::GraphicsPipelineDesc pipelineDesc = CreatePipelineDesc(
        context.backendType,
        passLayout->pipelineLayout,
        *resolvedShaderPass.shader,
        resolvedShaderPass.passName,
        material,
        surface,
        BuildCommonInputLayout());
    RHI::RHIPipelineState* pipelineState = context.device->CreatePipelineState(pipelineDesc);
    if (pipelineState == nullptr || !pipelineState->IsValid()) {
        if (pipelineState != nullptr) {
            pipelineState->Shutdown();
            delete pipelineState;
        }
        return nullptr;
    }

    m_pipelineStates.emplace(pipelineKey, pipelineState);
    return pipelineState;
}

bool BuiltinDepthStylePassBase::CreateOwnedDescriptorSet(
    const BuiltinPassSetLayoutMetadata& setLayout,
    OwnedDescriptorSet& descriptorSet) {
    RHI::DescriptorPoolDesc poolDesc = {};
    poolDesc.type = setLayout.heapType;
    poolDesc.descriptorCount = CountBuiltinPassHeapDescriptors(setLayout.heapType, setLayout.bindings);
    poolDesc.shaderVisible = setLayout.shaderVisible;

    descriptorSet.pool = m_device->CreateDescriptorPool(poolDesc);
    if (descriptorSet.pool == nullptr) {
        return false;
    }

    descriptorSet.set = descriptorSet.pool->AllocateSet(setLayout.layout);
    if (descriptorSet.set == nullptr) {
        DestroyOwnedDescriptorSet(descriptorSet);
        return false;
    }

    return true;
}

RHI::RHIDescriptorSet* BuiltinDepthStylePassBase::GetOrCreatePerObjectSet(
    const PassLayoutKey& passLayoutKey,
    const PassResourceLayout& passLayout,
    Core::uint64 objectId) {
    if (!passLayout.perObject.IsValid() ||
        passLayout.perObjectSetLayout.layout.bindingCount == 0) {
        return nullptr;
    }

    PerObjectSetKey key = {};
    key.passLayout = passLayoutKey;
    key.objectId = objectId;

    const auto existing = m_perObjectSets.find(key);
    if (existing != m_perObjectSets.end()) {
        return existing->second.set;
    }

    OwnedDescriptorSet descriptorSet = {};
    if (!CreateOwnedDescriptorSet(passLayout.perObjectSetLayout, descriptorSet)) {
        return nullptr;
    }

    const auto result = m_perObjectSets.emplace(key, descriptorSet);
    return result.first->second.set;
}

void BuiltinDepthStylePassBase::DestroyOwnedDescriptorSet(OwnedDescriptorSet& descriptorSet) {
    if (descriptorSet.set != nullptr) {
        descriptorSet.set->Shutdown();
        delete descriptorSet.set;
        descriptorSet.set = nullptr;
    }

    if (descriptorSet.pool != nullptr) {
        descriptorSet.pool->Shutdown();
        delete descriptorSet.pool;
        descriptorSet.pool = nullptr;
    }
}

void BuiltinDepthStylePassBase::DestroyPassResourceLayout(PassResourceLayout& passLayout) {
    if (passLayout.pipelineLayout != nullptr) {
        passLayout.pipelineLayout->Shutdown();
        delete passLayout.pipelineLayout;
        passLayout.pipelineLayout = nullptr;
    }

    passLayout.perObject = {};
    passLayout.perObjectSetLayout = {};
    passLayout.firstDescriptorSet = 0;
}

bool BuiltinDepthStylePassBase::DrawVisibleItem(
    const RenderContext& context,
    const RenderSurface& surface,
    const RenderSceneData& sceneData,
    const VisibleRenderItem& visibleItem) {
    if (visibleItem.mesh == nullptr || visibleItem.gameObject == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinDepthStylePassBase skipped visible item because mesh or game object was null");
        return false;
    }

    const RenderResourceCache::CachedMesh* cachedMesh =
        m_resourceCache.GetOrCreateMesh(m_device, visibleItem.mesh);
    if (cachedMesh == nullptr || cachedMesh->vertexBufferView == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase failed to cache mesh for ") +
                visibleItem.gameObject->GetName().c_str()).CStr());
        return false;
    }

    const Resources::Material* material = ResolveMaterial(visibleItem);
    const ResolvedShaderPass resolvedShaderPass = ResolveSurfaceShaderPass(material);
    if (resolvedShaderPass.shader == nullptr || resolvedShaderPass.pass == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase could not resolve shader pass for ") +
                visibleItem.gameObject->GetName().c_str()).CStr());
        return false;
    }

    PassLayoutKey passLayoutKey = {};
    passLayoutKey.shader = resolvedShaderPass.shader;
    passLayoutKey.passName = resolvedShaderPass.passName;

    PassResourceLayout* passLayout = GetOrCreatePassResourceLayout(context, resolvedShaderPass);
    if (passLayout == nullptr || passLayout->pipelineLayout == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase failed to create pass layout for ") +
                visibleItem.gameObject->GetName().c_str()).CStr());
        return false;
    }

    RHI::RHIPipelineState* pipelineState = GetOrCreatePipelineState(context, surface, material);
    if (pipelineState == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase failed to create pipeline state for ") +
                visibleItem.gameObject->GetName().c_str() +
                " using pass " + resolvedShaderPass.passName).CStr());
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;
    commandList->SetPipelineState(pipelineState);

    RHI::RHIResourceView* vertexBuffers[] = { cachedMesh->vertexBufferView };
    const uint64_t offsets[] = { 0 };
    const uint32_t strides[] = { cachedMesh->vertexStride };
    commandList->SetVertexBuffers(0, 1, vertexBuffers, offsets, strides);
    if (cachedMesh->indexBufferView != nullptr) {
        commandList->SetIndexBuffer(cachedMesh->indexBufferView, 0);
    }

    RHI::RHIDescriptorSet* constantSet = GetOrCreatePerObjectSet(
        passLayoutKey,
        *passLayout,
        visibleItem.gameObject->GetID());
    if (constantSet == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            (Containers::String("BuiltinDepthStylePassBase failed to allocate descriptor set for ") +
                visibleItem.gameObject->GetName().c_str()).CStr());
        return false;
    }

    const Math::Matrix4x4 projectionMatrix =
        m_passType == BuiltinMaterialPass::ShadowCaster
            ? sceneData.cameraData.viewProjection
            : sceneData.cameraData.projection;
    const Math::Matrix4x4 viewMatrix =
        m_passType == BuiltinMaterialPass::ShadowCaster
            ? Math::Matrix4x4::Identity()
            : sceneData.cameraData.view;
    const PerObjectConstants constants = {
        projectionMatrix,
        viewMatrix,
        visibleItem.localToWorld.Transpose()
    };
    constantSet->WriteConstant(passLayout->perObject.binding, &constants, sizeof(constants));

    RHI::RHIDescriptorSet* descriptorSets[] = { constantSet };
    commandList->SetGraphicsDescriptorSets(
        passLayout->firstDescriptorSet,
        1,
        descriptorSets,
        passLayout->pipelineLayout);

    if (visibleItem.hasSection) {
        const Containers::Array<Resources::MeshSection>& sections = visibleItem.mesh->GetSections();
        if (visibleItem.sectionIndex >= sections.Size()) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                (Containers::String("BuiltinDepthStylePassBase received invalid mesh section for ") +
                    visibleItem.gameObject->GetName().c_str()).CStr());
            return false;
        }

        const Resources::MeshSection& section = sections[visibleItem.sectionIndex];
        if (cachedMesh->indexBufferView != nullptr && section.indexCount > 0) {
            commandList->DrawIndexed(section.indexCount, 1, section.startIndex, 0, 0);
        } else if (section.vertexCount > 0) {
            commandList->Draw(section.vertexCount, 1, section.baseVertex, 0);
        }
        return true;
    }

    if (cachedMesh->indexBufferView != nullptr && cachedMesh->indexCount > 0) {
        commandList->DrawIndexed(cachedMesh->indexCount, 1, 0, 0, 0);
    } else if (cachedMesh->vertexCount > 0) {
        commandList->Draw(cachedMesh->vertexCount, 1, 0, 0);
    }

    return true;
}

} // namespace Passes
} // namespace Rendering
} // namespace XCEngine