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

#pragma once

#include "global.h"
#include "buffer.h"
#include "image.h"

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

struct Device;
struct Texture;
struct UniformStorageBuffer;

struct GpuResourceManager;

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 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 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,
    };

    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_Sampler;

    //FreeList m_BufferFreeList;
    BufferManager m_BufferManager;
    TextureManager m_TextureManager;

    void EraseWrites(u32 handleIndex, HandleType handleType);

public:
    const Device *m_Device;

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

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

    BufferHandle Commit(StorageBuffer *storageBuffer);
    void Write(BufferHandle handle, usize offset, usize size, const void *data);
    void Release(BufferHandle handle);
    TextureHandle Commit(Texture *texture);
    void Release(TextureHandle handle);

    void Update();

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

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

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

    DISALLOW_COPY_AND_ASSIGN(GpuResourceManager);
};