#include "Rendering/Passes/BuiltinObjectIdOutlinePass.h"

#include "Rendering/ObjectIdEncoding.h"
#include "RHI/RHICommandList.h"
#include "RHI/RHIDevice.h"

#include <algorithm>
#include <cstring>

namespace XCEngine {
namespace Rendering {
namespace Passes {

namespace {

const char kBuiltinObjectIdOutlineHlsl[] = R"(
cbuffer OutlineConstants : register(b0) {
    float4 gViewportSizeAndTexelSize;
    float4 gOutlineColor;
    float4 gSelectedInfo;
    float4 gSelectedObjectColors[256];
};

Texture2D gObjectIdTexture : register(t0);

struct VSOutput {
    float4 position : SV_POSITION;
};

VSOutput MainVS(uint vertexId : SV_VertexID) {
    static const float2 positions[3] = {
        float2(-1.0, -1.0),
        float2(-1.0,  3.0),
        float2( 3.0, -1.0)
    };

    VSOutput output;
    output.position = float4(positions[vertexId], 0.0, 1.0);
    return output;
}

int2 ClampPixelCoord(int2 pixelCoord) {
    const int2 maxCoord = int2(
        max((int)gViewportSizeAndTexelSize.x - 1, 0),
        max((int)gViewportSizeAndTexelSize.y - 1, 0));
    return clamp(pixelCoord, int2(0, 0), maxCoord);
}

float4 LoadObjectId(int2 pixelCoord) {
    return gObjectIdTexture.Load(int3(ClampPixelCoord(pixelCoord), 0));
}

bool IsSelectedObject(float4 objectIdColor) {
    if (objectIdColor.a <= 0.0) {
        return false;
    }

    const int selectedCount = min((int)gSelectedInfo.x, 256);
    [loop]
    for (int i = 0; i < selectedCount; ++i) {
        const float4 selectedColor = gSelectedObjectColors[i];
        if (all(abs(objectIdColor - selectedColor) <= float4(
                0.0025,
                0.0025,
                0.0025,
                0.0025))) {
            return true;
        }
    }

    return false;
}

float4 MainPS(VSOutput input) : SV_TARGET {
    const int2 pixelCoord = int2(input.position.xy);
    const bool debugSelectionMask = gSelectedInfo.y > 0.5;
    const bool centerSelected = IsSelectedObject(LoadObjectId(pixelCoord));

    if (debugSelectionMask) {
        return centerSelected ? float4(1.0, 1.0, 1.0, 1.0) : float4(0.0, 0.0, 0.0, 1.0);
    }

    if (centerSelected) {
        discard;
    }

    const int outlineWidth = max((int)gSelectedInfo.z, 1);
    float outline = 0.0;
    [loop]
    for (int y = -2; y <= 2; ++y) {
        [loop]
        for (int x = -2; x <= 2; ++x) {
            if (x == 0 && y == 0) {
                continue;
            }

            const float distancePixels = length(float2((float)x, (float)y));
            if (distancePixels > outlineWidth) {
                continue;
            }

            if (!IsSelectedObject(LoadObjectId(pixelCoord + int2(x, y)))) {
                continue;
            }

            const float weight = saturate(1.0 - ((distancePixels - 1.0) / max((float)outlineWidth, 1.0)));
            outline = max(outline, weight);
        }
    }

    if (outline <= 0.001) {
        discard;
    }

    return float4(gOutlineColor.rgb, gOutlineColor.a * outline);
}
)";

} // namespace

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

bool BuiltinObjectIdOutlinePass::Render(
    const RenderContext& renderContext,
    const RenderSurface& surface,
    RHI::RHIResourceView* objectIdTextureView,
    const std::vector<uint64_t>& selectedObjectIds,
    const ObjectIdOutlineStyle& style) {
    if (!renderContext.IsValid() ||
        renderContext.backendType != RHI::RHIType::D3D12 ||
        objectIdTextureView == nullptr ||
        selectedObjectIds.empty()) {
        return false;
    }

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

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

    OutlineConstants constants = {};
    constants.viewportSizeAndTexelSize = Math::Vector4(
        static_cast<float>(surface.GetWidth()),
        static_cast<float>(surface.GetHeight()),
        surface.GetWidth() > 0 ? 1.0f / static_cast<float>(surface.GetWidth()) : 0.0f,
        surface.GetHeight() > 0 ? 1.0f / static_cast<float>(surface.GetHeight()) : 0.0f);
    constants.outlineColor = Math::Vector4(
        style.outlineColor.r,
        style.outlineColor.g,
        style.outlineColor.b,
        style.outlineColor.a);

    const uint32_t selectedCount = (std::min)(
        static_cast<uint32_t>(selectedObjectIds.size()),
        kMaxSelectedObjectCount);
    constants.selectedInfo = Math::Vector4(
        static_cast<float>(selectedCount),
        style.debugSelectionMask ? 1.0f : 0.0f,
        style.outlineWidthPixels,
        0.0f);
    for (uint32_t index = 0; index < selectedCount; ++index) {
        constants.selectedObjectColors[index] =
            EncodeObjectIdToColor(selectedObjectIds[index]);
    }

    m_constantSet->WriteConstant(0, &constants, sizeof(constants));
    m_textureSet->Update(0, objectIdTextureView);

    RHI::RHICommandList* commandList = renderContext.commandList;
    RHI::RHIResourceView* renderTarget = colorAttachments[0];
    commandList->SetRenderTargets(1, &renderTarget, nullptr);

    const RHI::Viewport viewport = {
        0.0f,
        0.0f,
        static_cast<float>(surface.GetWidth()),
        static_cast<float>(surface.GetHeight()),
        0.0f,
        1.0f
    };
    const RHI::Rect scissorRect = {
        0,
        0,
        static_cast<int32_t>(surface.GetWidth()),
        static_cast<int32_t>(surface.GetHeight())
    };

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

    RHI::RHIDescriptorSet* descriptorSets[] = { m_constantSet, m_textureSet };
    commandList->SetGraphicsDescriptorSets(0, 2, descriptorSets, m_pipelineLayout);
    commandList->Draw(3, 1, 0, 0);
    return true;
}

bool BuiltinObjectIdOutlinePass::EnsureInitialized(const RenderContext& renderContext) {
    if (m_pipelineLayout != nullptr &&
        m_pipelineState != nullptr &&
        m_constantPool != nullptr &&
        m_constantSet != nullptr &&
        m_texturePool != nullptr &&
        m_textureSet != nullptr &&
        m_device == renderContext.device &&
        m_backendType == renderContext.backendType) {
        return true;
    }

    DestroyResources();
    return CreateResources(renderContext);
}

bool BuiltinObjectIdOutlinePass::CreateResources(const RenderContext& renderContext) {
    if (!renderContext.IsValid() || renderContext.backendType != RHI::RHIType::D3D12) {
        return false;
    }

    m_device = renderContext.device;
    m_backendType = renderContext.backendType;

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

    RHI::DescriptorSetLayoutDesc constantLayout = {};
    constantLayout.bindings = &setBindings[0];
    constantLayout.bindingCount = 1;

    RHI::DescriptorSetLayoutDesc textureLayout = {};
    textureLayout.bindings = &setBindings[1];
    textureLayout.bindingCount = 1;

    RHI::DescriptorSetLayoutDesc setLayouts[2] = {};
    setLayouts[0] = constantLayout;
    setLayouts[1] = textureLayout;

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

    RHI::DescriptorPoolDesc constantPoolDesc = {};
    constantPoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
    constantPoolDesc.descriptorCount = 1;
    constantPoolDesc.shaderVisible = false;
    m_constantPool = m_device->CreateDescriptorPool(constantPoolDesc);
    if (m_constantPool == nullptr) {
        DestroyResources();
        return false;
    }

    m_constantSet = m_constantPool->AllocateSet(constantLayout);
    if (m_constantSet == nullptr) {
        DestroyResources();
        return false;
    }

    RHI::DescriptorPoolDesc texturePoolDesc = {};
    texturePoolDesc.type = RHI::DescriptorHeapType::CBV_SRV_UAV;
    texturePoolDesc.descriptorCount = 1;
    texturePoolDesc.shaderVisible = true;
    m_texturePool = m_device->CreateDescriptorPool(texturePoolDesc);
    if (m_texturePool == nullptr) {
        DestroyResources();
        return false;
    }

    m_textureSet = m_texturePool->AllocateSet(textureLayout);
    if (m_textureSet == nullptr) {
        DestroyResources();
        return false;
    }

    RHI::GraphicsPipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = m_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::Unknown);
    pipelineDesc.sampleCount = 1;

    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 = true;
    pipelineDesc.blendState.srcBlend = static_cast<uint32_t>(RHI::BlendFactor::SrcAlpha);
    pipelineDesc.blendState.dstBlend = static_cast<uint32_t>(RHI::BlendFactor::InvSrcAlpha);
    pipelineDesc.blendState.srcBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::One);
    pipelineDesc.blendState.dstBlendAlpha = static_cast<uint32_t>(RHI::BlendFactor::InvSrcAlpha);
    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);

    pipelineDesc.vertexShader.source.assign(
        kBuiltinObjectIdOutlineHlsl,
        kBuiltinObjectIdOutlineHlsl + std::strlen(kBuiltinObjectIdOutlineHlsl));
    pipelineDesc.vertexShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
    pipelineDesc.vertexShader.entryPoint = L"MainVS";
    pipelineDesc.vertexShader.profile = L"vs_5_0";

    pipelineDesc.fragmentShader.source.assign(
        kBuiltinObjectIdOutlineHlsl,
        kBuiltinObjectIdOutlineHlsl + std::strlen(kBuiltinObjectIdOutlineHlsl));
    pipelineDesc.fragmentShader.sourceLanguage = RHI::ShaderLanguage::HLSL;
    pipelineDesc.fragmentShader.entryPoint = L"MainPS";
    pipelineDesc.fragmentShader.profile = L"ps_5_0";

    m_pipelineState = m_device->CreatePipelineState(pipelineDesc);
    if (m_pipelineState == nullptr || !m_pipelineState->IsValid()) {
        DestroyResources();
        return false;
    }

    return true;
}

void BuiltinObjectIdOutlinePass::DestroyResources() {
    if (m_pipelineState != nullptr) {
        m_pipelineState->Shutdown();
        delete m_pipelineState;
        m_pipelineState = nullptr;
    }

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

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

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

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

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

    m_device = nullptr;
    m_backendType = RHI::RHIType::D3D12;
}

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