// =============================================
//  Aster: gpu_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 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 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 GpuResourceManager
{
  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_ImmutableSampler;

    BufferManager m_BufferManager;
    TextureManager m_TextureManager;
    StorageTextureManager m_StorageTextureManager;

    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;

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

    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);        // 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); // What they said ^
    void Release(StorageTextureHandle handle);                                 //
    void Release(StorageTexture *texture, StorageTextureHandle handle);        //

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

    // Ctor/Dtor
    GpuResourceManager(Device *device, u16 maxSize);
    ~GpuResourceManager();

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

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

    DISALLOW_COPY_AND_ASSIGN(GpuResourceManager);
};