#include "arena.h"

#include <string.h>

#if COLLA_DEBUG
#include <stdlib.h>
#endif

#include "os.h"

static uptr arena__align(uptr ptr, usize align) {
    return (ptr + (align - 1)) & ~(align - 1);
}

static arena_t arena__make_virtual(usize size);
static arena_t arena__make_malloc(usize size);
static arena_t arena__make_static(u8 *buf, usize len);

static void *arena__alloc_common(const arena_alloc_desc_t *desc);
static void *arena__alloc_malloc_always(const arena_alloc_desc_t *desc);

static void arena__free_virtual(arena_t *arena);
static void arena__free_malloc(arena_t *arena);

arena_t malloc_arena = {
    .type = ARENA_MALLOC_ALWAYS,
};

arena_t arena_init(const arena_desc_t *desc) {
    arena_t out = {0};

#if COLLA_DEBUG
    out.file = desc->file;
    out.line = desc->line;
#endif
    
    if (desc) {
        switch (desc->type) {
            case ARENA_VIRTUAL:       out = arena__make_virtual(desc->size); break;
            case ARENA_MALLOC:        out = arena__make_malloc(desc->size); break;
            case ARENA_STATIC:        out = arena__make_static(desc->static_buffer, desc->size); break;
            case ARENA_MALLOC_ALWAYS: out = malloc_arena; break;
		    default: break;  
        }
    }

    return out;
}

void arena_cleanup(arena_t *arena) {
    if (!arena) {
        return;
    }
    
    switch (arena->type) {
        case ARENA_VIRTUAL: arena__free_virtual(arena); break;
        case ARENA_MALLOC:  arena__free_malloc(arena);  break;
        // ARENA_STATIC does not need to be freed
        default: break;  
    }
    
    memset(arena, 0, sizeof(arena_t));
}

arena_t arena_scratch(arena_t *arena, usize size) {
    u8 *buffer = alloc(arena, u8, size, ALLOC_SOFT_FAIL | ALLOC_NOZERO);
    return arena__make_static(buffer, buffer ? size : 0);
}

void *arena_alloc(const arena_alloc_desc_t *desc) {
    if (!desc || !desc->arena || desc->arena->type == ARENA_TYPE_NONE) {
        return NULL;
    }

    arena_t *arena = desc->arena;

    u8 *ptr = NULL;

    switch (arena->type) {
        case ARENA_MALLOC_ALWAYS:
            ptr = arena__alloc_malloc_always(desc);
            break;
        default:
            ptr = arena__alloc_common(desc);
            break; 
    }

    if (!ptr && desc->flags & ALLOC_SOFT_FAIL) {
        return NULL;
    }

    usize total = desc->size * desc->count;

    return desc->flags & ALLOC_NOZERO ? ptr : memset(ptr, 0, total);
}

usize arena_tell(arena_t *arena) {
    return arena ? arena->cur - arena->beg : 0;
}

usize arena_remaining(arena_t *arena) {
    return arena && (arena->cur < arena->end) ? arena->end - arena->cur : 0;
}

usize arena_capacity(arena_t *arena) {
	return arena ? arena->end - arena->beg : 0;
}

void arena_rewind(arena_t *arena, usize from_start) {
    if (!arena) {
        return;
    }

    colla_assert(arena_tell(arena) >= from_start);

    arena->cur = arena->beg + from_start;
}

void arena_pop(arena_t *arena, usize amount) {
    if (!arena) {
        return;
    }
    usize position = arena_tell(arena);
    if (!position) {
        return;
    }
    arena_rewind(arena, position - amount);
}

// == VIRTUAL ARENA ====================================================================================================

static arena_t arena__make_virtual(usize size) {
    usize alloc_size = 0;
    u8 *ptr = os_reserve(size, &alloc_size);
    if (!os_commit(ptr, 1)) {
        os_release(ptr, alloc_size);
        ptr = NULL;
    }

    return (arena_t){
        .beg = ptr,
        .cur = ptr,
        .end = ptr ? ptr + alloc_size : NULL,
        .type = ARENA_VIRTUAL,
    };
}

static void arena__free_virtual(arena_t *arena) {
    if (!arena->beg) {
        return;
    }

    os_release(arena->beg, arena_capacity(arena));
}
// == MALLOC ARENA =====================================================================================================

static arena_t arena__make_malloc(usize size) {
    u8 *ptr = os_alloc(size);
    colla_assert(ptr);
    return (arena_t) {
        .beg = ptr,
        .cur = ptr,
        .end = ptr ? ptr + size : NULL,
        .type = ARENA_MALLOC,
    };
}

static void arena__free_malloc(arena_t *arena) {
    os_free(arena->beg);
}

// == ARENA ALLOC ======================================================================================================

