#include <stdlib.h>

#include <stddef.h>
#include "os/tty.hpp"

namespace
{
inline const size_t INIT_MALLOC_SPACE_ELEMENTS = 1024 * 1024;
inline const size_t INIT_MALLOC_SPACE_BYTES = INIT_MALLOC_SPACE_ELEMENTS * alignof(max_align_t);

max_align_t gInitMallocSpace[INIT_MALLOC_SPACE_ELEMENTS];

struct MallocBlock
{
    size_t elements;
    MallocBlock* nextBlock;
};
struct AllocInfo
{
    size_t elements;
};
static_assert(sizeof(MallocBlock) <= alignof(max_align_t));
static_assert(sizeof(AllocInfo) <= alignof(max_align_t));

MallocBlock* gNextBlock = []()
{
    MallocBlock* initialBlock = reinterpret_cast<MallocBlock*>(gInitMallocSpace);
    initialBlock->elements = INIT_MALLOC_SPACE_ELEMENTS;
    initialBlock->nextBlock = nullptr;
    return initialBlock;
}();
}

void __ba_registerAllocatableMemory(void* memory, size_t size) noexcept
{
    if (size < sizeof(max_align_t))
    {
        return;
    }
    MallocBlock* newBlock = static_cast<MallocBlock*>(memory);
    newBlock->nextBlock = gNextBlock;
    newBlock->elements = size / sizeof(max_align_t);
    gNextBlock = newBlock;
}

extern "C"
{
void abort()
{
    tty::write("Abort called!");
    __asm__ volatile("hlt");
    while(true) {};
    __builtin_unreachable();
}

void* malloc(size_t size) noexcept
{
    if (size == 0)
    {
        return nullptr;
    }

    const size_t requiredElements = (size + sizeof(max_align_t) - 1) / sizeof(max_align_t) + 1;
    MallocBlock* prevBlock = nullptr;
    for (MallocBlock* block = gNextBlock; block != nullptr; block = block->nextBlock)
    {
        if (block->elements >= requiredElements)
        {
            MallocBlock* newBlock = nullptr;
            if (block->elements > requiredElements)
            {
                newBlock = reinterpret_cast<MallocBlock*>(reinterpret_cast<max_align_t*>(block) + requiredElements);
                newBlock->nextBlock = block->nextBlock;
                newBlock->elements = block->elements - requiredElements;
            }
            else
            {
                newBlock = block->nextBlock;
            }
            if (prevBlock != nullptr)
            {
                prevBlock->nextBlock = newBlock;
            }
            else
            {
                gNextBlock = newBlock;
            }
            AllocInfo* allocInfo = reinterpret_cast<AllocInfo*>(block);
            allocInfo->elements = requiredElements;
            return reinterpret_cast<max_align_t*>(block) + 1;
        }
        prevBlock = block;
    }

    return nullptr;
}

void free(void* memory) noexcept
{
    if (memory == nullptr)
    {
        return;
    }
    MallocBlock* block = reinterpret_cast<MallocBlock*>(static_cast<max_align_t*>(memory) - 1);
    block->nextBlock = gNextBlock;
    gNextBlock = block;
}

int __cxa_atexit(void (*func)(void*), void* arg, void* dsoHandle)
{
    (void) func;
    (void) arg;
    (void) dsoHandle;

    // NOOP for now
    // TODO
    return 0;
}
} // extern "C"