#include "deimos/core/temp_allocator.h"

#include "deimos/core/api_registry.h"
#include "deimos/core/allocator.h"
#include "deimos/core/os.h"

static deimos::OsApi* os_api;

namespace deimos
{

class TempAllocatorImpl : public ITempAllocator
{
    static constexpr int64_t kDefaultAlign    = 8;
    static constexpr int64_t kPageSize        = 1 * Megabytes;
    static constexpr int64_t kReserveSize     = 256 * Megabytes;
    
    gsl::owner<void*> m_base{};
    void*             m_commit_end{};
    void*             m_reserve_end{};

    // @Todo Last allocation optimization
    void*             m_current{};
    
public:
    gsl::owner<void*> Reallocate(TempAllocatorTag tag,
        gsl::owner<void*> old_ptr, int64_t old_size, int64_t new_size) override
    {
        void* tag_ptr = std::bit_cast<void*>(tag.tag);
        Expects(tag_ptr >= m_base && tag_ptr <= m_commit_end && tag_ptr <= m_current);
        
        if (new_size <= old_size) { return (new_size == 0) ? nullptr : old_ptr; }

        void* new_current = OffsetBytes(m_current, AlignUp(new_size, kDefaultAlign));
        if (new_current > m_reserve_end)
        {
            deimos_Panic("Ran out of temporary memory");
        }

        if (new_current > m_commit_end)
        {
            const int64_t new_commit_size = AlignUp(
                (int64_t)(std::bit_cast<uintptr_t>(new_current) - std::bit_cast<uintptr_t>(m_base)), // NOLINT
                kPageSize);
            os_api->virtual_memory->Commit(m_base, new_commit_size);
            m_commit_end = OffsetBytes(m_base, new_commit_size);
            Ensures(m_commit_end <= m_reserve_end);
        }

        if (old_ptr != nullptr)
        {
            MemoryCopy(m_current, old_ptr, old_size);
        }
        
        return std::exchange(m_current, new_current);
    }
    
    TempAllocator Acquire()
    {
        if (m_base == nullptr)
        {
            m_base        = os_api->virtual_memory->Reserve(kReserveSize);
            m_current     = m_base;
            m_commit_end  = m_base;
            m_reserve_end = OffsetBytes(m_base, kReserveSize);
        }
        
        return TempAllocator(this, {std::bit_cast<uintptr_t>(m_current)});
    }
    
    void Release(TempAllocatorTag tag) override
    {
        void* rewind_base = std::bit_cast<void*>(tag.tag);
        Expects(rewind_base >= m_base && rewind_base <= m_commit_end);
        Expects(rewind_base <= m_current);
        m_current = rewind_base;
    }
};

static thread_local TempAllocatorImpl g_impl;

class TempAllocatorApiImpl : public TempAllocatorApi
{
public:
    TempAllocator Acquire() override
    {
        return g_impl.Acquire();
    }
};

void RegisterTempAllocatorApi(ApiRegistry* api_registry)
{
    os_api = api_registry->Get<OsApi>();
    
    auto* allocator = api_registry->Get<AllocatorApi>()->system;
    gsl::owner<TempAllocatorApi*> temp_allocator_api = allocator->New<TempAllocatorApiImpl>();
    api_registry->Set(temp_allocator_api);
}

} // namespace deimos