#include "Rendering/Passes/BuiltinObjectIdPass.h"

#include "Components/GameObject.h"
#include "Core/Asset/ResourceManager.h"
#include "Debug/Logger.h"
#include "RHI/RHICommandList.h"
#include "RHI/RHIPipelineLayout.h"
#include "RHI/RHIPipelineState.h"
#include "Rendering/Detail/ShaderVariantUtils.h"
#include "Rendering/RenderMaterialUtility.h"
#include "Resources/BuiltinResources.h"
#include "Resources/Mesh/Mesh.h"

#include <cstddef>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

const Resources::ShaderPass* FindObjectIdCompatiblePass(
    const Resources::Shader& shader,
    Resources::ShaderBackend backend) {
    for (const Resources::ShaderPass& shaderPass : shader.GetPasses()) {
        if (ShaderPassMatchesBuiltinPass(shaderPass, BuiltinMaterialPass::ObjectId) &&
            ::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(shader, shaderPass.name, backend)) {
            return &shaderPass;
        }
    }

    const Resources::ShaderPass* objectIdPass = shader.FindPass("ObjectId");
    if (objectIdPass != nullptr &&
        ::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(shader, objectIdPass->name, backend)) {
        return objectIdPass;
    }

    const Resources::ShaderPass* editorObjectIdPass = shader.FindPass("EditorObjectId");
    if (editorObjectIdPass != nullptr &&
        ::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(shader, editorObjectIdPass->name, backend)) {
        return editorObjectIdPass;
    }

    if (shader.GetPassCount() > 0 &&
        ::XCEngine::Rendering::Detail::ShaderPassHasGraphicsVariants(shader, shader.GetPasses()[0].name, backend)) {
        return &shader.GetPasses()[0];
    }

    return nullptr;
}

RHI::GraphicsPipelineDesc CreatePipelineDesc(
    RHI::RHIType backendType,
    RHI::RHIPipelineLayout* pipelineLayout,
    const Resources::Shader& shader,
    const Containers::String& passName) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    pipelineDesc.renderTargetCount = 1;
    pipelineDesc.renderTargetFormats[0] = static_cast<uint32_t>(RHI::Format::R8G8B8A8_UNorm);
    pipelineDesc.depthStencilFormat = static_cast<uint32_t>(RHI::Format::D24_UNorm_S8_UInt);
    pipelineDesc.sampleCount = 1;
    pipelineDesc.inputLayout = BuiltinObjectIdPass::BuildInputLayout();

    pipelineDesc.rasterizerState.fillMode = static_cast<uint32_t>(RHI::FillMode::Solid);
    pipelineDesc.rasterizerState.cullMode = static_cast<uint32_t>(RHI::CullMode::None);
    pipelineDesc.rasterizerState.frontFace = static_cast<uint32_t>(RHI::FrontFace::CounterClockwise);
    pipelineDesc.rasterizerState.depthClipEnable = true;

    pipelineDesc.blendState.blendEnable = false;
    pipelineDesc.blendState.colorWriteMask = static_cast<uint8_t>(RHI::ColorWriteMask::All);

    pipelineDesc.depthStencilState.depthTestEnable = true;
    pipelineDesc.depthStencilState.depthWriteEnable = false;
    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

BuiltinObjectIdPass::~BuiltinObjectIdPass() {
    Shutdown();
}

RHI::InputLayoutDesc BuiltinObjectIdPass::BuildInputLayout() {
    RHI::InputLayoutDesc inputLayout = {};

    RHI::InputElementDesc position = {};
    position.semanticName = "POSITION";
    position.semanticIndex = 0;
    position.format = static_cast<uint32_t>(RHI::Format::R32G32B32_Float);
    position.inputSlot = 0;
    position.alignedByteOffset = static_cast<uint32_t>(offsetof(Resources::StaticMeshVertex, position));
    inputLayout.elements.push_back(position);
    return inputLayout;
}

bool BuiltinObjectIdPass::Render(
    const RenderContext& context,
    const RenderSurface& surface,
    const RenderSceneData& sceneData) {
    if (!context.IsValid()) {
        return false;
    }

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

    const Math::RectInt renderArea = surface.GetRenderArea();
    if (renderArea.width <= 0 || renderArea.height <= 0) {
        return false;
    }

    if (!EnsureInitialized(context)) {
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];

    if (surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            surface.GetColorStateBefore(),
            RHI::ResourceStates::RenderTarget);
    }

    commandList->SetRenderTargets(1, &renderTarget, surface.GetDepthAttachment());

    const RHI::Viewport viewport = {
        static_cast<float>(renderArea.x),
        static_cast<float>(renderArea.y),
        static_cast<float>(renderArea.width),
        static_cast<float>(renderArea.height),
        0.0f,
        1.0f
    };
    const RHI::Rect scissorRect = {
        renderArea.x,
        renderArea.y,
        renderArea.x + renderArea.width,
        renderArea.y + renderArea.height
    };
    const float clearColor[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
    const RHI::Rect clearRects[] = { scissorRect };

    commandList->SetViewport(viewport);
    commandList->SetScissorRect(scissorRect);
    commandList->ClearRenderTarget(renderTarget, clearColor, 1, clearRects);
    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
    commandList->SetPipelineState(m_pipelineState);

    for (const VisibleRenderItem& visibleItem : sceneData.visibleItems) {
        DrawVisibleItem(context, sceneData, visibleItem);
    }

    if (surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            RHI::ResourceStates::RenderTarget,
            surface.GetColorStateAfter());
    }

    return true;
}

