#include <stdlib.h>

#define BLOCK_SIZE (64 * 1024)
#define BLOCK_1MB 16

typedef struct header {
    size_t size : 31;
    uint8_t is_free : 1;
} header_t;
static header_t* head;

static uintptr_t bump_pointer = (uintptr_t)&__heap_base;
static uintptr_t heap_end = BLOCK_1MB * BLOCK_SIZE;
IMPORT(grow) void grow(uint32_t blocks);

static void* bump_alloc(uintptr_t n) {
    uintptr_t r = bump_pointer;
    bump_pointer += n;

    while (bump_pointer >= heap_end) {
        grow(heap_end / BLOCK_SIZE);
        heap_end *= 2;
    }

    return (void*)r;
}

static header_t* get_next_header(header_t* block) {
    if (!block) {
        return NULL;
    }
    uint8_t* data = (uint8_t*)(block + 1);
    uint8_t* next_header = data + block->size;
    if ((uintptr_t)next_header >= bump_pointer) {
        return NULL;
    }
    return (header_t*)next_header;
}

static header_t* get_prev_header(header_t* block) {
    header_t* curr = head;
    while (curr) {
        header_t* next = get_next_header(curr);
        if (next == block) {
            return curr;
        }
        curr = next;
    }
    return NULL;
}

static header_t* get_free_block(size_t size) {
    header_t* curr = head;
    while (curr) {
        if (curr->is_free && curr->size >= size) {
            return curr;
        }
        curr = get_next_header(curr);
    }
    return NULL;
}

void* malloc(size_t size) {
    if (!size) {
        return NULL;
    }

    header_t* header = get_free_block(size);
    if (header) {
        header->is_free = 0;
        if (header->size > size + sizeof(header_t) * 2) {
            size_t old_size = header->size;
            header->size = size;

            header_t* next = get_next_header(header);
            next->size = old_size - size - sizeof(header_t);
            next->is_free = 1;
        }
        return (void*)(header + 1);
    }
    size_t total_size = sizeof(header_t) + size;

    void* block = bump_alloc(total_size);
    if (!block) {
        return NULL;
    }

    header = block;
    header->size = size;
    header->is_free = 0;

    if (!head) {
        head = header;
    }

    return (void*)(header + 1);
}

void free(void* block) {
    if (!block) {
        return;
    }

    header_t* header = ((header_t*)block) - 1;
    if (header->is_free) {
        return;
    }
    header->is_free = 1;

    header_t* next_header = get_next_header(header);
    if (next_header && next_header->is_free) {
        header->size += next_header->size + sizeof(header_t);
    }

    header_t* prev_header = get_prev_header(header);
    if (prev_header && prev_header->is_free) {
        prev_header->size += header->size + sizeof(header_t);
    }
}

void* realloc(void* block, size_t size) {
    header_t* header;
    void* ret;

    if (!block || !size)
        return malloc(size);

    header = (header_t*)block - 1;
    if (header->size >= size)
        return block;

    ret = malloc(size);
    if (ret) {
        memcpy(ret, block, header->size);
        free(block);
    }
    return ret;
}

void* memset(void* s, uint8_t c, uint32_t n) {
    uint8_t* p = (uint8_t*)s;
    while (n--) {
        *p++ = c;
    }
    return s;
}

void* memcpy(void* dest, const void* src, uint32_t n) {
    uint8_t* d = (uint8_t*)dest;
    const uint8_t* s = (const uint8_t*)src;
    while (n--) {
        *d++ = *s++;
    }
    return dest;
}

int memcmp(const void* s1, const void* s2, uint32_t n) {
    const uint8_t* p1 = (const uint8_t*)s1;
    const uint8_t* p2 = (const uint8_t*)s2;

    while (n--) {
        if (*p1 != *p2) {
            return (*p1 - *p2);
        }
        p1++;
        p2++;
    }

    return 0; // Memory blocks are equal
}

static uint64_t rand_state;

void srand(uint64_t seed) { rand_state = seed; }

uint64_t rand(void) {
    uint64_t z = (rand_state += 0x9e3779b97f4a7c15);
    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9;
    z = (z ^ (z >> 27)) * 0x94d049bb133111eb;
    return z ^ (z >> 31);
}