#include "Rendering/Passes/Internal/BuiltinGaussianSplatPassResources.h"

#include "Components/GaussianSplatRendererComponent.h"
#include "Debug/Logger.h"
#include "Resources/GaussianSplat/GaussianSplat.h"

namespace XCEngine {
namespace Rendering {
namespace Passes {
namespace Internal {

namespace {

bool CreateStructuredBufferViews(
    RHI::RHIDevice* device,
    Core::uint32 elementCount,
    Core::uint32 elementStride,
    BuiltinGaussianSplatPassResources::CachedBufferView& bufferView) {
    if (device == nullptr || elementCount == 0u || elementStride == 0u) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::CreateStructuredBufferViews failed: invalid parameters");
        return false;
    }

    RHI::BufferDesc bufferDesc = {};
    bufferDesc.size = static_cast<Core::uint64>(elementCount) * static_cast<Core::uint64>(elementStride);
    bufferDesc.stride = elementStride;
    bufferDesc.bufferType = static_cast<uint32_t>(RHI::BufferType::Storage);
    bufferDesc.flags = static_cast<Core::uint64>(RHI::BufferFlags::AllowUnorderedAccess);

    bufferView.buffer = device->CreateBuffer(bufferDesc);
    if (bufferView.buffer == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::CreateStructuredBufferViews failed: CreateBuffer returned null");
        return false;
    }

    bufferView.buffer->SetStride(elementStride);
    bufferView.buffer->SetBufferType(RHI::BufferType::Storage);
    bufferView.buffer->SetState(RHI::ResourceStates::Common);

    RHI::ResourceViewDesc viewDesc = {};
    viewDesc.dimension = RHI::ResourceViewDimension::StructuredBuffer;
    viewDesc.firstElement = 0u;
    viewDesc.elementCount = elementCount;
    viewDesc.structureByteStride = elementStride;

    bufferView.shaderResourceView = device->CreateShaderResourceView(bufferView.buffer, viewDesc);
    if (bufferView.shaderResourceView == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::CreateStructuredBufferViews failed: CreateShaderResourceView returned null");
        return false;
    }

    bufferView.unorderedAccessView = device->CreateUnorderedAccessView(bufferView.buffer, viewDesc);
    if (bufferView.unorderedAccessView == nullptr) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::CreateStructuredBufferViews failed: CreateUnorderedAccessView returned null");
        return false;
    }

    bufferView.elementCount = elementCount;
    bufferView.elementStride = elementStride;
    bufferView.currentState = RHI::ResourceStates::Common;
    return true;
}

} // namespace

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

bool BuiltinGaussianSplatPassResources::EnsureWorkingSet(
    RHI::RHIDevice* device,
    const VisibleGaussianSplatItem& visibleGaussianSplat,
    WorkingSet*& outWorkingSet) {
    outWorkingSet = nullptr;
    if (device == nullptr ||
        visibleGaussianSplat.gaussianSplatRenderer == nullptr ||
        visibleGaussianSplat.gaussianSplat == nullptr ||
        !visibleGaussianSplat.gaussianSplat->IsValid()) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::EnsureWorkingSet failed: invalid input");
        return false;
    }

    const Core::uint32 splatCapacity = visibleGaussianSplat.gaussianSplat->GetSplatCount();
    const Core::uint32 sortCapacity = ComputeSortCapacity(splatCapacity);
    if (splatCapacity == 0u) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::EnsureWorkingSet failed: splat capacity is zero");
        return false;
    }

    if (!ResetForDevice(device)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::EnsureWorkingSet failed: ResetForDevice returned false");
        return false;
    }

    WorkingSet& workingSet = m_workingSets[visibleGaussianSplat.gaussianSplatRenderer];
    if (workingSet.splatCapacity < splatCapacity) {
        DestroyBufferView(workingSet.sortDistances);
        DestroyBufferView(workingSet.orderIndices);
        DestroyBufferView(workingSet.viewData);
        if (!RecreateWorkingSet(device, splatCapacity, sortCapacity, workingSet)) {
            Debug::Logger::Get().Error(
                Debug::LogCategory::Rendering,
                "BuiltinGaussianSplatPassResources::EnsureWorkingSet failed: RecreateWorkingSet returned false");
            m_workingSets.erase(visibleGaussianSplat.gaussianSplatRenderer);
            return false;
        }
    }

