#include <shadow-tag.h>
#include <context.h>
#include <string.h>
#include <stdlib.h>

void mt_init_shadow_tag_pool_allocation(MtShadowTagPool *pool, uint32_t capacity)
{
    uint32_t slot_index = 0;
    while(slot_index < capacity)
    {
        pool->allocation[slot_index].pool = pool;
        pool->allocation[slot_index].next = &pool->allocation[slot_index+1];
        ++slot_index;
    }
    pool->allocation[capacity - 1].next = NULL;
}

void mt_init_shadow_tag_pool(MtShadowTagPool *pool, uint32_t capacity)
{
    pool->capacity = capacity;
    pool->allocation = calloc(sizeof(MtShadowTag), capacity);
    pool->first_free = pool->allocation;
    pool->continuation = NULL;
    mt_init_shadow_tag_pool_allocation(pool, capacity);
}

MtShadowTagPool * mt_new_shadow_tag_pool(uint32_t capacity)
{
    MtShadowTagPool *pool = malloc(sizeof(MtShadowTagPool));
    mt_init_shadow_tag_pool(pool, capacity * 2);
    return pool;
}

void mt_free_shadow_tag_pool(MtShadowTagPool *pool)
{
    if(pool->continuation != NULL)
        mt_free_shadow_tag_pool(pool->continuation);
    free(pool->allocation);
    free(pool);
}

MtShadowTag * mt_alloc_shadow_tag_in_pool(MtShadowTagPool *pool)
{
    if(pool->first_free == NULL)
    {
        if(pool->continuation == NULL)
            pool->continuation = mt_new_shadow_tag_pool(pool->capacity * 2);
        return mt_alloc_shadow_tag_in_pool(pool->continuation);
    }
    MtShadowTag *allocated_tag = pool->first_free;
    pool->first_free = allocated_tag->next;

    // Reset the tag's old data
    memset(allocated_tag, sizeof(MtShadowTag), 0x00);
    allocated_tag->pool = pool;
    
    return allocated_tag;
}

MtShadowTag * mt_alloc_shadow_tag(void *context_ptr)
{
    MtContext *context = context_ptr;
    return mt_alloc_shadow_tag_in_pool(&context->shadow_tag_system.root_pool);
}

MtShadowTag * mt_alloc_shadow_tag_before(MtShadowTag *chain_start)
{
    MtShadowTag *tag = mt_alloc_shadow_tag_in_pool(chain_start->pool);
    tag->next = chain_start;
    return tag;
}

void mt_free_shadow_tag_chain(MtShadowTag *tag)
{
    MtShadowTagPool *pool = tag->pool;
    MtShadowTag *last_tag = tag;
    while(last_tag->next != NULL)
    {
        last_tag = last_tag->next;
    }
    last_tag->next = pool->first_free;
    pool->first_free = tag;
}



void mt_init_shadow_tag_name_arena(MtShadowTagNameArena *arena, uint32_t size)
{
    arena->usage = 0;
    arena->capacity = size;
    arena->allocation = malloc(size);
    arena->continuation = NULL;
}

MtShadowTagNameArena * mt_new_shadow_tag_name_arena(uint32_t size)
{
    MtShadowTagNameArena *arena = malloc(sizeof(MtShadowTagNameArena));
    mt_init_shadow_tag_name_arena(arena, size);
    return arena;
}

void mt_free_shadow_tag_name_arena(MtShadowTagNameArena *arena)
{
    if(arena->continuation != NULL)
        mt_free_shadow_tag_name_arena(arena->continuation);
    free(arena->allocation);
    free(arena);
}

void mt_init_shadow_tag_name_registry(MtShadowTagNameRegistry *name_registry)
{
    name_registry->root_arena = mt_new_shadow_tag_name_arena(4096);
    name_registry->id_counter = 1;
    name_registry->num_names = 0;
    name_registry->names_capacity = 256;
    name_registry->names = calloc(sizeof(MtShadowTagName), name_registry->names_capacity);
}

void mt_cleanup_shadow_tag_name_registry(MtShadowTagNameRegistry *name_registry)
{
    mt_free_shadow_tag_name_arena(name_registry->root_arena);
    free(name_registry->names);
}

void mt_init_shadow_tag_system(MtShadowTagSystem *tag_system)
{
    mt_init_shadow_tag_name_registry(&tag_system->name_registry);
    mt_init_shadow_tag_pool(&tag_system->root_pool, 8192);
}

void mt_cleanup_shadow_tag_system(MtShadowTagSystem *tag_system)
{
    mt_free_shadow_tag_pool(&tag_system->root_pool);
    mt_cleanup_shadow_tag_name_registry(&tag_system->name_registry);
}



char * mt_shadow_tag_name_arena_alloc(MtShadowTagNameArena *arena, uint32_t num_bytes)
{
    uint32_t end_size = arena->usage + num_bytes;
    if(end_size >= arena->capacity)
    {
        if(arena->continuation == NULL)
            arena->continuation = mt_new_shadow_tag_name_arena(arena->capacity * 2);
        return mt_shadow_tag_name_arena_alloc(arena->continuation, num_bytes);
    }
    char *allocation = arena->allocation + arena->usage;
    arena->usage = end_size;
    return allocation;
}

uint32_t mt_force_create_new_tag_name_entry(MtShadowTagNameRegistry *name_registry, const char *tag_name)
{
    if(name_registry->num_names >= name_registry->names_capacity)
    {
        name_registry->names_capacity *= 2;
        name_registry->names = realloc(name_registry->names, sizeof(MtShadowTagName) * name_registry->names_capacity);
    }
    MtShadowTagName *name_entry = &name_registry->names[name_registry->num_names++];

    uint32_t len_tag_name = strlen(tag_name);
    char *tag_name_copy = mt_shadow_tag_name_arena_alloc(name_registry->root_arena, len_tag_name+1);
    memcpy(tag_name_copy, tag_name, len_tag_name + 1);

    name_entry->string_in_arena = tag_name_copy;
    name_entry->identifier = name_registry->id_counter++;
    
    return name_entry->identifier;
}

uint32_t mt_tag_name_to_id(void *context_ptr, const char *tag_name)
{
    MtContext *context = context_ptr;
    MtShadowTagNameRegistry *name_registry = &context->shadow_tag_system.name_registry;

    uint32_t name_index = 0;
    while(name_index < name_registry->num_names)
    {
        if(!strcmp(name_registry->names[name_index].string_in_arena, tag_name))
            return name_registry->names[name_index].identifier;
        ++name_index;
    }
    return mt_force_create_new_tag_name_entry(name_registry, tag_name);
}