#pragma once

#include <vk/forward.h>

#include <vulkan/vulkan.h>

#include <memory>
#include <vector>


namespace vk {


class MemoryBlock {
public:
    MemoryBlock() noexcept;
    MemoryBlock(
        VkDeviceMemory device_memory,
        VkDeviceSize size,
        VkDeviceSize offset,
        MemoryPage* memory_page,
        uint32_t memory_type
    ) noexcept;
    MemoryBlock(const MemoryBlock&) = delete;
    MemoryBlock(MemoryBlock&& other) noexcept;
    ~MemoryBlock() noexcept;

    MemoryBlock& operator=(const MemoryBlock&) = delete;
    MemoryBlock& operator=(MemoryBlock&& other) noexcept;

    inline explicit operator bool() const noexcept {
        return m_size != 0;
    }

    inline bool is_valid() const noexcept {
        return m_size != 0;
    }

    inline VkDeviceSize size() const noexcept {
        return m_size;
    }

    inline VkDeviceSize offset() const noexcept {
        return m_offset;
    }

    inline VkDeviceMemory device_memory() const noexcept {
        return m_device_memory;
    }

    inline uint32_t memory_type() const noexcept {
        return m_memory_type;
    }

    inline MemoryPage* memory_page() const noexcept {
        return m_memory_page;
    }

    VkMemoryType memory_type_info(Context& context) const noexcept;

    void free() noexcept;

    void* map(Context& context);
    void* map(Context& context, VkDeviceSize offset, VkDeviceSize size);
    void unmap(Context& context) noexcept;

    void flush(Context& context);
    void invalidate(Context& context);

private:
    VkDeviceSize m_size = 0;
    VkDeviceSize m_offset = 0;
    VkDeviceMemory m_device_memory = VK_NULL_HANDLE;
    MemoryPage* m_memory_page = nullptr;
    uint32_t m_memory_type = 0;
};


class MemoryPage {
public:
    MemoryPage(
        VkDeviceSize size,
        VkDeviceMemory device_memory,
        uint32_t memory_type
    ) noexcept;
    MemoryPage(const MemoryPage&) = delete;
    MemoryPage(MemoryPage&&);
    ~MemoryPage();

    MemoryPage& operator=(const MemoryPage&) = delete;
    MemoryPage& operator=(MemoryPage&&);

    explicit inline operator bool() const noexcept {
        return is_valid();
    }

    inline bool is_valid() const noexcept {
        return m_size != 0;
    }

    inline VkDeviceSize size() const noexcept {
        return m_size;
    }

    inline VkDeviceMemory device_memory() const noexcept {
        return m_device_memory;
    }

    inline uint32_t memory_type() const noexcept {
        return m_memory_type;
    }

    MemoryBlock allocate(VkDeviceSize size) noexcept;
    void p_free(MemoryBlock& block) noexcept;

    void free_device_memory(Context& context) noexcept;

private:
    struct Block {
        VkDeviceSize offset;
        bool is_free;
    };
    using BlockList = std::vector<Block>;

private:
    BlockList::iterator find_free_block(VkDeviceSize size);

private:
    VkDeviceSize m_size = 0;
    VkDeviceMemory m_device_memory = VK_NULL_HANDLE;
    BlockList m_blocks;
    uint32_t m_memory_type = 0;
};


class Allocator {
public:
    Allocator(Context* context) noexcept;
    Allocator(const Allocator&) = delete;
    Allocator(Allocator&&) = delete;
    ~Allocator();

    Allocator& operator=(const Allocator&) = delete;
    Allocator& operator=(Allocator&) = delete;

    int32_t find_memory_type(
        uint32_t type_mask,
        VkMemoryPropertyFlags properties
    ) noexcept;
    int32_t find_memory_type(
        uint32_t type_mask,
        std::initializer_list<VkMemoryPropertyFlags> properties_list
    ) noexcept;

    MemoryBlock allocate(
        VkDeviceSize size,
        uint32_t memory_type
    ) noexcept;
    MemoryBlock allocate(
        VkDeviceSize size,
        uint32_t type_mask,
        VkMemoryPropertyFlags properties
    ) noexcept;
    MemoryBlock allocate(
        VkDeviceSize size,
        uint32_t type_mask,
        std::initializer_list<VkMemoryPropertyFlags> properties_list
    ) noexcept;

    void free_all_pages() noexcept;

private:
    using PageList = std::vector<MemoryPage>;

    static constexpr VkDeviceSize BasePageSize = 1 << 24;

private:
    Context* m_context = nullptr;
    VkPhysicalDeviceMemoryProperties m_memory_properties;
    PageList m_page_map[VK_MAX_MEMORY_TYPES];
    VkDeviceSize m_next_page_sizes[VK_MAX_MEMORY_TYPES];
};


}