    workingSet.renderer = visibleGaussianSplat.gaussianSplatRenderer;
    outWorkingSet = &workingSet;
    return true;
}

bool BuiltinGaussianSplatPassResources::EnsureAccumulationSurface(
    RHI::RHIDevice* device,
    Core::uint32 width,
    Core::uint32 height,
    RHI::Format format,
    AccumulationSurface*& outSurface) {
    outSurface = nullptr;
    if (device == nullptr ||
        width == 0u ||
        height == 0u ||
        format == RHI::Format::Unknown) {
        return false;
    }

    if (!ResetForDevice(device)) {
        return false;
    }

    if (m_accumulationSurface.texture == nullptr ||
        m_accumulationSurface.width != width ||
        m_accumulationSurface.height != height ||
        m_accumulationSurface.format != format) {
        if (!RecreateAccumulationSurface(device, width, height, format)) {
            return false;
        }
    }

    outSurface = &m_accumulationSurface;
    return true;
}

void BuiltinGaussianSplatPassResources::Shutdown() {
    for (auto& workingSetPair : m_workingSets) {
        DestroyBufferView(workingSetPair.second.sortDistances);
        DestroyBufferView(workingSetPair.second.orderIndices);
        DestroyBufferView(workingSetPair.second.viewData);
    }

    m_workingSets.clear();
    DestroyAccumulationSurface(m_accumulationSurface);
    m_device = nullptr;
}

const BuiltinGaussianSplatPassResources::WorkingSet* BuiltinGaussianSplatPassResources::FindWorkingSet(
    const Components::GaussianSplatRendererComponent* renderer) const {
    const auto workingSetIt = m_workingSets.find(renderer);
    return workingSetIt != m_workingSets.end() ? &workingSetIt->second : nullptr;
}

const BuiltinGaussianSplatPassResources::AccumulationSurface* BuiltinGaussianSplatPassResources::GetAccumulationSurface() const {
    return m_accumulationSurface.texture != nullptr ? &m_accumulationSurface : nullptr;
}

void BuiltinGaussianSplatPassResources::DestroyBufferView(CachedBufferView& bufferView) {
    if (bufferView.shaderResourceView != nullptr) {
        bufferView.shaderResourceView->Shutdown();
        delete bufferView.shaderResourceView;
        bufferView.shaderResourceView = nullptr;
    }

    if (bufferView.unorderedAccessView != nullptr) {
        bufferView.unorderedAccessView->Shutdown();
        delete bufferView.unorderedAccessView;
        bufferView.unorderedAccessView = nullptr;
    }

    if (bufferView.buffer != nullptr) {
        bufferView.buffer->Shutdown();
        delete bufferView.buffer;
        bufferView.buffer = nullptr;
    }

    bufferView.elementCount = 0u;
    bufferView.elementStride = 0u;
    bufferView.currentState = RHI::ResourceStates::Common;
}

void BuiltinGaussianSplatPassResources::DestroyAccumulationSurface(AccumulationSurface& surface) {
    if (surface.renderTargetView != nullptr) {
        surface.renderTargetView->Shutdown();
        delete surface.renderTargetView;
        surface.renderTargetView = nullptr;
    }

    if (surface.shaderResourceView != nullptr) {
        surface.shaderResourceView->Shutdown();
        delete surface.shaderResourceView;
        surface.shaderResourceView = nullptr;
    }

    if (surface.texture != nullptr) {
        surface.texture->Shutdown();
        delete surface.texture;
        surface.texture = nullptr;
    }

    surface.width = 0u;
    surface.height = 0u;
    surface.format = RHI::Format::Unknown;
    surface.currentColorState = RHI::ResourceStates::Common;
}

bool BuiltinGaussianSplatPassResources::ResetForDevice(RHI::RHIDevice* device) {
    if (m_device == nullptr) {
        m_device = device;
        return true;
    }

    if (m_device == device) {
        return true;
    }

    Shutdown();
    m_device = device;
    return true;
}

