#pragma once

// Windows
#include <heapapi.h>

// STD
#include <span>
#include <algorithm>
#include <type_traits>

#include "bitmask.hpp"

#ifdef min
#undef min
#endif

#ifdef ZeroMemory
#undef ZeroMemory
#endif

enum class HeapFlags : uint32_t
{
    None               = 0x00000000,
    NoSerialize        = 0x00000001, // HEAP_NO_SERIALIZE
    GenerateExceptions = 0x00000004, // HEAP_GENERATE_EXCEPTIONS
    ZeroMemory         = 0x00000008, // HEAP_ZERO_MEMORY
    ReAllocInPlaceOnly = 0x00000010, // HEAP_REALLOC_IN_PLACE_ONLY
    // Custom.
    AllowShrinking     = 0x40000000,
    NoPreserveContent  = 0x80000000,
    // Masks.
    AllocMask          = NoSerialize | ZeroMemory | GenerateExceptions,
    ReAllocMask        = AllocMask | ReAllocInPlaceOnly,
    FreeMask           = NoSerialize,
};

PVOID HeapAlloc(HANDLE hHeap, size_t size, HeapFlags flags)
{
    return HeapAlloc(hHeap, DWORD(flags & HeapFlags::AllocMask), size);
}

PVOID HeapReAlloc(HANDLE hHeap, PVOID pMem, size_t size, HeapFlags flags)
{
    return HeapReAlloc(hHeap, DWORD(flags & HeapFlags::ReAllocMask), pMem, size);
}

BOOL HeapFree(HANDLE hHeap, PVOID pMem, HeapFlags flags)
{
    return HeapFree(hHeap, DWORD(flags & HeapFlags::FreeMask), pMem);
}

class Heap final
{
public:
    Heap() noexcept = default;

    explicit Heap(size_t capacity, size_t size = 0, HeapFlags flags = HeapFlags::None)
    {
        if (!realloc(size, capacity, flags))
            throw std::bad_alloc{};
    }

    Heap(const void* ptr, size_t size, size_t capacity, HeapFlags flags = HeapFlags::None)
    {
        if (!assign(ptr, size, capacity, flags))
            throw std::bad_alloc{};
    }

    Heap(const void* ptr, size_t size, HeapFlags flags = HeapFlags::None)
        : Heap(ptr, size, size, flags)
    { }

    Heap(const Heap& other, size_t capacity, HeapFlags flags = HeapFlags::None)
        : Heap(other.m_ptr, other.m_size, capacity, flags)
    { }

    Heap(const Heap& other, HeapFlags flags = HeapFlags::None)
        : Heap(other, other.m_size, flags)
    { }

    Heap(Heap&& other) noexcept
        : m_ptr(std::exchange(other.m_ptr, nullptr))
        , m_capacity(std::exchange(other.m_capacity, 0))
        , m_size(std::exchange(other.m_size, 0))
    { }

    ~Heap()
    {
        reset();
    }

    void reset() noexcept
    {
        if (m_ptr) HeapFree(m_heap, m_ptr, HeapFlags::None);
        m_ptr = nullptr;
        m_capacity = m_size = 0;
    }

    [[nodiscard]]
    uint8_t* release() noexcept
    {
        m_capacity = m_size = 0;
        return std::exchange(m_ptr, nullptr);
    }

    [[nodiscard]]
    auto data(this auto& self) noexcept
    {
        using byte_t = std::conditional_t<std::is_const_v<std::remove_reference_t<decltype(self)>>, const uint8_t, uint8_t>;
        return static_cast<byte_t*>(self.m_ptr);
    }

    [[nodiscard]]
    auto bytes(this auto& self) noexcept
    {
        return std::span{self.data(), self.m_size};
    }

    [[nodiscard]]
    std::size_t capacity() const noexcept
    {
        return m_capacity;
    }

    [[nodiscard]]
    std::size_t size() const noexcept
    {
        return m_size;
    }

    [[nodiscard]]
    bool empty() const noexcept
    {
        return m_size == 0;
    }

