#include <foundation/types.h>
#include <foundation/malloc.h>
#include <foundation/platform.h>

//------------------------------------------------------------------------------
struct Header
{
    enum : uint32 { Magic = 0xaa55aa55 };
    size_t  size;
    uint32  bias;
    uint32  canary;
    uint8   ptr[];
};
static bool g_add_canaries = false;

//------------------------------------------------------------------------------
#ifdef _WIN32

template <typename T=uint8>
static T* tt_malloc_norm(size_t size, size_t alignment)
{
    alignment = uint32(std::max(alignment, alignof(std::max_align_t)));
    if (!g_add_canaries)
        return (T*)_aligned_malloc(size, alignment);

    uint32 lead_size = uint32(sizeof(Header) + alignment - 1) & ~uint32(alignment - 1);
    size_t alloc_size = lead_size + size + sizeof(uint32);
    auto* alloc_ptr = (uint8*)_aligned_malloc(alloc_size, alignment);

    auto* header = (Header*)(alloc_ptr + lead_size - sizeof(Header));
    header->bias = lead_size;
    header->size = size;
    header->canary = Header::Magic;

    std::memcpy(header->ptr + size, &header->canary, sizeof(uint32));

    return header->ptr;
}

static void tt_free_norm(void* address)
{
    if (address == nullptr)
        return;

    if (!g_add_canaries)
        return _aligned_free(address);

    auto* header = (Header*)address - 1;

    auto is_canary = [] (void* ptr) {
        uint32 canary = Header::Magic;
        return !std::memcmp(ptr, &canary, sizeof(uint32));
    };
    if (!is_canary(&header->canary))            __debugbreak();
    if (!is_canary(header->ptr + header->size)) __debugbreak();

    _aligned_free(header->ptr - header->bias);
}

#else // POSIX

template <typename T=uint8>
static T* tt_malloc_norm(size_t size, size_t alignment)
{
    alignment = std::max(alignment, alignof(std::max_align_t));
    // std::aligned_alloc requires size to be a multiple of alignment
    size_t aligned_size = (size + alignment - 1) & ~(alignment - 1);
    if (aligned_size == 0)
        aligned_size = alignment;

    if (!g_add_canaries)
        return (T*)std::aligned_alloc(alignment, aligned_size);

    uint32 lead_size = uint32(sizeof(Header) + alignment - 1) & ~uint32(alignment - 1);
    size_t alloc_size = lead_size + size + sizeof(uint32);
    size_t alloc_aligned = (alloc_size + alignment - 1) & ~(alignment - 1);
    auto* alloc_ptr = (uint8*)std::aligned_alloc(alignment, alloc_aligned);

    auto* header = (Header*)(alloc_ptr + lead_size - sizeof(Header));
    header->bias = lead_size;
    header->size = size;
    header->canary = Header::Magic;

    std::memcpy(header->ptr + size, &header->canary, sizeof(uint32));

    return header->ptr;
}

static void tt_free_norm(void* address)
{
    if (address == nullptr)
        return;

    if (!g_add_canaries)
        return std::free(address);

    auto* header = (Header*)address - 1;

    auto is_canary = [] (void* ptr) {
        uint32 canary = Header::Magic;
        return !std::memcmp(ptr, &canary, sizeof(uint32));
    };
    if (!is_canary(&header->canary))            __builtin_trap();
    if (!is_canary(header->ptr + header->size)) __builtin_trap();

    std::free(header->ptr - header->bias);
}

#endif

//------------------------------------------------------------------------------
void tt_memory_canaries_on()
{
    g_add_canaries = true;
}


//------------------------------------------------------------------------------
void*    (*tt_malloc_impl)(size_t, size_t)  = tt_malloc_norm;
void     (*tt_free)(void*)                  = tt_free_norm;



//------------------------------------------------------------------------------
#ifdef _WIN32

template <typename T=uint8>
static T* tt_malloc_stomp(size_t size, size_t alignment)
{
    enum { PAGE_SIZE = 64 << 10 };
    size = (size + (alignment - 1)) & ~(alignment - 1);
    size_t alloc_size = (size + (2 * PAGE_SIZE)) & ~(PAGE_SIZE - 1);
    auto* ptr = (uint8*)VirtualAlloc(nullptr, alloc_size, MEM_COMMIT, PAGE_READWRITE);
    ptr += alloc_size - PAGE_SIZE;
    DWORD unused;
    VirtualProtect(ptr, PAGE_SIZE, PAGE_NOACCESS, &unused);

    return (T*)(ptr - size);
}

static void tt_free_stomp(void* address)
{
    if (address == nullptr)
        return;

    *(uint32*)address = 0x30493049;

    MEMORY_BASIC_INFORMATION info;
    VirtualQuery(address, &info, sizeof(info));
    DWORD unused;
    VirtualProtect(info.AllocationBase, info.RegionSize, PAGE_NOACCESS, &unused);
}

#else // POSIX

static size_t tt_get_page_size()
{
    static size_t page_size = size_t(sysconf(_SC_PAGESIZE));
    return page_size;
}

struct StompHeader
{
    void*  base;
    size_t total_size;
};

template <typename T=uint8>
static T* tt_malloc_stomp(size_t size, size_t alignment)
{
    size_t page_size = tt_get_page_size();
    size = (size + (alignment - 1)) & ~(alignment - 1);
    size_t needed = sizeof(StompHeader) + size + page_size;
    size_t alloc_size = (needed + (page_size - 1)) & ~(page_size - 1);
    auto* base = (uint8*)mmap(nullptr, alloc_size, PROT_READ | PROT_WRITE,
                              MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (base == MAP_FAILED)
        return nullptr;

    // Guard page at the end
    uint8* guard = base + alloc_size - page_size;
    mprotect(guard, page_size, PROT_NONE);

    // User allocation right before the guard, header right before that
    T* user_ptr = (T*)(guard - size);
    auto* hdr = (StompHeader*)((uint8*)user_ptr - sizeof(StompHeader));
    hdr->base = base;
    hdr->total_size = alloc_size;

    return user_ptr;
}

static void tt_free_stomp(void* address)
{
    if (address == nullptr)
        return;

    *(uint32*)address = 0x30493049;

    auto* hdr = (StompHeader*)((uint8*)address - sizeof(StompHeader));
    mprotect(hdr->base, hdr->total_size, PROT_NONE);
}

#endif

//------------------------------------------------------------------------------
void tt_memory_stomp_on()
{
    tt_malloc_impl = tt_malloc_stomp;
    tt_free = tt_free_stomp;
}