#include "Rendering/Passes/BuiltinVectorFullscreenPass.h"

#include "Core/Asset/ResourceManager.h"
#include "Debug/Logger.h"
#include <XCEngine/Rendering/RenderPassGraphContract.h>
#include "Rendering/Internal/RenderSurfacePipelineUtils.h"
#include "Rendering/Internal/ShaderVariantUtils.h"
#include "Rendering/RenderSurface.h"
#include "Resources/BuiltinResources.h"
#include "RHI/RHICommandList.h"
#include "RHI/RHIDescriptorPool.h"
#include "RHI/RHIDescriptorSet.h"
#include "RHI/RHIDevice.h"
#include "RHI/RHIPipelineLayout.h"
#include "RHI/RHIPipelineState.h"
#include "RHI/RHIResourceView.h"
#include "RHI/RHISampler.h"

#include <utility>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

struct PostProcessConstants {
    Math::Vector4 vectorPayload = Math::Vector4::One();
};

const Resources::ShaderPass* FindCompatiblePass(
    const Resources::Shader& shader,
    Resources::ShaderBackend backend,
    const Containers::String& preferredPassName) {
    if (!preferredPassName.Empty()) {
        const Resources::ShaderPass* preferredPass =
            shader.FindPass(preferredPassName);
        if (preferredPass != nullptr &&
            ::XCEngine::Rendering::Internal::ShaderPassHasGraphicsVariants(
                shader,
                preferredPass->name,
                backend)) {
            return preferredPass;
        }
        return nullptr;
    }

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

    if (shader.GetPassCount() > 0 &&
        ::XCEngine::Rendering::Internal::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,
    const RenderSurface& surface) {
    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.topologyType = static_cast<uint32_t>(RHI::PrimitiveTopologyType::Triangle);
    ::XCEngine::Rendering::Internal::ApplySingleColorAttachmentPropertiesToGraphicsPipelineDesc(
        surface,
        pipelineDesc);
    pipelineDesc.depthStencilFormat = static_cast<uint32_t>(RHI::Format::Unknown);

    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.srcBlend = static_cast<uint32_t>(RHI::BlendFactor::One);
    pipelineDesc.blendState.dstBlend = static_cast<uint32_t>(RHI::BlendFactor::Zero);
    pipelineDesc.blendState.srcBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::One);
    pipelineDesc.blendState.dstBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::Zero);
    pipelineDesc.blendState.blendOp = static_cast<uint32_t>(RHI::BlendOp::Add);
    pipelineDesc.blendState.blendOpAlpha = static_cast<uint32_t>(RHI::BlendOp::Add);
    pipelineDesc.blendState.colorWriteMask = static_cast<uint8_t>(RHI::ColorWriteMask::All);

    pipelineDesc.depthStencilState.depthTestEnable = false;
    pipelineDesc.depthStencilState.depthWriteEnable = false;
    pipelineDesc.depthStencilState.depthFunc = static_cast<uint32_t>(RHI::ComparisonFunc::Always);

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

    return pipelineDesc;
}

} // namespace

BuiltinVectorFullscreenPass::BuiltinVectorFullscreenPass(
    const Math::Vector4& vectorPayload,
    Containers::String shaderPath,
    Containers::String preferredPassName)
    : m_shaderPath(std::move(shaderPath))
    , m_preferredPassName(std::move(preferredPassName))
    , m_vectorPayload(vectorPayload) {
    if (m_shaderPath.Empty()) {
        m_shaderPath = Resources::GetBuiltinColorScalePostProcessShaderPath();
    }
}

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

const char* BuiltinVectorFullscreenPass::GetName() const {
    return "BuiltinVectorFullscreenPass";
}

bool BuiltinVectorFullscreenPass::SupportsRenderGraph() const {
    return true;
}

bool BuiltinVectorFullscreenPass::RecordRenderGraph(
    const RenderPassRenderGraphContext& context) {
    return RecordSourceColorFullscreenRasterPass(
        *this,
        context);
}