void BuiltinObjectIdPass::Shutdown() {
    DestroyResources();
}

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

    if (m_pipelineLayout != nullptr &&
        m_pipelineState != nullptr &&
        m_device == context.device &&
        m_backendType == context.backendType) {
        return true;
    }

    DestroyResources();
    return CreateResources(context);
}

bool BuiltinObjectIdPass::CreateResources(const RenderContext& context) {
    m_device = context.device;
    m_backendType = context.backendType;
    m_builtinObjectIdShader = Resources::ResourceManager::Get().Load<Resources::Shader>(
        Resources::GetBuiltinObjectIdShaderPath());
    if (!m_builtinObjectIdShader.IsValid()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinObjectIdPass failed to load builtin object-id shader resource");
        DestroyResources();
        return false;
    }

    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Detail::ToShaderBackend(m_backendType);
    const Resources::ShaderPass* objectIdPass =
        FindObjectIdCompatiblePass(*m_builtinObjectIdShader.Get(), backend);
    if (objectIdPass == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinObjectIdPass could not resolve a valid ObjectId shader pass");
        DestroyResources();
        return false;
    }

    RHI::DescriptorSetLayoutBinding constantBinding = {};
    constantBinding.binding = 0;
    constantBinding.type = static_cast<uint32_t>(RHI::DescriptorType::CBV);
    constantBinding.count = 1;

    RHI::DescriptorSetLayoutDesc constantLayout = {};
    constantLayout.bindings = &constantBinding;
    constantLayout.bindingCount = 1;

    RHI::RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = &constantLayout;
    pipelineLayoutDesc.setLayoutCount = 1;
    m_pipelineLayout = m_device->CreatePipelineLayout(pipelineLayoutDesc);
    if (m_pipelineLayout == nullptr) {
        DestroyResources();
        return false;
    }

    m_pipelineState = m_device->CreatePipelineState(
        CreatePipelineDesc(
            m_backendType,
            m_pipelineLayout,
            *m_builtinObjectIdShader.Get(),
            objectIdPass->name));
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        if (m_pipelineState != nullptr) {
            m_pipelineState->Shutdown();
            delete m_pipelineState;
            m_pipelineState = nullptr;
        }
        DestroyResources();
        return false;
    }

    return true;
}

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

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

    if (m_pipelineState != nullptr) {
        m_pipelineState->Shutdown();
        delete m_pipelineState;
        m_pipelineState = nullptr;
    }

    if (m_pipelineLayout != nullptr) {
        m_pipelineLayout->Shutdown();
        delete m_pipelineLayout;
        m_pipelineLayout = nullptr;
    }

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

RHI::RHIDescriptorSet* BuiltinObjectIdPass::GetOrCreatePerObjectSet(uint64_t objectId) {
    const auto existing = m_perObjectSets.find(objectId);
    if (existing != m_perObjectSets.end()) {
        return existing->second.set;
    }

    RHI::DescriptorPoolDesc poolDesc = {};
    poolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 1;
    poolDesc.shaderVisible = false;

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

    RHI::DescriptorSetLayoutBinding binding = {};
    binding.binding = 0;
    binding.type = static_cast<uint32_t>(RHI::DescriptorType::CBV);
    binding.count = 1;

    RHI::DescriptorSetLayoutDesc layout = {};
    layout.bindings = &binding;
    layout.bindingCount = 1;
    descriptorSet.set = descriptorSet.pool->AllocateSet(layout);
    if (descriptorSet.set == nullptr) {
        DestroyOwnedDescriptorSet(descriptorSet);
        return nullptr;
    }

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

void BuiltinObjectIdPass::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;
    }
}

bool BuiltinObjectIdPass::DrawVisibleItem(
    const RenderContext& context,
    const RenderSceneData& sceneData,
    const VisibleRenderItem& visibleItem) {
    if (visibleItem.mesh == nullptr || visibleItem.gameObject == nullptr) {
        return false;
    }

    const RenderResourceCache::CachedMesh* cachedMesh =
        m_resourceCache.GetOrCreateMesh(m_device, visibleItem.mesh);
    if (cachedMesh == nullptr || cachedMesh->vertexBufferView == nullptr) {
        return false;
    }

    RHI::RHICommandList* commandList = context.commandList;

    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);
    }

    const uint64_t objectId = visibleItem.gameObject->GetID();
    RHI::RHIDescriptorSet* constantSet = GetOrCreatePerObjectSet(objectId);
    if (constantSet == nullptr) {
        return false;
    }

    const PerObjectConstants constants = {
        sceneData.cameraData.projection,
        sceneData.cameraData.view,
        visibleItem.localToWorld.Transpose(),
        EncodeObjectIdToColor(objectId)
    };
    constantSet->WriteConstant(0, &constants, sizeof(constants));

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

    if (visibleItem.hasSection) {
        const Containers::Array<Resources::MeshSection>& sections = visibleItem.mesh->GetSections();
        if (visibleItem.sectionIndex >= sections.Size()) {
            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