//
// Created by atdunbg on 2026/9/20.
//

#include <Engine/Vulkan/MemoryAllocator.h>
#include <Engine/Vulkan/Device.h>
#include <Engine/Core/Log.h>

#include <algorithm>

#include "Engine/Core/Stats.h"

namespace Engine {

MemoryAllocator::MemoryAllocator(Device& device)
    : m_Device(device)
{
    ENGINE_INFO("MemoryAllocator created");
}

MemoryAllocator::~MemoryAllocator()
{
    const auto& vkDevice = m_Device.GetDevice();

    for (auto& blocks : m_Blocks) {
        for (auto& block : blocks) {
            if (block.IsMapped) {
                vkDevice.unmapMemory(block.Memory);
            }
            vkDevice.freeMemory(block.Memory);

            STATS.RemoveMemoryBlock(block.Size);
        }
        blocks.clear();
    }
}

vk::DeviceSize MemoryAllocator::GetBlockSize(vk::MemoryPropertyFlags props) const
{
    if (props & vk::MemoryPropertyFlagBits::eHostVisible) {
        return 16ull * 1024 * 1024;   // 16 MB
    }
    return 64ull * 1024 * 1024;       // 64 MB
}

MemoryAllocator::Allocation MemoryAllocator::Allocate(
    vk::MemoryRequirements reqs, vk::MemoryPropertyFlags props)
{
    u32 typeIndex = m_Device.FindMemoryType(reqs.memoryTypeBits, props);
    auto& blocks = m_Blocks[typeIndex];

    // 1. 扫描已有块的 FreeList，找能放下的区域（first-fit）
    for (u32 i = 0; i < blocks.size(); ++i) {
        auto& block = blocks[i];
        for (auto it = block.FreeList.begin(); it != block.FreeList.end(); ++it) {
            vk::DeviceSize alignedOffset =
                (it->Offset + reqs.alignment - 1) & ~(reqs.alignment - 1);
            vk::DeviceSize end = alignedOffset + reqs.size;

            if (end <= it->Offset + it->Size) {
                // 找到了
                Allocation alloc;
                alloc.Memory = block.Memory;
                alloc.Offset = alignedOffset;
                alloc.Size = reqs.size;
                alloc.MappedPtr = block.IsMapped
                    ? static_cast<char*>(block.MappedPtr) + alignedOffset
                    : nullptr;
                alloc.BlockIndex = i;
                alloc.TypeIndex = typeIndex;
                alloc.Valid = true;

                // 更新 FreeList
                if (alignedOffset == it->Offset && reqs.size == it->Size) {
                    // 整个区域被用了，移除
                    block.FreeList.erase(it);
                } else if (alignedOffset == it->Offset) {
                    // 从头部切
                    it->Offset += reqs.size;
                    it->Size -= reqs.size;
                } else if (end == it->Offset + it->Size) {
                    // 从尾部切
                    it->Size = alignedOffset - it->Offset;
                } else {
                    // 从中间切，拆成两段
                    FreeRange tail;
                    tail.Offset = end;
                    tail.Size = (it->Offset + it->Size) - end;
                    it->Size = alignedOffset - it->Offset;
                    block.FreeList.insert(it + 1, tail);
                }

                return alloc;
            }
        }
    }

    // 2. 在已有块的尾部 bump 分配
    for (u32 i = 0; i < blocks.size(); ++i) {
        auto& block = blocks[i];
        vk::DeviceSize alignedOffset =
            (block.Used + reqs.alignment - 1) & ~(reqs.alignment - 1);

        if (alignedOffset + reqs.size <= block.Size) {
            Allocation alloc;
            alloc.Memory = block.Memory;
            alloc.Offset = alignedOffset;
            alloc.Size = reqs.size;
            alloc.MappedPtr = block.IsMapped
                ? static_cast<char*>(block.MappedPtr) + alignedOffset
                : nullptr;
            alloc.BlockIndex = i;
            alloc.TypeIndex = typeIndex;
            alloc.Valid = true;

            block.Used = alignedOffset + reqs.size;
            return alloc;
        }
    }

    // 3. 创建新块
    vk::DeviceSize blockSize = std::max(GetBlockSize(props), reqs.size);

    vk::MemoryAllocateInfo allocInfo;
    allocInfo.setAllocationSize(blockSize)
             .setMemoryTypeIndex(typeIndex);

    Block block;
    block.Memory = m_Device.GetDevice().allocateMemory(allocInfo);
    block.Size = blockSize;
    block.Used = 0;

    if (props & vk::MemoryPropertyFlagBits::eHostVisible) {
        block.MappedPtr = m_Device.GetDevice().mapMemory(
            block.Memory, 0, VK_WHOLE_SIZE);
        block.IsMapped = true;
    }

    blocks.push_back(std::move(block));
    m_TotalAllocated[typeIndex] += blockSize;

    u32 newIndex = static_cast<u32>(blocks.size() - 1);
    auto& newBlock = blocks[newIndex];

    vk::DeviceSize alignedOffset = (reqs.alignment - 1) & ~(reqs.alignment - 1);

    Allocation alloc;
    alloc.Memory = newBlock.Memory;
    alloc.Offset = alignedOffset;
    alloc.Size = reqs.size;
    alloc.MappedPtr = newBlock.IsMapped
        ? static_cast<char*>(newBlock.MappedPtr) + alignedOffset
        : nullptr;
    alloc.BlockIndex = newIndex;
    alloc.TypeIndex = typeIndex;
    alloc.Valid = true;

    newBlock.Used = alignedOffset + reqs.size;

    ENGINE_INFO("MemoryAllocator: new block of {} bytes (type {})",
                blockSize, typeIndex);

    STATS.AddMemoryBlock(blockSize);

    STATS.AddMemoryUsed(reqs.size);
    return alloc;
}

void MemoryAllocator::Free(const Allocation& alloc)
{
    if (!alloc.Valid) return;
    STATS.RemoveMemoryUsed(alloc.Size);

    auto& block = m_Blocks[alloc.TypeIndex][alloc.BlockIndex];

    FreeRange range;
    range.Offset = alloc.Offset;
    range.Size = alloc.Size;

    // 插入 FreeList，保持按 offset 排序，方便合并
    auto it = std::lower_bound(
        block.FreeList.begin(), block.FreeList.end(), range,
        [](const FreeRange& a, const FreeRange& b) {
            return a.Offset < b.Offset;
        });

    it = block.FreeList.insert(it, range);

    // 向前合并
    if (it != block.FreeList.begin()) {
        auto prev = it - 1;
        if (prev->Offset + prev->Size == it->Offset) {
            prev->Size += it->Size;
            block.FreeList.erase(it);
            it = prev;
        }
    }

    // 向后合并
    auto next = it + 1;
    if (next != block.FreeList.end()) {
        if (it->Offset + it->Size == next->Offset) {
            it->Size += next->Size;
            block.FreeList.erase(next);
        }
    }
}

void MemoryAllocator::PrintStats() const
{
    for (u32 type = 0; type < VK_MAX_MEMORY_TYPES; ++type) {
        if (m_Blocks[type].empty()) continue;

        vk::DeviceSize totalUsed = 0;
        for (auto& b : m_Blocks[type]) totalUsed += b.Used;

        ENGINE_INFO("  Type {}: {} blocks, total allocated {} MB, used {} MB",
                    type,
                    m_Blocks[type].size(),
                    m_TotalAllocated[type] / (1024 * 1024),
                    totalUsed / (1024 * 1024));
    }
}

} // namespace Engine