// =============================================
//  Aster: render_resource_manager.h
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#pragma once

#include "global.h"

#include <EASTL/deque.h>
#include <EASTL/vector_map.h>

struct Device;
struct Texture;
struct StorageTexture;
struct StorageBuffer;
struct Buffer;

struct GpuResourceHandle
{
    constexpr static u32 INVALID_HANDLE = MaxValue<u32>;
    u32 m_Index = INVALID_HANDLE; // Default = invalid

    [[nodiscard]] bool
    IsInvalid() const
    {
        return m_Index == INVALID_HANDLE;
    }
};

struct BufferHandle : GpuResourceHandle
{
};

struct TextureHandle : GpuResourceHandle
{
};

struct StorageTextureHandle : GpuResourceHandle
{
};

struct SamplerHandle : GpuResourceHandle
{
};

struct VirtualizedBufferHandle : GpuResourceHandle
{
};

struct GeometryHandle : VirtualizedBufferHandle
{
};

struct IndexHandle : VirtualizedBufferHandle
{
};

struct MaterialHandle : VirtualizedBufferHandle
{
};

struct TextureManager
{
    eastl::vector<Texture> m_Textures;
    u32 m_MaxCapacity;
    u32 m_FreeHead;

    void Init(u32 maxCapacity);
    TextureHandle Commit(Texture *texture);
    Texture *Fetch(TextureHandle handle);
    void Release(const Device *device, TextureHandle handle);
    void Destroy(const Device *device);
};

struct BufferManager
{
    eastl::vector<StorageBuffer> m_Buffers;
    u32 m_MaxCapacity;
    u32 m_FreeHead;

    void Init(u32 maxCapacity);
    BufferHandle Commit(StorageBuffer *buffer);
    StorageBuffer *Fetch(BufferHandle handle);
    void Release(const Device *device, BufferHandle handle);
    void Destroy(const Device *device);
};

struct StorageTextureManager : TextureManager
{
    StorageTextureHandle Commit(StorageTexture *texture);
    StorageTexture *Fetch(StorageTextureHandle handle);
    void Release(const Device *device, StorageTextureHandle handle);
};

struct SamplerManager
{
    // There can only be so many samplers.
    eastl::vector<vk::Sampler> m_Samplers;
    eastl::vector<usize> m_SamplerHashes;

    void Init(usize size);
    SamplerHandle Create(const Device *device, const vk::SamplerCreateInfo *createInfo);
    vk::Sampler Fetch(SamplerHandle handle);
    void Destroy(const Device *device);
};

struct VirtualizedBufferPool
{
    // TODO: Use buffer device address
    std::unique_ptr<StorageBuffer> m_BackingBuffer;
    uptr m_BufferPtr;
    VmaVirtualBlock m_Block;

    struct VirtualBuffer
    {
        VmaVirtualAllocation m_Allocation;
        usize m_Offset;
        usize m_Size;
    };

    eastl::vector<VirtualBuffer> m_VirtualBuffers;
    u32 m_FreeHead = GpuResourceHandle::INVALID_HANDLE;
    bool m_Dirty = false;

    void InitStorage(const Device *device, usize bufferMaxSize);
    void InitIndex(const Device *device, usize bufferMaxSize);

    // Sync the offset buffer if required.
    // FUTURE(Bob): Handle the writes for non-ReBAR system.
    void UpdateToGpu(const Device *device);

    VirtualizedBufferHandle Create(usize size, usize alignment);
    usize FetchOffset(VirtualizedBufferHandle handle);
    void Release(VirtualizedBufferHandle handle);
    void Write(VirtualizedBufferHandle handle, usize offset, usize size, const void *data);
    void Destroy(const Device *device);
};

struct RenderResourceManager
{
  private:
    union WriteInfo {
        vk::DescriptorBufferInfo uBufferInfo;
        vk::DescriptorImageInfo uImageInfo;
        vk::BufferView uBufferView;

        WriteInfo()
        {
        }

        explicit WriteInfo(vk::DescriptorBufferInfo info);
        explicit WriteInfo(vk::DescriptorImageInfo info);
        explicit WriteInfo(vk::BufferView info);
    };