    bool realloc(size_t size, size_t capacity, HeapFlags flags = HeapFlags::None)
    {
        if (capacity == m_capacity || (capacity < m_capacity && !BITALL(flags, HeapFlags::AllowShrinking)))
        {
            m_size = std::min(size, m_capacity);
            return m_ptr;
        }

        if (capacity == 0)
        {
            reset();
            return true;
        }

        uint8_t* ptr;

        if (!m_ptr)
        {
            ptr = (uint8_t*)HeapAlloc(m_heap, capacity, flags);
        }
        else if (!BITALL(flags, HeapFlags::NoPreserveContent) || BITALL(flags, HeapFlags::ReAllocInPlaceOnly))
        {
            ptr = (uint8_t*)HeapReAlloc(m_heap, m_ptr, capacity, flags);
        }
        else if ((ptr = (uint8_t*)HeapAlloc(m_heap, capacity, flags)))
        {
            HeapFree(m_heap, m_ptr, flags);
        }

        if (!ptr) return false;

        m_ptr = ptr;
        m_capacity = capacity;
        m_size = std::min(size, m_capacity);

        return true;
    }

    bool reserve(size_t capacity, HeapFlags flags = HeapFlags::None)
    {
        return realloc(m_size, capacity, flags);
    }

    bool resize(size_t size, HeapFlags flags = HeapFlags::None)
    {
        return realloc(size, size, flags);
    }

    void fill(size_t offset, size_t count, uint8_t value = 0)
    {
        if (offset >= m_capacity) return;
        std::fill_n(m_ptr + offset, std::min(count, m_capacity - offset), value);
    }

    // Moves the offset back if necessary.
    void fill_shifted(size_t offset, size_t count, uint8_t value = 0)
    {
        count = std::min(count, m_capacity);
        std::fill_n(m_ptr + std::min(offset, m_capacity - count), count, value);
    }

    bool assign(const void* ptr, size_t size, size_t capacity, HeapFlags flags = HeapFlags::None)
    {
        const auto src = static_cast<const std::uint8_t*>(ptr);
        const auto aliasing = m_ptr && src && size && std::less_equal<>{}(m_ptr, src) && std::less<>{}(src, m_ptr + m_capacity);
        flags = aliasing ? (flags | HeapFlags::NoPreserveContent) : (flags & ~HeapFlags::NoPreserveContent);
        if (!realloc(size, capacity, flags)) return false;
        std::memmove(m_ptr, src, m_size);
        return true;
    }

    bool assign(const void* ptr, size_t size, HeapFlags flags = HeapFlags::None)
    {
        return assign(ptr, size, size, flags);
    }

    bool assign(const Heap& other, size_t capacity, HeapFlags flags = HeapFlags::None)
    {
        return assign(other.m_ptr, other.size(), capacity, flags);
    }

    bool assign(const Heap& other, HeapFlags flags = HeapFlags::None)
    {
        return assign(other, other.m_size, flags);
    }

    Heap& operator=(const Heap& other)
    {
        if (this != &other && !assign(other.m_ptr, other.m_size, other.m_size))
            throw std::bad_alloc{};
        return *this;
    }

    Heap& operator=(Heap&& other) noexcept
    {
        if (this != &other)
        {
            reset();
            m_ptr = std::exchange(other.m_ptr, nullptr);
            m_capacity = std::exchange(other.m_capacity, 0);
            m_size = std::exchange(other.m_size, 0);
        }
        return *this;
    }

    auto& operator[](this auto& self, std::size_t idx) noexcept
    {
        return self.data()[idx];
    }

    [[nodiscard]]
    friend bool operator==(const Heap& lhs, const Heap& rhs)
    {
        const auto size = lhs.size();
        if (size != rhs.size()) return false;
        if (size == 0) return true;
        return !std::memcmp(lhs.m_ptr, rhs.m_ptr, size);
    }
private:
    uint8_t* m_ptr = nullptr;
    size_t m_capacity = 0;
    size_t m_size = 0;

    inline static const auto m_heap = GetProcessHeap();
};