bool BuiltinGaussianSplatPassResources::RecreateWorkingSet(
    RHI::RHIDevice* device,
    Core::uint32 splatCapacity,
    Core::uint32 sortCapacity,
    WorkingSet& workingSet) {
    if (!CreateStructuredBufferView(device, sortCapacity, kSortDistanceStride, workingSet.sortDistances)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::RecreateWorkingSet failed: sort-distance buffer view creation failed");
        DestroyBufferView(workingSet.sortDistances);
        DestroyBufferView(workingSet.orderIndices);
        DestroyBufferView(workingSet.viewData);
        workingSet.renderer = nullptr;
        workingSet.splatCapacity = 0u;
        workingSet.sortCapacity = 0u;
        return false;
    }

    if (!CreateStructuredBufferView(device, sortCapacity, kOrderIndexStride, workingSet.orderIndices)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::RecreateWorkingSet failed: order-index buffer view creation failed");
        DestroyBufferView(workingSet.sortDistances);
        DestroyBufferView(workingSet.orderIndices);
        DestroyBufferView(workingSet.viewData);
        workingSet.renderer = nullptr;
        workingSet.splatCapacity = 0u;
        workingSet.sortCapacity = 0u;
        return false;
    }

    if (!CreateStructuredBufferView(device, splatCapacity, kViewDataStride, workingSet.viewData)) {
        Debug::Logger::Get().Error(
            Debug::LogCategory::Rendering,
            "BuiltinGaussianSplatPassResources::RecreateWorkingSet failed: view-data buffer view creation failed");
        DestroyBufferView(workingSet.sortDistances);
        DestroyBufferView(workingSet.orderIndices);
        DestroyBufferView(workingSet.viewData);
        workingSet.renderer = nullptr;
        workingSet.splatCapacity = 0u;
        workingSet.sortCapacity = 0u;
        return false;
    }

    workingSet.splatCapacity = splatCapacity;
    workingSet.sortCapacity = sortCapacity;
    return true;
}

bool BuiltinGaussianSplatPassResources::CreateStructuredBufferView(
    RHI::RHIDevice* device,
    Core::uint32 elementCount,
    Core::uint32 elementStride,
    CachedBufferView& bufferView) {
    return CreateStructuredBufferViews(device, elementCount, elementStride, bufferView);
}

bool BuiltinGaussianSplatPassResources::RecreateAccumulationSurface(
    RHI::RHIDevice* device,
    Core::uint32 width,
    Core::uint32 height,
    RHI::Format format) {
    DestroyAccumulationSurface(m_accumulationSurface);

    RHI::TextureDesc textureDesc = {};
    textureDesc.width = width;
    textureDesc.height = height;
    textureDesc.depth = 1u;
    textureDesc.mipLevels = 1u;
    textureDesc.arraySize = 1u;
    textureDesc.format = static_cast<uint32_t>(format);
    textureDesc.textureType = static_cast<uint32_t>(RHI::TextureType::Texture2D);
    textureDesc.sampleCount = 1u;
    textureDesc.sampleQuality = 0u;
    textureDesc.flags = 0u;

    m_accumulationSurface.texture = device->CreateTexture(textureDesc);
    if (m_accumulationSurface.texture == nullptr) {
        return false;
    }

    RHI::ResourceViewDesc viewDesc = {};
    viewDesc.format = static_cast<uint32_t>(format);
    viewDesc.dimension = RHI::ResourceViewDimension::Texture2D;
    viewDesc.mipLevel = 0u;

    m_accumulationSurface.renderTargetView = device->CreateRenderTargetView(m_accumulationSurface.texture, viewDesc);
    if (m_accumulationSurface.renderTargetView == nullptr) {
        DestroyAccumulationSurface(m_accumulationSurface);
        return false;
    }

    m_accumulationSurface.shaderResourceView = device->CreateShaderResourceView(m_accumulationSurface.texture, viewDesc);
    if (m_accumulationSurface.shaderResourceView == nullptr) {
        DestroyAccumulationSurface(m_accumulationSurface);
        return false;
    }

    m_accumulationSurface.width = width;
    m_accumulationSurface.height = height;
    m_accumulationSurface.format = format;
    m_accumulationSurface.currentColorState = RHI::ResourceStates::Common;
    return true;
}

Core::uint32 BuiltinGaussianSplatPassResources::ComputeSortCapacity(Core::uint32 splatCapacity) {
    if (splatCapacity == 0u) {
        return 0u;
    }

    Core::uint32 sortCapacity = 1u;
    while (sortCapacity < splatCapacity && sortCapacity <= (0x80000000u >> 1u)) {
        sortCapacity <<= 1u;
    }
    return sortCapacity < splatCapacity ? splatCapacity : sortCapacity;
}

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