#ifndef __UH_ALLOCATOR
#define __UH_ALLOCATOR

#include <windows.h>

#include "std.h"
#include "errors.h"
#include "synchronization.h"
#include "timers.h"

UHEXTERN_C_BEGIN

#ifndef UH_USE_ASAN
#define UH_USE_ASAN 0
#endif

#if UH_USE_ASAN
void __asan_poison_memory_region(void const volatile* addr, usize size);
void __asan_unpoison_memory_region(void const volatile* addr, usize size);

#define UH_ASAN_POISON(addr, size) __asan_poison_memory_region((addr), (size))
#define UH_ASAN_UNPOISON(addr, size) __asan_unpoison_memory_region((addr), (size))
#else
#define UH_ASAN_POISON(addr, size) ((void) 0)
#define UH_ASAN_UNPOISON(addr, size) ((void) 0)
#endif

// it uses a linear allocation strategy. when the initial 8k 
// page is used up, it allocates new 64k pages by default from the global allocator, or exact size pages for new allocations

// binned allocations in pages for individual blocks of allocations 
typedef struct uha_block_bin {
    void** pages; // 64K Pages
    i32 page_count; // 64K Pages
    i32 page_cap; // 64K Pages
    i32 block_size; // 8, 16, 64, 512, 2048, 8192, 16 * 1024, 32 * 1024

    i32 cap;
    i32 next;

    i32* free_list;
    i32 free_count;
} uha_block_bin;

typedef struct uh_page_header {
    i32 kind;
    i32 bin_index;
    u32 payload_offset;
    u32 rsvd;
} uh_page_header;

typedef struct uh_large_page_node uh_large_page_node;

typedef struct uh_large_page_node {
    uh_large_page_node* next;
    uh_large_page_node* prev;
    void* alloc_base;
    void* page;
    i64 size;
    i64 free_time;
} uh_large_page_node;


#define UH_WPSIZEINIT 8 * 1024
#define UH_BINPAGESIZE 64 * 1024

#define UH_SLABMAX 32752
#define UH_ALIGN16(x) (((x) + 15) & ~((usize) 15))
#define UH_PAGE_KIND_SLAB 0
#define UH_PAGE_KIND_LARGE -1
#define UH_SLAB_OFFSET 16
#define UH_LARGE_TRIM_NS 1000000000LL

static const usize uh_bins_size_list[] = {
    8, 16, 64, 512, 
    2 * 1024, 

    // these 3 will use tiled blocks strategy
    8176,
    16368,
    UH_SLABMAX,
};
#define UH_BINSCOUNT sizeof(uh_bins_size_list) / sizeof(uh_bins_size_list[0])

// allocations up to 8KB in size.
// small arenas, we do not pre-create any bins until we hit an allocation
// but they are pre-slotted
typedef struct uh_allocator {
    uh_mutex mutex;
    void* work_page;     // 4k small work page for all internal memory allocation functions
    i32 wp_size;
    i32 wp_allocator;
                    
    uha_block_bin* bins;
    i32 bin_count;
    
    void** large;
    uh_large_page_node* lpn;

    uh_large_page_node* lpn_free;
    uh_large_page_node* lpn_free_tail;
    
    i64 commit_amount; // bytes requested through VirtualAlloc
    i64 alloc_amount; // bytes handed out to subAllocations
    i64 alloc_count; // bytes handed out to subAllocations
                    // slab overhead = commit_amount - alloc_amount
} uh_allocator;


static b8 uh_allocator_init(uh_allocator* allocator);
static void uh_allocator_deinit(uh_allocator* a);

static void uh_allocator_trim_large_pages(uh_allocator* a, i64 now);
static void uha_free_unlocked(uh_allocator* a, void* ptr);
static void* uha_alloc_unlocked(uh_allocator* a, usize size);

extern uh_allocator* g_allocator;

uh_allocator* g_allocator;

static uh_allocator* uh_get_global_allocator(void)
{
    return g_allocator;
}

