#ifndef FREERTOS_CPP_MEMORY_RESOURCE_HPP
#define FREERTOS_CPP_MEMORY_RESOURCE_HPP

#include <stddef.h>
#include <stdint.h>

extern "C"
{
#include "FreeRTOS.h"
}

namespace freertos
{

class MemoryResource final
{
public:
    using Allocate = void *( * )( void *, size_t, size_t ) noexcept;
    using Deallocate = void ( * )( void *, void *, size_t, size_t ) noexcept;

    constexpr MemoryResource( void * context,
                              Allocate allocate,
                              Deallocate deallocate ) noexcept
        : context_ { context },
          allocate_ { allocate },
          deallocate_ { deallocate }
    {
    }

    void * allocate_bytes( size_t bytes, size_t alignment ) const noexcept
    {
        return allocate_( context_, bytes, alignment );
    }

    void deallocate_bytes( void * memory,
                           size_t bytes,
                           size_t alignment ) const noexcept
    {
        deallocate_( context_, memory, bytes, alignment );
    }

    constexpr void * context() const noexcept
    {
        return context_;
    }

    constexpr Allocate allocate_function() const noexcept
    {
        return allocate_;
    }

    constexpr Deallocate deallocate_function() const noexcept
    {
        return deallocate_;
    }

private:
    void * context_;
    Allocate allocate_;
    Deallocate deallocate_;
};

constexpr bool valid_alignment( size_t alignment ) noexcept
{
    return alignment != 0U &&
           ( alignment & ( alignment - 1U ) ) == 0U;
}

class HeapResource final
{
public:
    MemoryResource resource() noexcept
    {
        return MemoryResource { this, &allocate, &deallocate };
    }

    constexpr size_t successful_allocations() const noexcept
    {
        return successful_allocations_;
    }

    constexpr size_t failed_allocations() const noexcept
    {
        return failed_allocations_;
    }

    constexpr size_t deallocations() const noexcept
    {
        return deallocations_;
    }

private:
    static void * allocate( void * context,
                            size_t bytes,
                            size_t alignment ) noexcept
    {
        auto & self = *static_cast< HeapResource * >( context );
        if( bytes == 0U || !valid_alignment( alignment ) ||
            alignment > portBYTE_ALIGNMENT )
        {
            ++self.failed_allocations_;
            return nullptr;
        }

        void * const memory = pvPortMalloc( bytes );
        if( memory == nullptr )
        {
            ++self.failed_allocations_;
        }
        else
        {
            ++self.successful_allocations_;
        }
        return memory;
    }

    static void deallocate( void * context,
                            void * memory,
                            size_t,
                            size_t ) noexcept
    {
        if( memory != nullptr )
        {
            auto & self = *static_cast< HeapResource * >( context );
            vPortFree( memory );
            ++self.deallocations_;
        }
    }

    size_t successful_allocations_ {};
    size_t failed_allocations_ {};
    size_t deallocations_ {};
};

class StaticArenaResource final
{
public:
    constexpr StaticArenaResource( uint8_t * storage,
                                   size_t capacity,
                                   size_t maximum_alignment ) noexcept
        : storage_ { storage },
          capacity_ { capacity },
          maximum_alignment_ { maximum_alignment }
    {
    }

    MemoryResource resource() noexcept
    {
        return MemoryResource { this, &allocate, &deallocate };
    }

    void reset() noexcept
    {
        used_ = 0U;
        ++resets_;
    }

    constexpr size_t used() const noexcept
    {
        return used_;
    }

    constexpr size_t remaining() const noexcept
    {
        return capacity_ - used_;
    }

    constexpr size_t successful_allocations() const noexcept
    {
        return successful_allocations_;
    }

    constexpr size_t failed_allocations() const noexcept
    {
        return failed_allocations_;
    }

    constexpr size_t resets() const noexcept
    {
        return resets_;
    }

    bool owns( const void * memory ) const noexcept
    {
        const uintptr_t address = reinterpret_cast< uintptr_t >( memory );
        const uintptr_t first = reinterpret_cast< uintptr_t >( storage_ );
        return address >= first && address < first + capacity_;
    }

private:
    static void * allocate( void * context,
                            size_t bytes,
                            size_t alignment ) noexcept
    {
        auto & self = *static_cast< StaticArenaResource * >( context );
        if( bytes == 0U || !valid_alignment( alignment ) ||
            alignment > self.maximum_alignment_ )
        {
            ++self.failed_allocations_;
            return nullptr;
        }

        const uintptr_t base =
            reinterpret_cast< uintptr_t >( self.storage_ );
        const uintptr_t current = base + self.used_;
        if( current > UINTPTR_MAX - ( alignment - 1U ) )
        {
            ++self.failed_allocations_;
            return nullptr;
        }
        const uintptr_t aligned =
            ( current + alignment - 1U ) & ~( alignment - 1U );
        const size_t padding = static_cast< size_t >( aligned - current );
        const size_t remaining = self.capacity_ - self.used_;
        if( padding > remaining || bytes > remaining - padding )
        {
            ++self.failed_allocations_;
            return nullptr;
        }

        self.used_ += padding + bytes;
        ++self.successful_allocations_;
        return reinterpret_cast< void * >( aligned );
    }

    static void deallocate( void *, void *, size_t, size_t ) noexcept
    {
        /* Monotonic arena: individual deallocation deliberately does nothing. */
    }

    uint8_t * storage_;
    size_t capacity_;
    size_t maximum_alignment_;
    size_t used_ {};
    size_t successful_allocations_ {};
    size_t failed_allocations_ {};
    size_t resets_ {};
};

} // namespace freertos

#endif