    enum class HandleType
    {
        eBuffer,
        eTexture,
        eStorageTexture,
    };

    using WriteOwner = eastl::pair<HandleType, u32>;

    eastl::deque<WriteInfo> m_WriteInfos;
    eastl::vector<vk::WriteDescriptorSet> m_Writes;
    eastl::vector<WriteOwner> m_WriteOwner;

    vk::Sampler m_DefaultSampler;

    VirtualizedBufferPool m_Geometry;
    VirtualizedBufferPool m_Index;
    VirtualizedBufferPool m_Material;
    BufferManager m_BufferManager;
    TextureManager m_TextureManager;
    StorageTextureManager m_StorageTextureManager;
    SamplerManager m_SamplerManager;

    void EraseWrites(u32 handleIndex, HandleType handleType);

  public:
    Device *m_Device;

    constexpr static u32 BUFFER_BINDING_INDEX = 0;
    constexpr static u32 TEXTURE_BINDING_INDEX = 1;
    constexpr static u32 STORAGE_TEXTURE_BINDING_INDEX = 2;

    constexpr static u32 UNIFIED_GEOMETRY_DATA_HANDLE_INDEX = 0;
    constexpr static u32 UNIFIED_GEOMETRY_OFFSET_HANDLE_INDEX = 1;
    constexpr static u32 MATERIAL_HANDLE_INDEX = 2;

    vk::SamplerCreateInfo m_DefaultSamplerCreateInfo;

    vk::DescriptorPool m_DescriptorPool;
    vk::DescriptorSetLayout m_SetLayout;
    vk::DescriptorSet m_DescriptorSet;

    bool m_UseBufferAddr;

    BufferHandle Commit(StorageBuffer *storageBuffer);                           // Commit to GPU and take Ownership
    void Write(BufferHandle handle, usize offset, usize size, const void *data); // Write to buffer
    void Release(BufferHandle handle);                                           // Release and Destroy
    void Release(StorageBuffer *storageBuffer, BufferHandle handle);             // Release and Return

    TextureHandle CommitTexture(Texture *texture, SamplerHandle sampler = {}); // Commit to GPU and take Ownership
    void Release(TextureHandle handle);                                        // Release and Destroy
    void Release(Texture *texture, TextureHandle handle);                      // Release and Return

    StorageTextureHandle CommitStorageTexture(StorageTexture *storageTexture,
                                              SamplerHandle sampler = {}); // Commit to GPU and take Ownership
    void Release(StorageTextureHandle handle);                             // Release and Destroy
    void Release(StorageTexture *texture, StorageTextureHandle handle);    // Release and Return

    SamplerHandle CreateSampler(const vk::SamplerCreateInfo *samplerCreateInfo);

    GeometryHandle CreateGeometryBuffer(usize size, usize alignment, uptr *addr = nullptr);
    uptr FetchAddress(GeometryHandle handle);
    void Write(GeometryHandle handle, usize offset, usize size, const void *data);
    void Release(GeometryHandle handle);

    MaterialHandle CreateMaterialBuffer(usize size, usize alignment, uptr *addr = nullptr);
    uptr FetchAddress(MaterialHandle handle);
    void Write(MaterialHandle handle, usize offset, usize size, const void *data);
    void Release(MaterialHandle handle);

    IndexHandle CreateIndexBuffer(usize size, usize alignment, u32 *firstIndex = nullptr);
    u32 FetchIndex(IndexHandle handle);
    void Write(IndexHandle handle, usize offset, usize size, const void *data);
    void Release(IndexHandle handle);

    vk::Buffer GetIndexBuffer() const;

    void Update(); // Update all the descriptors required.

    // Ctor/Dtor
    RenderResourceManager(Device *device, u16 maxSize, bool useBufferAddress = true);
    ~RenderResourceManager();

    RenderResourceManager(RenderResourceManager &&other) noexcept;
    RenderResourceManager &operator=(RenderResourceManager &&other) noexcept;

#if !defined(ASTER_NDEBUG)
    usize m_CommitedBufferCount = 0;
    usize m_CommitedTextureCount = 0;
    usize m_CommitedStorageTextureCount = 0;
#endif

    DISALLOW_COPY_AND_ASSIGN(RenderResourceManager);
};