static void *arena__alloc_common(const arena_alloc_desc_t *desc) {
    usize total = desc->size * desc->count;
    arena_t *arena = desc->arena;

    arena->cur = (u8 *)arena__align((uptr)arena->cur, desc->align);
    bool soft_fail = desc->flags & ALLOC_SOFT_FAIL;

    if (total > arena_remaining(arena)) {
        if (!soft_fail) {
#if COLLA_DEBUG
            arena__print_crash(arena);
#endif
            fatal("finished space in arena, tried to allocate %_$$$dB out of %_$$$dB (total: %_$$$dB)\n", total, arena_remaining(arena), (usize)(arena->end - arena->beg));
        }
        return NULL;
    }

    if (arena->type == ARENA_VIRTUAL) {
        usize allocated = arena_tell(arena);
        usize page_end = os_pad_to_page(allocated);
        usize new_cur = allocated + total;

        if (new_cur > page_end) {
            usize extra_mem = os_pad_to_page(new_cur - page_end);
            usize page_size = os_get_system_info().page_size;
            // TODO is this really correct?
            usize num_of_pages = (extra_mem / page_size) + 1;

            colla_assert(num_of_pages > 0);

            if (!os_commit(arena->cur, num_of_pages + 1)) {
                if (!soft_fail) {
#if COLLA_DEBUG
                    arena__print_crash(arena);
#endif
                    fatal("failed to commit memory for virtual arena, tried to commit %zu pages\n", num_of_pages);
                }
                return NULL;
            }
        }
    }

    u8 *ptr = arena->cur;
    arena->cur += total;

#if COLLA_DEBUG
    alloc_header_t *header = calloc(1, sizeof(alloc_header_t));
    header->size = desc->size;
    header->count = desc->count;
    memcpy(header->type_name, desc->type_name.buf, MIN(desc->type_name.len, sizeof(header->type_name)));
    list_push(arena->head, header);
#endif

    return ptr;
}

static void *arena__alloc_malloc_always(const arena_alloc_desc_t *desc) {
    usize total = desc->size * desc->count;

    // TODO: alignment?
    u8 *ptr = os_alloc(total);
    if (!ptr && !(desc->flags & ALLOC_SOFT_FAIL)) {
        fatal("alloc call failed for %_$$$dB", total);
    }

    return ptr;
}


// == STATIC ARENA =====================================================================================================

static arena_t arena__make_static(u8 *buf, usize len) {
    return (arena_t) {
        .beg = buf,
        .cur = buf,
        .end = buf ? buf + len : NULL,
        .type = ARENA_STATIC,
    };
}

// == DEBUG ============================================================================================================
#if COLLA_DEBUG

#define ARENA_HMAP_SIZE (1024)

typedef struct arena_hnode_t arena_hnode_t;
struct arena_hnode_t {
    u32 size;
    u32 count;
};

typedef struct arena_hmap_t arena_hmap_t;
struct arena_hmap_t {
    u64 hashes[ARENA_HMAP_SIZE];
    char keys[ARENA_HMAP_SIZE][16];
    arena_hnode_t values[ARENA_HMAP_SIZE];
};

u64 arena_hmap_hash(char key[16]) {
    const u8 *data = (const u8 *)key;
    u64 hash = 0;

    for (usize i = 0; i < 16; ++i) {
		hash = data[i] + (hash << 6) + (hash << 16) - hash;
    }

    return hash;
}

arena_hnode_t *arena_hmap_find_or_add(arena_hmap_t *map, char key[16]) {
    u64 hash = arena_hmap_hash(key);
    usize index = hash & (ARENA_HMAP_SIZE - 1);

    usize beg = index;
    usize end = ARENA_HMAP_SIZE;
    
    for (usize k = 0; k < 2; ++k) {
        for (usize i = index; i < ARENA_HMAP_SIZE; ++i) {
            if (map->hashes[i] == 0){
                arena_hnode_t *node = &map->values[i];
                map->hashes[i] = hash;
                memcpy(map->keys[i], key, 16);
                return node;
            }

            if (map->hashes[i] == hash && memcmp(map->keys[i], key, 16) == 0) {
                return &map->values[i];
            }
        }

        beg = 0;
        end = index;
    }

    return NULL;
}

void arena__print_crash(arena_t *arena) {
    arena_hmap_t hmap = {0};

    debug("arena %v:%d", arena->file, arena->line);

    for_each (header, arena->head) {
        arena_hnode_t *node = arena_hmap_find_or_add(&hmap, header->type_name);
        colla_assert(node->size == 0 || node->size == header->size);
        node->size = header->size;
        node->count += header->count;
    }

    print("type name       | size\t| count\t| total\n");
    for (usize i = 0; i < ARENA_HMAP_SIZE; ++i) {
        if (hmap.hashes[i] == 0) continue;
        arena_hnode_t n = hmap.values[i];
        print("%16s| %_$$$dB\t| %d\t| %_$$$dB\n", hmap.keys[i], n.size, n.count, n.size * n.count);
    }
}
#endif