XCEngine/engine/Runtime/Rendering/Caches/RenderResourceCache.h

#pragma once

#include <XCEngine/RHI/RHIBuffer.h>
#include <XCEngine/RHI/RHIDevice.h>
#include <XCEngine/RHI/RHIResourceView.h>
#include <XCEngine/RHI/RHITexture.h>
#include <XCEngine/Resources/GaussianSplat/GaussianSplat.h>
#include <XCEngine/Resources/Mesh/Mesh.h>
#include <XCEngine/Resources/Texture/Texture.h>
#include <XCEngine/Resources/Volume/VolumeField.h>

#include <unordered_map>

namespace XCEngine {
namespace Rendering {

class RenderResourceCache {
public:
    enum class GaussianSplatResidencyState {
        Uninitialized = 0,
        CpuReady,
        GpuUploading,
        GpuReady,
        Failed
    };

    struct CachedMesh {
        RHI::RHIBuffer* vertexBuffer = nullptr;
        RHI::RHIResourceView* vertexBufferView = nullptr;
        RHI::RHIBuffer* indexBuffer = nullptr;
        RHI::RHIResourceView* indexBufferView = nullptr;
        uint32_t vertexCount = 0;
        uint32_t indexCount = 0;
        uint32_t vertexStride = 0;
        bool uses32BitIndices = false;
    };

    struct CachedTexture {
        RHI::RHITexture* texture = nullptr;
        RHI::RHIResourceView* shaderResourceView = nullptr;
        uint32_t width = 0;
        uint32_t height = 0;
    };

    struct CachedBufferView {
        RHI::RHIResourceView* resourceView = nullptr;
    };

    struct CachedVolumeField {
        RHI::RHIBuffer* payloadBuffer = nullptr;
        RHI::RHIResourceView* shaderResourceView = nullptr;
        uint32_t elementStride = 0;
        uint32_t elementCount = 0;
        uint64_t payloadSize = 0;
        Resources::VolumeStorageKind storageKind = Resources::VolumeStorageKind::Unknown;
    };

    struct CachedGaussianSplatSection {
        RHI::RHIBuffer* buffer = nullptr;
        RHI::RHIResourceView* shaderResourceView = nullptr;
        Resources::GaussianSplatSectionType type = Resources::GaussianSplatSectionType::Unknown;
        Resources::GaussianSplatSectionFormat format = Resources::GaussianSplatSectionFormat::Unknown;
        uint32_t elementStride = 0;
        uint32_t elementCount = 0;
        uint64_t payloadSize = 0;
    };

    struct CachedGaussianSplat {
        GaussianSplatResidencyState residencyState = GaussianSplatResidencyState::Uninitialized;
        uint32_t contentVersion = 0;
        uint32_t splatCount = 0;
        uint32_t chunkCount = 0;
        uint32_t cameraCount = 0;
        Math::Bounds bounds;
        CachedGaussianSplatSection positions;
        CachedGaussianSplatSection other;
        CachedGaussianSplatSection color;
        CachedGaussianSplatSection sh;
        CachedGaussianSplatSection chunks;
    };

    ~RenderResourceCache();

    void Shutdown();

    const CachedMesh* GetOrCreateMesh(RHI::RHIDevice* device, const Resources::Mesh* mesh);
    const CachedTexture* GetOrCreateTexture(RHI::RHIDevice* device, const Resources::Texture* texture);
    const CachedVolumeField* GetOrCreateVolumeField(
        RHI::RHIDevice* device,
        const Resources::VolumeField* volumeField);
    const CachedGaussianSplat* GetOrCreateGaussianSplat(
        RHI::RHIDevice* device,
        const Resources::GaussianSplat* gaussianSplat);
    const CachedBufferView* GetOrCreateBufferView(
        RHI::RHIDevice* device,
        RHI::RHIBuffer* buffer,
        RHI::ResourceViewType viewType,
        const RHI::ResourceViewDesc& viewDesc);

private:
    struct BufferViewCacheKey {
        const RHI::RHIBuffer* buffer = nullptr;
        RHI::ResourceViewType viewType = RHI::ResourceViewType::ShaderResource;
        uint32_t format = 0;
        RHI::ResourceViewDimension dimension = RHI::ResourceViewDimension::Unknown;
        uint64_t bufferLocation = 0;
        uint32_t firstElement = 0;
        uint32_t elementCount = 0;
        uint32_t structureByteStride = 0;

        bool operator==(const BufferViewCacheKey& other) const {
            return buffer == other.buffer &&
                   viewType == other.viewType &&
                   format == other.format &&
                   dimension == other.dimension &&
                   bufferLocation == other.bufferLocation &&
                   firstElement == other.firstElement &&
                   elementCount == other.elementCount &&
                   structureByteStride == other.structureByteStride;
        }
    };

    struct BufferViewCacheKeyHash {
        size_t operator()(const BufferViewCacheKey& key) const noexcept;
    };

    bool UploadMesh(RHI::RHIDevice* device, const Resources::Mesh* mesh, CachedMesh& cachedMesh);
    bool UploadTexture(RHI::RHIDevice* device, const Resources::Texture* texture, CachedTexture& cachedTexture);
    bool UploadVolumeField(
        RHI::RHIDevice* device,
        const Resources::VolumeField* volumeField,
        CachedVolumeField& cachedVolumeField);
    bool UploadGaussianSplat(
        RHI::RHIDevice* device,
        const Resources::GaussianSplat* gaussianSplat,
        CachedGaussianSplat& cachedGaussianSplat);
    bool UploadGaussianSplatSection(
        RHI::RHIDevice* device,
        const Resources::GaussianSplat* gaussianSplat,
        Resources::GaussianSplatSectionType sectionType,
        Resources::GaussianSplatSectionFormat requiredFormat,
        uint32_t requiredStride,
        bool requiredSection,
        CachedGaussianSplatSection& cachedSection);
    bool CreateBufferView(
        RHI::RHIDevice* device,
        RHI::RHIBuffer* buffer,
        RHI::ResourceViewType viewType,
        const RHI::ResourceViewDesc& viewDesc,
        CachedBufferView& cachedBufferView);

    std::unordered_map<const Resources::Mesh*, CachedMesh> m_meshCache;
    std::unordered_map<const Resources::Texture*, CachedTexture> m_textureCache;
    std::unordered_map<const Resources::VolumeField*, CachedVolumeField> m_volumeFieldCache;
    std::unordered_map<const Resources::GaussianSplat*, CachedGaussianSplat> m_gaussianSplatCache;
    std::unordered_map<BufferViewCacheKey, CachedBufferView, BufferViewCacheKeyHash> m_bufferViewCache;
};

} // namespace Rendering
} // namespace XCEngine