#pragma once

#include "deimos/core/base.h"
#include "deimos/core/id_name.h"

namespace deimos
{

struct MemoryScope { uint32_t id; };

struct IAllocator
{
    IAllocator() = default;

    deimos_NO_COPY_MOVE(IAllocator);

    virtual ~IAllocator() = default;

    [[nodiscard]]
    virtual gsl::owner<void*> Reallocate(
        gsl::owner<void*> old_ptr, int64_t old_size, int64_t new_size,
        MemoryScope scope, const SourceLocation& source_location = {}) = 0;
};


class Allocator
{
    gsl::owner<IAllocator*> m_allocator;
    const MemoryScope m_scope;

public:
    constexpr Allocator(IAllocator* allocator, MemoryScope scope) :
        m_allocator{allocator}, m_scope{scope}
    {}

    deimos_NO_COPY_MOVE(Allocator);

    ~Allocator() = default;

    constexpr IAllocator* allocator() const { return m_allocator; }

    [[nodiscard]]
    constexpr gsl::owner<void*> Allocate(
        int64_t new_size,
        const SourceLocation& source_location = {})
    {
        return m_allocator->Reallocate(nullptr, 0, new_size, m_scope, source_location);
    }

    [[nodiscard]]
    gsl::owner<void*> Reallocate(
        gsl::owner<void*> old_ptr, int64_t old_size, int64_t new_size,
        const SourceLocation& source_location = {})
    {
        return m_allocator->Reallocate(old_ptr, old_size, new_size, m_scope, source_location);
    }

    constexpr void Free(
        gsl::owner<void*> old_ptr, int64_t old_size,
        const SourceLocation& source_location = {})
    {
        (void)m_allocator->Reallocate(old_ptr, old_size, 0, m_scope, source_location);
    }

    template<typename T, typename... Args>
    gsl::owner<T*> NewInner(
        const SourceLocation& source_location,
        Args&&... args)
    {
        gsl::owner<void*> ptr = Allocate(sizeof(T), source_location);
        return new(ptr) T(std::forward<Args>(args)...);
    }

    template<typename T>
    constexpr gsl::owner<T*> New(const SourceLocation& source_location = {})
    {
        return NewInner<T>(source_location);
    }

    template<typename T, typename A0>
    constexpr gsl::owner<T*> New(A0&& arg0, const SourceLocation& source_location = {})
    {
        return NewInner<T>(source_location, std::forward<A0>(arg0));
    }

    template<typename T, typename A0, typename A1>
    constexpr gsl::owner<T*> New(
        A0&& arg0, A1&& arg1,
        const SourceLocation& source_location = {})
    {
        return NewInner<T>(source_location,
            std::forward<A0>(arg0),
            std::forward<A1>(arg1));
    }

    template<typename T, typename A0, typename A1, typename A2>
    constexpr gsl::owner<T*> New(
        A0&& arg0, A1&& arg1, A2&& arg2,
        const SourceLocation& source_location = {})
    {
        return NewInner<T>(source_location,
            std::forward<A0>(arg0),
            std::forward<A1>(arg1),
            std::forward<A2>(arg2));
    }

    template<typename T, typename A0, typename A1, typename A2, typename A3>
    constexpr gsl::owner<T*> New(
        A0&& arg0, A1&& arg1, A2&& arg2, A3&& arg3,
        const SourceLocation& source_location = {})
    {
        return NewInner<T>(source_location,
            std::forward<A0>(arg0),
            std::forward<A1>(arg1),
            std::forward<A2>(arg2),
            std::forward<A3>(arg3));
    }

    template<typename T>
    void Delete(gsl::owner<T*> t, const SourceLocation& source_location = {})
    {
        if constexpr (!std::is_trivially_destructible_v<T>)
        {
            t->~T();
        }
        Free(t, sizeof(T), source_location);
    }

    template<typename T>
    gsl::owner<Span<T>> NewArray(int64_t count, const SourceLocation& source_location = {})
        requires std::is_default_constructible_v<T>
    {
        Expects(count > 0);
        
        auto* raw = Allocate((int64_t)sizeof(T) * count, source_location);
        if constexpr (std::is_trivially_default_constructible_v<T>)
        {
            MemoryZero(raw, (int64_t)sizeof(T) * count);
        }
        else
        {
            for (int64_t i = 0; i < count; ++i)
            {
                new((T*)raw + i) T{};
            }
        }
        return Span<T>{ (T*)raw, count };
    }
};

class AllocatorApi
{
public:
    AllocatorApi() = default;

    deimos_NO_COPY_MOVE(AllocatorApi);

    virtual ~AllocatorApi() = default;

    static constexpr IdName kApiName{"deimos::AllocatorApi"};

    Allocator* system{};

    virtual gsl::owner<Allocator*> CreateChild(Allocator* parent, gsl::czstring description) = 0;
    virtual void DestroyChild(gsl::owner<Allocator*>) = 0;
};

} // namespace deimos