bool BuiltinVectorFullscreenPass::Execute(const RenderPassContext& context) {
    if (!context.renderContext.IsValid() ||
        context.sourceSurface == nullptr ||
        context.sourceColorView == nullptr) {
        return false;
    }
    if (!::XCEngine::Rendering::Internal::HasSingleColorAttachment(*context.sourceSurface)) {
        return false;
    }

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

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

    if (!EnsureInitialized(context.renderContext, context.surface)) {
        return false;
    }

    PostProcessConstants constants = {};
    constants.vectorPayload = m_vectorPayload;
    m_constantsSet.set->WriteConstant(0, &constants, sizeof(constants));

    if (m_boundSourceColorView != context.sourceColorView) {
        m_textureSet.set->Update(0, context.sourceColorView);
        m_boundSourceColorView = context.sourceColorView;
    }

    RHI::RHICommandList* commandList = context.renderContext.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];
    const bool autoTransitionSource = context.sourceSurface->IsAutoTransitionEnabled();

    if (context.surface.IsAutoTransitionEnabled()) {
        commandList->TransitionBarrier(
            renderTarget,
            context.surface.GetColorStateBefore(),
            RHI::ResourceStates::RenderTarget);
    }
    if (autoTransitionSource) {
        commandList->TransitionBarrier(
            context.sourceColorView,
            context.sourceColorState,
            RHI::ResourceStates::PixelShaderResource);
    } else if (context.sourceColorState != RHI::ResourceStates::PixelShaderResource) {
        return false;
    }

    commandList->SetRenderTargets(1, &renderTarget, nullptr);

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

    commandList->SetViewport(viewport);
    commandList->SetScissorRect(scissorRect);
    commandList->SetPrimitiveTopology(RHI::PrimitiveTopology::TriangleList);
    commandList->SetPipelineState(m_pipelineState);

    RHI::RHIDescriptorSet* descriptorSets[] = {
        m_constantsSet.set,
        m_textureSet.set,
        m_samplerSet.set
    };
    commandList->SetGraphicsDescriptorSets(0, 3, descriptorSets, m_pipelineLayout);
    commandList->Draw(3, 1, 0, 0);
    commandList->EndRenderPass();

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

    return true;
}

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

void BuiltinVectorFullscreenPass::SetVectorPayload(const Math::Vector4& vectorPayload) {
    m_vectorPayload = vectorPayload;
}

const Math::Vector4& BuiltinVectorFullscreenPass::GetVectorPayload() const {
    return m_vectorPayload;
}

void BuiltinVectorFullscreenPass::SetShaderPath(const Containers::String& shaderPath) {
    if (m_shaderPath == shaderPath) {
        return;
    }

    DestroyResources();
    m_shaderPath = shaderPath;
}

const Containers::String& BuiltinVectorFullscreenPass::GetShaderPath() const {
    return m_shaderPath;
}

void BuiltinVectorFullscreenPass::SetPreferredPassName(
    const Containers::String& preferredPassName) {
    if (m_preferredPassName == preferredPassName) {
        return;
    }

    DestroyResources();
    m_preferredPassName = preferredPassName;
}

const Containers::String& BuiltinVectorFullscreenPass::GetPreferredPassName() const {
    return m_preferredPassName;
}

bool BuiltinVectorFullscreenPass::EnsureInitialized(
    const RenderContext& renderContext,
    const RenderSurface& surface) {
    const RHI::Format renderTargetFormat =
        ::XCEngine::Rendering::Internal::ResolveSurfaceColorFormat(surface, 0u);
    const uint32_t renderTargetSampleCount =
        ::XCEngine::Rendering::Internal::ResolveSurfaceSampleCount(surface);
    const uint32_t renderTargetSampleQuality =
        ::XCEngine::Rendering::Internal::ResolveSurfaceSampleQuality(surface);
    if (m_device == renderContext.device &&
        m_backendType == renderContext.backendType &&
        m_renderTargetFormat == renderTargetFormat &&
        m_renderTargetSampleCount == renderTargetSampleCount &&
        m_renderTargetSampleQuality == renderTargetSampleQuality &&
        m_pipelineLayout != nullptr &&
        m_pipelineState != nullptr &&
        m_sampler != nullptr &&
        m_constantsSet.set != nullptr &&
        m_textureSet.set != nullptr &&
        m_samplerSet.set != nullptr) {
        return true;
    }

    DestroyResources();
    return CreateResources(renderContext, surface);
}

