#include <stdint.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <magenta/vmo.h>
#include <magenta/bootstrap.h>
#include <magenta/errors.h>
#include <magenta/misc.h>
#include <magenta/vmar.h>

/* 1 GiB */
#define HEAP_REGION_SZ 0x40000000

#define ROUND_UP(x, b) x += (b) - ((x) % (b))

static uintptr_t heap_start = 0;
static uintptr_t heap_end = 0;
static uintptr_t heap_alloc_point = 0;
static size_t heap_sz = 0;
static mx_handle_t heap_vmo = MX_NULL_HANDLE;
static mx_handle_t heap_vmar = MX_NULL_HANDLE;

void __crt_heap_init(size_t sz)
{
    mx_handle_t root_vmar = mx_bootstrap_handle_get(MX_B_VMAR_ROOT);
    mx_vaddr_t heap = 0;

    mx_status_t status = mx_vmar_allocate(root_vmar,
            MX_VM_CAN_MAP_READ | MX_VM_CAN_MAP_WRITE | MX_VM_CAN_MAP_SPECIFIC,
            0, HEAP_REGION_SZ, &heap_vmar, &heap);
    if (status != MX_OK) {
        fprintf(stderr, "fatal: cannot allocate heap virtual region (%s)\n", mx_status_to_string(status));
        abort();
    }

    status = mx_vmo_create(
            sz,
            MX_VM_CAN_MAP_READ | MX_VM_CAN_MAP_WRITE | MX_VM_CAN_MAP_SPECIFIC,
            &heap_vmo);
    if (status != MX_OK) {
        fprintf(stderr, "fatal: cannot allocate heap (%s)\n", mx_status_to_string(status));
        abort();
    }

    status = mx_vmar_map(heap_vmar,
            MX_VM_PERM_READ | MX_VM_PERM_WRITE | MX_VM_SPECIFIC,
            0, heap_vmo, 0, sz,
            &heap);
    if (status != MX_OK) {
        fprintf(stderr, "fatal: cannot map heap (%s)\n", mx_status_to_string(status));
        abort();
    }

    heap_sz = sz;
    heap_start = heap;
    heap_end = heap + heap_sz;
    heap_alloc_point = heap_start;

    memset((void *)heap, 0x0, heap_sz);
}

void *__crt_heap_extend(size_t sz)
{
    if (sz % MX_PAGE_SIZE) {
        sz &= (MX_PAGE_SIZE - 1);
        sz += MX_PAGE_SIZE;
    }

    if (!sz) {
        return (void *)heap_end;
    }

    if (!heap_start) {
        __crt_heap_init(sz);
        return (void *)heap_start;
    }

    mx_vmo_set_size(heap_vmo, heap_sz + sz);

    mx_vaddr_t vmar_base, alloc_base;
    mx_vmar_bounds(heap_vmar, &vmar_base, NULL);
    size_t offset = heap_end - vmar_base;

    mx_status_t err = mx_vmar_map(heap_vmar,
            MX_VM_PERM_READ | MX_VM_PERM_WRITE | MX_VM_SPECIFIC,
            offset, heap_vmo, heap_sz, sz, &alloc_base);

    if (err != MX_OK) {
        fprintf(stderr, "fatal: cannot map extended heap (%s)\n", mx_status_to_string(err));
        abort();
    }

    heap_sz += sz;

    void *out = (void *)heap_end;
    memset(out, 0x0, sz);
    heap_end += sz;

    return out;
}