#pragma once

#include "forward.hpp"
#include "../../../core/flags.hpp"
#include "api/Memory.hpp"
#include <cassert>
#include <vector>
#include <optional>
#include <functional>

namespace render::vk {

struct Allocation {
    Allocation(VkDevice, VkDeviceMemory, VkDeviceSize, uint32_t, bool optimalTiling, void *ptr);
    ~Allocation();

    const VkDevice device;
    const VkDeviceMemory memory;
    const VkDeviceSize size;
    const uint32_t memoryType;
    const bool optimalTiling;
    void *const ptr = nullptr;
    const memory::Deleter deleter;

    struct area { VkDeviceSize size; VkDeviceSize offset; };
    std::vector<area> areas;
};
class Allocator final {
public:
    Allocator(VkDevice, const PhysicalDeviceInfo&);
    ~Allocator();

    memory::ptr allocate(VkMemoryRequirements, VkMemoryPropertyFlags, bool optimalTiling = false);
    bool deallocate(const memory::area&);

    void transfer(std::function<void(VkCommandBuffer)>);
    void copyBuffer(VkBuffer src, VkBuffer dst, VkDeviceSize size);
    void transitionImageLayout(VkImage image, VkFormat format, VkImageLayout oldLayout, VkImageLayout newLayout, uint32_t mipLevels, uint32_t arrayLayers);
    void copyBufferToImage(VkBuffer src, VkImage dst, uint32_t width, uint32_t height, uint32_t mipLevels = 1, uint32_t arrayLayer = 0);

    struct Capabilities {
        std::optional<float> maxAnisotropy;
        float maxLodBias;
    };

    constexpr VkDevice getDevice() const { return device; }
    constexpr Capabilities getCapabilities() const { return capabilities; }

    static void MakeDefault(Allocator *ptr) { sInstance = ptr; }
    static _FORCE_INLINE_ Allocator *GetDefault() {
        assert(sInstance != nullptr && "No default allocator");
        return sInstance;
    }

private:
    std::optional<uint32_t> findMemory(uint32_t, VkMemoryPropertyFlags, VkDeviceSize size = 0);
    constexpr bool hasBudget() const { return properties2.pNext != nullptr; }
    constexpr const VkPhysicalDeviceMemoryProperties &getProperties() const { return hasBudget() ? properties2.memoryProperties : properties; }
    void updateProperties();

    VkPhysicalDevice const physicalDevice;
    Capabilities const capabilities;
    VkDevice const device;
    VkPhysicalDeviceMemoryProperties properties{};
    VkPhysicalDeviceMemoryProperties2 properties2{};
    VkPhysicalDeviceMemoryBudgetPropertiesEXT budget{};

    VkQueue transferQueue;
    VkCommandPool transferPool;
    VkCommandBuffer transferBuffer; // MAYBE: parallel upload
    VkQueue graphicsQueue;
    VkCommandPool graphicsPool;
    VkCommandBuffer graphicsBuffer;

    std::vector<std::unique_ptr<Allocation>> allocations;

    static Allocator *sInstance;
};
}