bool BuiltinVectorFullscreenPass::CreateResources(
    const RenderContext& renderContext,
    const RenderSurface& surface) {
    if (!renderContext.IsValid()) {
        return false;
    }

    RHI::Format renderTargetFormat = RHI::Format::Unknown;
    if (!::XCEngine::Rendering::Internal::TryResolveSingleColorAttachmentFormat(surface, renderTargetFormat)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinVectorFullscreenPass requires exactly one valid destination color attachment");
        return false;
    }

    if (m_shaderPath.Empty()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinVectorFullscreenPass requires a shader path before resource creation");
        return false;
    }

    m_shader = Resources::ResourceManager::Get().Load<Resources::Shader>(m_shaderPath);
    if (!m_shader.IsValid()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinVectorFullscreenPass failed to load configured fullscreen shader resource");
        DestroyResources();
        return false;
    }

    const Resources::ShaderBackend backend = ::XCEngine::Rendering::Internal::ToShaderBackend(renderContext.backendType);
    const Resources::ShaderPass* selectedPass =
        FindCompatiblePass(*m_shader.Get(), backend, m_preferredPassName);
    if (selectedPass == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            m_preferredPassName.Empty()
                ? "BuiltinVectorFullscreenPass could not resolve a compatible fullscreen shader pass"
                : "BuiltinVectorFullscreenPass could not resolve the configured fullscreen shader pass");
        DestroyResources();
        return false;
    }

    m_device = renderContext.device;
    m_backendType = renderContext.backendType;
    m_renderTargetFormat = renderTargetFormat;
    m_renderTargetSampleCount = ::XCEngine::Rendering::Internal::ResolveSurfaceSampleCount(surface);
    m_renderTargetSampleQuality = ::XCEngine::Rendering::Internal::ResolveSurfaceSampleQuality(surface);

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

    RHI::DescriptorSetLayoutBinding textureBinding = {};
    textureBinding.binding = 0;
    textureBinding.type = static_cast<uint32_t>(RHI::DescriptorType::SRV);
    textureBinding.count = 1;
    textureBinding.visibility = static_cast<uint32_t>(RHI::ShaderVisibility::All);

    RHI::DescriptorSetLayoutBinding samplerBinding = {};
    samplerBinding.binding = 0;
    samplerBinding.type = static_cast<uint32_t>(RHI::DescriptorType::Sampler);
    samplerBinding.count = 1;
    samplerBinding.visibility = static_cast<uint32_t>(RHI::ShaderVisibility::All);

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

    RHI::DescriptorSetLayoutDesc textureLayout = {};
    textureLayout.bindings = &textureBinding;
    textureLayout.bindingCount = 1;

    RHI::DescriptorSetLayoutDesc samplerLayout = {};
    samplerLayout.bindings = &samplerBinding;
    samplerLayout.bindingCount = 1;

    RHI::DescriptorSetLayoutDesc setLayouts[] = {
        constantLayout,
        textureLayout,
        samplerLayout
    };

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

    RHI::SamplerDesc samplerDesc = {};
    samplerDesc.filter = static_cast<uint32_t>(RHI::FilterMode::Linear);
    samplerDesc.addressU = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.addressV = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.addressW = static_cast<uint32_t>(RHI::TextureAddressMode::Clamp);
    samplerDesc.mipLodBias = 0.0f;
    samplerDesc.maxAnisotropy = 1;
    samplerDesc.comparisonFunc = static_cast<uint32_t>(RHI::ComparisonFunc::Always);
    samplerDesc.minLod = 0.0f;
    samplerDesc.maxLod = 1000.0f;
    m_sampler = m_device->CreateSampler(samplerDesc);
    if (m_sampler == nullptr) {
        DestroyResources();
        return false;
    }

    auto createOwnedDescriptorSet =
        [this](const RHI::DescriptorSetLayoutDesc& layout,
               RHI::DescriptorHeapType heapType,
               bool shaderVisible,
               OwnedDescriptorSet& ownedSet) -> bool {
            RHI::DescriptorPoolDesc poolDesc = {};
            poolDesc.type = heapType;
            poolDesc.descriptorCount = 1;
            poolDesc.shaderVisible = shaderVisible;
            ownedSet.pool = m_device->CreateDescriptorPool(poolDesc);
            if (ownedSet.pool == nullptr) {
                return false;
            }

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

            return true;
        };

    if (!createOwnedDescriptorSet(
            constantLayout,
            RHI::DescriptorHeapType::CBV_SRV_UAV,
            false,
            m_constantsSet) ||
        !createOwnedDescriptorSet(
            textureLayout,
            RHI::DescriptorHeapType::CBV_SRV_UAV,
            true,
            m_textureSet) ||
        !createOwnedDescriptorSet(
            samplerLayout,
            RHI::DescriptorHeapType::Sampler,
            true,
            m_samplerSet)) {
        DestroyResources();
        return false;
    }

    m_samplerSet.set->UpdateSampler(0, m_sampler);

    m_pipelineState = m_device->CreatePipelineState(
        CreatePipelineDesc(
            m_backendType,
            m_pipelineLayout,
            *m_shader.Get(),
            selectedPass->name,
            surface));
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        DestroyResources();
        return false;
    }

    return true;
}

void BuiltinVectorFullscreenPass::DestroyResources() {
    m_boundSourceColorView = nullptr;

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

    DestroyOwnedDescriptorSet(m_samplerSet);
    DestroyOwnedDescriptorSet(m_textureSet);
    DestroyOwnedDescriptorSet(m_constantsSet);

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

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

    m_shader.Reset();
    m_device = nullptr;
    m_backendType = RHI::RHIType::D3D12;
    m_renderTargetFormat = RHI::Format::Unknown;
    m_renderTargetSampleCount = 1u;
    m_renderTargetSampleQuality = 0u;
}

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

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