static b8 uh_init_global_allocator(void)
{
    g_allocator = VirtualAlloc(NULL, sizeof(uh_allocator), MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    UH_ASAN_UNPOISON(g_allocator, sizeof(uh_allocator));
    uh_logf("arena allocator init.");
    return uh_allocator_init(uh_get_global_allocator());
}

static void uh_deinit_global_allocator(void)
{
    uh_allocator_deinit(uh_get_global_allocator());
    VirtualFree(uh_get_global_allocator(), 0, MEM_RELEASE);
}

static void uh_allocator_deinit(uh_allocator* a)
{
    uh_mutex_lock(&a->mutex);

    if(a->commit_amount || a->alloc_amount || a->alloc_count)
    {
        const i64 overhead_amount = a->commit_amount - a->alloc_amount;

        uh_logf("==== LEAKED MEMORY report ====");
        uh_logf("live alloc_amount=%d", a->alloc_amount);
        uh_logf("live alloc_count=%d", a->alloc_count);
        uh_logf("committed bytes=%d", a->commit_amount);
        uh_logf("allocator overhead=%d", overhead_amount);
    }

    for (i32 i = 0; i < a->bin_count; i += 1)
    {
        uha_block_bin* bin = &a->bins[i];

        for(i32 j = 0; j < bin->page_count; j += 1)
        {
            VirtualFree(bin->pages[j], 0, MEM_RELEASE);
        }
        if (bin->free_list != NULL)
        {
            VirtualFree(bin->free_list, 0, MEM_RELEASE);
            bin->free_list = NULL;
            bin->free_count = 0;
        }
    }


    uh_large_page_node* lpn_dead = a->lpn;
    while(lpn_dead != NULL)
    {
        uh_large_page_node* lpn_next = lpn_dead->next;
        VirtualFree(lpn_dead->alloc_base, 0, MEM_RELEASE);
        lpn_dead = lpn_next;
    }

    lpn_dead = a->lpn_free;
    while(lpn_dead != NULL)
    {
        uh_large_page_node* lpn_next = lpn_dead->next;
        VirtualFree(lpn_dead->alloc_base, 0, MEM_RELEASE);
        lpn_dead = lpn_next;
    }

    VirtualFree(a->work_page, 0, MEM_RELEASE);
    uh_mutex_unlock(&a->mutex);
    uh_mutex_deinit(&a->mutex);
}

// allocate from the workpage linear allocator
static void* uha_wp_alloc(uh_allocator* a, i64 size)
{
    void* p = (void*) ((usize) a->work_page + a->wp_allocator);
    a->wp_allocator += size;

    if (a->wp_allocator > a->wp_size)
    {
        uh_panic("Workpage Size exceeded, adding additional workpages... (not implemented yet)");
    }

    return p;
}

static b8 uh_allocator_init(uh_allocator* a)
{
    uh_mutex_init(&a->mutex);
    a->work_page = VirtualAlloc(NULL, UH_WPSIZEINIT, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    UH_ASAN_UNPOISON(a->work_page, UH_WPSIZEINIT);
    a->wp_size = UH_WPSIZEINIT;
    a->wp_allocator = 0;
    
    a->lpn = NULL;
    a->lpn_free = NULL;
    a->lpn_free_tail = NULL;

    a->commit_amount = UH_WPSIZEINIT;
    a->alloc_amount = 0;
    a->alloc_count = 0;

    // initialize bins
    a->bin_count = (i32) UH_BINSCOUNT;
    a->bins = (uha_block_bin*) uha_wp_alloc(a, sizeof(uha_block_bin) * a->bin_count);

    // bins initialization
    for (i32 i = 0; i < a->bin_count; i += 1)
    {
        uha_block_bin* bin = a->bins + i;

        bin->pages = (void**) uha_wp_alloc(a, sizeof(void*) * 32);
        bin->page_cap = 32;
        bin->page_count = 0;

        bin->block_size = uh_bins_size_list[i];

        bin->cap = 0;
        bin->next = 0;

        bin->free_list = NULL;
        bin->free_count = 0;
    }

    return 1;
}

static void uh_allocator_trim_large_pages(uh_allocator* a, i64 now)
{
    while (a->lpn_free_tail != NULL)
    {
        uh_large_page_node* node = a->lpn_free_tail;

        if (now - node->free_time <= UH_LARGE_TRIM_NS)
        {
            break;
        }

        a->lpn_free_tail = node->prev;
        if (a->lpn_free_tail != NULL)
        {
            a->lpn_free_tail->next = NULL;
        }
        else
        {
            a->lpn_free = NULL;
        }

        a->commit_amount -= (i64) (((usize) node->page - (usize) node->alloc_base) + (usize) node->size);
        VirtualFree(node->alloc_base, 0, MEM_RELEASE);
    }
}

static void uha_free_unlocked(uh_allocator* a, void* ptr)
{
    uh_page_header* header;

    if (ptr == NULL)
    {
        return;
    }

    header = (uh_page_header*) ((usize) ptr & ~((usize) UH_BINPAGESIZE - 1));
    if (header->kind == UH_PAGE_KIND_LARGE)
    {
        uh_large_page_node* node = (uh_large_page_node*) ((usize) header + UH_ALIGN16(sizeof(uh_page_header)));

        if (node->prev != NULL)
        {
            node->prev->next = node->next;
        }
        else
        {
            a->lpn = node->next;
        }

        if (node->next != NULL)
        {
            node->next->prev = node->prev;
        }
        else
        {
            a->lpn_free_tail = node->prev;
        }

        node->prev = NULL;
        node->next = a->lpn_free;
        if (a->lpn_free != NULL)
        {
            a->lpn_free->prev = node;
        }
        else
        {
            a->lpn_free_tail = node;
        }
        a->lpn_free = node;
        node->free_time = uh_time();

        UH_ASAN_POISON(node->page, (usize) node->size);
        a->alloc_amount -= node->size;
        a->alloc_count -= 1;
        uh_allocator_trim_large_pages(a, node->free_time);
        return;
    }

    if (header->kind == UH_PAGE_KIND_SLAB)
    {
        uha_block_bin* bin = &a->bins[header->bin_index];
        const i32 slots_per_page = (UH_BINPAGESIZE - UH_SLAB_OFFSET) / bin->block_size;
        const usize offset = (usize) ptr - ((usize) header + UH_SLAB_OFFSET);
        const i32 slot_index = (i32) (offset / (usize) bin->block_size);
        const i32 page_index = (i32) header->rsvd;
        const usize free_list_size = sizeof(i32) * (usize) (slots_per_page * bin->page_cap);

        if (bin->free_list == NULL)
        {
            bin->free_list = (i32*) VirtualAlloc(NULL, free_list_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
            if (bin->free_list == NULL)
            {
                uh_panic("failed to alloc slab free_list");
            }
            UH_ASAN_UNPOISON(bin->free_list, free_list_size);
            a->commit_amount += (i64) free_list_size;
        }

        bin->free_list[bin->free_count] = page_index * slots_per_page + slot_index;
        bin->free_count += 1;

        UH_ASAN_POISON(ptr, (usize) bin->block_size);
        a->alloc_amount -= bin->block_size;
        a->alloc_count -= 1;
        return;
    }

    uh_panic("invalid page kind in uha_free");
}

static void uha_free(uh_allocator* a, void* ptr)
{
    uh_mutex_lock(&a->mutex);
    uha_free_unlocked(a, ptr);
    uh_mutex_unlock(&a->mutex);
}

static void* uha_alloc_linkedlist(uh_allocator* a, usize size)
{
    uh_page_header* header;
    uh_large_page_node* node;
    uh_large_page_node* free_node;
    usize node_offset;
    usize payload_offset;
    usize total_size;

    free_node = a->lpn_free;
    while (free_node != NULL)
    {
        if ((usize) free_node->size >= size)
        {
            if (free_node->prev != NULL)
            {
                free_node->prev->next = free_node->next;
            }
            else
            {
                a->lpn_free = free_node->next;
            }

            if (free_node->next != NULL)
            {
                free_node->next->prev = free_node->prev;
            }
            else
            {
                a->lpn_free_tail = free_node->prev;
            }

            free_node->prev = NULL;
            free_node->next = a->lpn;
            if (a->lpn != NULL)
            {
                a->lpn->prev = free_node;
            }
            a->lpn = free_node;
            free_node->size = (i64) size;

            UH_ASAN_UNPOISON(free_node->page, size);
            a->alloc_amount += (i64) size;
            a->alloc_count += 1;
            return free_node->page;
        }

        free_node = free_node->next;
    }

    node_offset = UH_ALIGN16(sizeof(uh_page_header));
    payload_offset = UH_ALIGN16(node_offset + sizeof(uh_large_page_node));
    total_size = payload_offset + size;
    header = (uh_page_header*) VirtualAlloc(NULL, total_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    if (header == NULL)
    {
        return NULL;
    }
    UH_ASAN_UNPOISON(header, total_size);

    header->kind = UH_PAGE_KIND_LARGE;
    header->bin_index = -1;
    header->payload_offset = (u32) payload_offset;
    header->rsvd = 0;

    node = (uh_large_page_node*) ((usize) header + node_offset);
    node->alloc_base = header;
    node->page = (void*) ((usize) header + payload_offset);
    node->size = (i64) size;
    node->free_time = 0;
    node->next = a->lpn;
    node->prev = NULL;
    if (a->lpn != NULL)
    {
        a->lpn->prev = node;
    }
    a->lpn = node;

    a->commit_amount += (i64) total_size;
    a->alloc_amount += (i64) size;
    a->alloc_count += 1;
    UH_ASAN_POISON(node->page, size);
    UH_ASAN_UNPOISON(node->page, size);

    return node->page;
}

static uha_block_bin* uha_ensure_block(uh_allocator* a, usize block_index)
{
    // for a given block index, ensure that there is a new page allocated if need be
    uha_block_bin* bin = &a->bins[block_index];
    const i32 slots_per_page = (UH_BINPAGESIZE - UH_SLAB_OFFSET) / bin->block_size;
    const i32 total_slots = bin->page_count * slots_per_page;

    if (!bin)
    {
        uh_panic("invalid bin setup");
    }

    // early out if something is available in the free_count
    if (bin->free_count > 0)
    {
        return bin;
    }

    if (bin->next >= total_slots)
    {
        // append a new page
        if(bin->page_count + 1 > bin->page_cap)
        {
            uh_panic("pages exceeded TODO, implement");
        }

        bin->pages[bin->page_count] = VirtualAlloc(NULL, UH_BINPAGESIZE, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
        if(!bin->pages[bin->page_count])
        {
            uh_panic("failed to alloc page");
        }
        UH_ASAN_UNPOISON(bin->pages[bin->page_count], UH_BINPAGESIZE);

        ((uh_page_header*) bin->pages[bin->page_count])->kind = UH_PAGE_KIND_SLAB;
        ((uh_page_header*) bin->pages[bin->page_count])->bin_index = (i32) block_index;
        ((uh_page_header*) bin->pages[bin->page_count])->payload_offset = UH_SLAB_OFFSET;
        ((uh_page_header*) bin->pages[bin->page_count])->rsvd = (u32) bin->page_count;
        UH_ASAN_POISON(
            (void*) ((usize) bin->pages[bin->page_count] + UH_SLAB_OFFSET),
            (usize) slots_per_page * (usize) bin->block_size);

        bin->page_count += 1;
        bin->cap += slots_per_page;
        a->commit_amount += UH_BINPAGESIZE;
    }

    return bin;
}

static void* uha_block_alloc(uh_allocator* a, uha_block_bin* bin)
{
    a->alloc_count += 1;
    a->alloc_amount += bin->block_size;
    i32 index = -1;

    if (bin->free_count > 0)
    {
        bin->free_count -= 1;
        index = bin->free_list[bin->free_count];
    }

    if (index == -1)
    {
        index = bin->next;
        bin->next += 1;
    }
    
    if (index == -1)
    {
        uh_panic("unable to assign index?");
        return NULL;
    }

    const i32 slots_per_page = (UH_BINPAGESIZE - UH_SLAB_OFFSET) / bin->block_size;
    const i32 page_index = index / slots_per_page;
    const i32 slot_index = index % slots_per_page;
    void* ptr = (void*) ((usize)(bin->pages[page_index]) + UH_SLAB_OFFSET + (slot_index * bin->block_size));

    UH_ASAN_UNPOISON(ptr, (usize) bin->block_size);

    return ptr;
}

static void* uha_alloc_unlocked(uh_allocator* a, usize size) 
{
    if (size > UH_SLABMAX)
    {
        // appends an allocation to the linkedlist
        return uha_alloc_linkedlist(a, size);
    }

    void* rv = NULL;

    for (i32 i = 0; i < (i32)(UH_BINSCOUNT); i += 1)
    {
        if(size <= uh_bins_size_list[i])
        {
            // bin found calculate slot index
            uha_block_bin* bin = uha_ensure_block(a, i);

            // grab the latest slot from the bin and return that 
            rv = uha_block_alloc(a, bin);
            break;
        }
    }

    return rv;
}

static void* uha_alloc(uh_allocator* a, usize size)
{
    void* rv;

    uh_mutex_lock(&a->mutex);
    rv = uha_alloc_unlocked(a, size);
    uh_mutex_unlock(&a->mutex);
    return rv;
}

// a linear allocator that uha can hand to you whenever.
// starts at 8k and hands out 64k pages after the first page is exhausted

#define UH_LA_PAGESIZE0 8176
#define UH_LA_PAGESIZE 64 * 1024
#define UH_LA_INITIAL_PAGECOUNT 8


typedef struct uh_linear_arena {
    uh_allocator* backing;
    void** pages;
    i32 page_count;
    i32 page_cap;
    i32 current_page;
    usize current_offset;

    uh_large_page_node* lpn;
    uh_large_page_node* free_list;

    i64 amount; // amount allocated, not amount in page
    i64 commit; // total bytes requested from the backing allocator

}uh_linear_arena;

static usize uh_linear_arena_page_size_for_index(i32 page_index)
{
    return page_index == 0 ? UH_LA_PAGESIZE0 : UH_LA_PAGESIZE;
}

static usize uh_linear_arena_align_forward(usize value, usize align)
{
    const usize mask = align - 1;
    return (value + mask) & ~mask;
}

static void uh_linear_arena_reset(uh_linear_arena* arena)
{
    arena->amount = 0;
    arena->current_page = 0;
    arena->current_offset = 0;

    if (arena->free_list != NULL)
    {
        uh_large_page_node* tail = arena->free_list;
        while (tail->next != NULL)
        {
            tail = tail->next;
        }
        tail->next = arena->lpn;
        if (arena->lpn != NULL)
        {
            arena->lpn->prev = tail;
        }
    }
    else
    {
        arena->free_list = arena->lpn;
    }

    if (arena->lpn != NULL)
    {
        arena->lpn->prev = NULL;
    }
    arena->lpn = NULL;
}

static b8 uh_linear_arena_push_page(uh_linear_arena* arena)
{
    if (arena->page_count + 1 > arena->page_cap)
    {
        i32 new_cap = arena->page_cap * 2;
        void** new_pages = uha_alloc(arena->backing, sizeof(void*) * new_cap);
        if (!new_pages)
        {
            return BAD;
        }

        for (i32 i = 0; i < arena->page_count; i += 1)
        {
            new_pages[i] = arena->pages[i];
        }

        uha_free(arena->backing, arena->pages);
        arena->pages = new_pages;
        arena->page_cap = new_cap;
    }

    usize page_size = uh_linear_arena_page_size_for_index(arena->page_count);

    void* page = uha_alloc(arena->backing, page_size);

    if (!page)
    {
        return BAD;
    }

    arena->pages[arena->page_count] = page;
    arena->page_count += 1;
    arena->commit += page_size;

    return OK;
}

static void* uh_linear_arena_alloc_large(uh_linear_arena* arena, usize size)
{
    uh_large_page_node* node = arena->free_list;

    while (node != NULL)
    {
        if ((usize) node->size >= size)
        {
            if (node->prev != NULL)
            {
                node->prev->next = node->next;
            }
            else
            {
                arena->free_list = node->next;
            }

            if (node->next != NULL)
            {
                node->next->prev = node->prev;
            }

            node->prev = NULL;
            node->next = arena->lpn;
            if (arena->lpn != NULL)
            {
                arena->lpn->prev = node;
            }
            arena->lpn = node;
            arena->amount += (i64) size;
            return node->page;
        }

        node = node->next;
    }

    node = (uh_large_page_node*) uha_alloc(arena->backing, sizeof(uh_large_page_node));
    if (node == NULL)
    {
        return NULL;
    }

    node->page = uha_alloc(arena->backing, size);
    if (node->page == NULL)
    {
        uha_free(arena->backing, node);
        return NULL;
    }

    node->alloc_base = node->page;
    node->size = (i64) size;
    node->prev = NULL;
    node->next = arena->lpn;
    if (arena->lpn != NULL)
    {
        arena->lpn->prev = node;
    }
    arena->lpn = node;

    arena->amount += (i64) size;
    arena->commit += (i64) size + (i64) sizeof(uh_large_page_node);
    return node->page;
}

static void* uh_linear_arena_alloc(uh_linear_arena* arena, usize size, usize align)
{
    usize offset;
    usize page_size;
    void* page;

    if (align == 0)
    {
        align = 1;
    }

    if ((align & (align - 1)) != 0)
    {
        uh_panic("uh_linear_arena_alloc requires power-of-two alignment");
    }

    if (size > UH_LA_PAGESIZE)
    {
        return uh_linear_arena_alloc_large(arena, size);
    }

    if (arena->page_count == 0)
    {
        if (!uh_linear_arena_push_page(arena))
        {
            return NULL;
        }
    }

    for (;;)
    {
        page_size = uh_linear_arena_page_size_for_index(arena->current_page);
        offset = uh_linear_arena_align_forward(arena->current_offset, align);

        if (offset + size <= page_size)
        {
            page = arena->pages[arena->current_page];
            arena->current_offset = offset + size;
            arena->amount += (i64) size;
            return (void*) ((usize) page + offset);
        }

        arena->current_page += 1;
        arena->current_offset = 0;

        if (arena->current_page >= arena->page_count)
        {
            if (!uh_linear_arena_push_page(arena))
            {
                return NULL;
            }
        }
    }
}

static b8 uh_new_arena(uh_allocator* backing, uh_linear_arena* arena)
{
    arena->backing = backing;
    arena->pages = uha_alloc(backing, UH_LA_INITIAL_PAGECOUNT * sizeof(void*)); // init 8 page slots first
    arena->page_cap = UH_LA_INITIAL_PAGECOUNT;

    arena->page_count = 0;
    arena->current_page = 0;
    arena->current_offset = 0;
    arena->amount = 0;
    arena->commit = 0;
    arena->lpn = NULL;
    arena->free_list = NULL;

    if(!arena->pages)
    {
        uh_panic("unable to init arena");
    }

    arena->commit += (i64) (UH_LA_INITIAL_PAGECOUNT * sizeof(void*));

    if (!uh_linear_arena_push_page(arena))
    {
        uh_panic("unable to push initial arena page");
    }

    return OK;
}

static void uh_destroy_arena(uh_linear_arena* arena)
{
    uh_allocator* a = arena->backing;
    i32 i;
    uh_large_page_node* lpn_dead;

    for (i = 0; i < arena->page_count; i += 1)
    {
        uha_free(a, arena->pages[i]);
    }

    lpn_dead = arena->lpn;
    while(lpn_dead != NULL)
    {
        uh_large_page_node* lpn_next = lpn_dead->next;
        uha_free(a, lpn_dead->page);
        uha_free(a, lpn_dead);
        lpn_dead = lpn_next;
    }

    lpn_dead = arena->free_list;
    while(lpn_dead != NULL)
    {
        uh_large_page_node* lpn_next = lpn_dead->next;
        uha_free(a, lpn_dead->page);
        uha_free(a, lpn_dead);
        lpn_dead = lpn_next;
    }

    uha_free(a, arena->pages);
}

UHEXTERN_C_END
#endif // __UH_ALLOCATOR