#include "Memory/LinearAllocator.h"
#include "Memory/PoolAllocator.h"
#include "Memory/ProxyAllocator.h"
#include "Memory/MemoryManager.h"
#include "Threading/Mutex.h"

#include <cstdlib>
#include <cstring>
#include <iostream>

namespace XCEngine {
namespace Memory {

class SystemAllocator : public IAllocator {
public:
    void* Allocate(size_t size, size_t alignment) override {
        if (alignment == 0) {
            return std::malloc(size);
        }
#ifdef _WIN32
        return _aligned_malloc(size, alignment);
#else
        return std::aligned_alloc(alignment, size);
#endif
    }

    void Free(void* ptr) override {
#ifdef _WIN32
        _aligned_free(ptr);
#else
        std::free(ptr);
#endif
    }

    void* Reallocate(void* ptr, size_t newSize) override {
        return std::realloc(ptr, newSize);
    }

    size_t GetTotalAllocated() const override { return 0; }
    size_t GetTotalFreed() const override { return 0; }
    size_t GetPeakAllocated() const override { return 0; }
    size_t GetAllocationCount() const override { return 0; }

    const char* GetName() const override { return "SystemAllocator"; }
};

LinearAllocator::LinearAllocator(size_t size, IAllocator* parent)
    : m_capacity(size), m_parent(parent) {
    if (parent) {
        m_buffer = static_cast<uint8_t*>(parent->Allocate(size, 8));
    } else {
        m_buffer = static_cast<uint8_t*>(_aligned_malloc(size, 8));
    }
}

LinearAllocator::~LinearAllocator() {
    if (m_parent) {
        m_parent->Free(m_buffer);
    } else {
#ifdef _WIN32
        _aligned_free(m_buffer);
#else
        std::free(m_buffer);
#endif
    }
}

void* LinearAllocator::Allocate(size_t size, size_t alignment) {
    if (size == 0) {
        return nullptr;
    }

    uintptr_t address = reinterpret_cast<uintptr_t>(m_buffer) + m_offset;
    
    if (alignment > 0) {
        size_t misalignment = address % alignment;
        if (misalignment != 0) {
            size += alignment - misalignment;
        }
    }

    if (m_offset + size > m_capacity) {
        return nullptr;
    }

    void* ptr = &m_buffer[m_offset];
    m_offset += size;
    return ptr;
}

void LinearAllocator::Free(void* ptr) {
}

void* LinearAllocator::Reallocate(void* ptr, size_t newSize) {
    return nullptr;
}

void LinearAllocator::Clear() {
    m_offset = 0;
}

void* LinearAllocator::GetMarker() const {
    return reinterpret_cast<void*>(m_offset);
}

void LinearAllocator::SetMarker(void* marker) {
    m_offset = reinterpret_cast<size_t>(marker);
}

PoolAllocator::PoolAllocator(size_t blockSize, size_t poolSize, size_t alignment)
    : m_blockSize(blockSize), m_alignment(alignment), m_totalBlocks(poolSize), m_freeBlocks(poolSize) {
    
    size_t actualBlockSize = blockSize;
    if (alignment > 0) {
        actualBlockSize = (blockSize + alignment - 1) & ~(alignment - 1);
    }

    m_memory = std::malloc(actualBlockSize * poolSize);
    
    uint8_t* memory = static_cast<uint8_t*>(m_memory);
    m_freeList = reinterpret_cast<FreeNode*>(memory);
    
    FreeNode* current = m_freeList;
    for (size_t i = 1; i < poolSize; ++i) {
        uint8_t* block = memory + (i * actualBlockSize);
        current->next = reinterpret_cast<FreeNode*>(block);
        current = current->next;
    }
    current->next = nullptr;
}

PoolAllocator::~PoolAllocator() {
    std::free(m_memory);
}

void* PoolAllocator::Allocate(size_t size, size_t alignment) {
    if (!m_freeList) {
        return nullptr;
    }

    if (size > m_blockSize) {
        return nullptr;
    }

    FreeNode* node = m_freeList;
    m_freeList = m_freeList->next;
    --m_freeBlocks;
    
    return node;
}

void PoolAllocator::Free(void* ptr) {
    if (!ptr) {
        return;
    }

    FreeNode* node = static_cast<FreeNode*>(ptr);
    node->next = m_freeList;
    m_freeList = node;
    ++m_freeBlocks;
}

void* PoolAllocator::Reallocate(void* ptr, size_t newSize) {
    return nullptr;
}

bool PoolAllocator::Contains(void* ptr) const {
    if (!ptr || !m_memory) {
        return false;
    }
    uint8_t* memory = static_cast<uint8_t*>(m_memory);
    uint8_t* p = static_cast<uint8_t*>(ptr);
    uintptr_t offset = p - memory;
    size_t blockSize = m_blockSize;
    if (m_alignment > 0) {
        blockSize = (m_blockSize + m_alignment - 1) & ~(m_alignment - 1);
    }
    return offset < blockSize * m_totalBlocks;
}

size_t PoolAllocator::GetFreeBlockCount() const {
    return m_freeBlocks;
}

ProxyAllocator::ProxyAllocator(IAllocator* underlying, const char* name)
    : m_underlying(underlying), m_name(name) {
}

void* ProxyAllocator::Allocate(size_t size, size_t alignment) {
    std::lock_guard<Threading::Mutex> lock(m_mutex);
    void* ptr = m_underlying->Allocate(size, alignment);
    if (ptr) {
        m_stats.totalAllocated += size;
        m_stats.allocationCount++;
        if (m_stats.totalAllocated - m_stats.totalFreed > m_stats.peakAllocated) {
            m_stats.peakAllocated = m_stats.totalAllocated - m_stats.totalFreed;
        }
    }
    return ptr;
}

void ProxyAllocator::Free(void* ptr) {
    std::lock_guard<Threading::Mutex> lock(m_mutex);
    m_underlying->Free(ptr);
    m_stats.totalFreed += m_stats.allocationCount;
    m_stats.allocationCount--;
}

void* ProxyAllocator::Reallocate(void* ptr, size_t newSize) {
    std::lock_guard<Threading::Mutex> lock(m_mutex);
    void* newPtr = m_underlying->Reallocate(ptr, newSize);
    return newPtr;
}

const ProxyAllocator::Stats& ProxyAllocator::GetStats() const {
    return m_stats;
}

MemoryManager& MemoryManager::Get() {
    static MemoryManager instance;
    return instance;
}

void MemoryManager::Initialize() {
    if (m_initialized) {
        return;
    }

    m_systemAllocator = new SystemAllocator();
    m_initialized = true;
}

void MemoryManager::Shutdown() {
    if (!m_initialized) {
        return;
    }

    delete static_cast<SystemAllocator*>(m_systemAllocator);
    m_systemAllocator = nullptr;
    m_initialized = false;
}

IAllocator* MemoryManager::GetSystemAllocator() {
    return m_systemAllocator;
}

std::unique_ptr<LinearAllocator> MemoryManager::CreateLinearAllocator(size_t size) {
    return std::make_unique<LinearAllocator>(size, m_systemAllocator);
}

std::unique_ptr<PoolAllocator> MemoryManager::CreatePoolAllocator(size_t blockSize, size_t count) {
    return std::make_unique<PoolAllocator>(blockSize, count);
}

std::unique_ptr<ProxyAllocator> MemoryManager::CreateProxyAllocator(const char* name) {
    return std::make_unique<ProxyAllocator>(m_systemAllocator, name);
}

void MemoryManager::SetTrackAllocations(bool track) {
    m_trackAllocations = track;
}

void MemoryManager::DumpMemoryLeaks() {
    std::cout << "Memory Leak Report:" << std::endl;
}

void MemoryManager::GenerateMemoryReport() {
    std::cout << "Memory Report:" << std::endl;
}

} // namespace Memory
} // namespace XCEngine