#include "tsk/tasking.h"
#include "mem/memory.h"
#include "io/log.h"
#include "lib/hashtable.h"
#include "lib/string.h"
#include "paging.h"
#include "cpu.h"
#include "tss.h"

struct TaskBody {
    struct Task base;
    struct PageDirectory *pd;
    struct Registers state;
};

struct Task *task_current;
uintptr_t _kernel_task_bp = 0;

//static struct SLinkedList s_tasks;
static struct HashTable s_tasks;

static struct SlabCache s_task_cache;
static struct SlabCache s_kbp_cache;

static tid_t s_task_id_next = 1;

static size_t
s_task_hash_func(const void* tid, size_t _)
{
    return (tid_t)tid;
}

void
tasking_setup(void)
{
    hashtable_new(&s_tasks, struct TaskBody*, tid_t);
    s_tasks.hash_function = s_task_hash_func;
 
    mem_slabcache_new(&s_task_cache, "tasks", sizeof(struct TaskBody));
    mem_slabcache_new(&s_kbp_cache, "kernel stacks", 4096);

    struct TaskBody *ktask = mem_slab_alloc(&s_task_cache);
    *ktask = (struct TaskBody){
        .base.id = s_task_id_next++,
        .base.kbp = ((uintptr_t)mem_slab_alloc(&s_kbp_cache)) + 0xFF0,
        .base.perm = (size_t)-1,
        .pd = mem_current_pd
    };
    hashtable_insert(&s_tasks, 0, ktask);

    task_current = (struct Task*)ktask;
    tss_set_rsp(0, task_current->kbp);
    _kernel_task_bp = task_current->kbp;
}

struct Task*
task_new(struct Task *parent)
{
    struct TaskBody *new = mem_slab_alloc(&s_task_cache);
    memcpy(new, parent, sizeof(struct TaskBody));
    new->base.id = s_task_id_next++;
    
    uintptr_t ksp_base = (uintptr_t)mem_slab_alloc(&s_kbp_cache);
    new->base.kbp = ksp_base + 0xFFF;
    new->base.perm = parent->perm;

    hashtable_insert(&s_tasks, (void*)new->base.id, new);
    return (struct Task*)new;
}

struct Task*
task_get(tid_t id)
{
    struct Task **task = hashtable_get(&s_tasks, (void*)id, struct Task*);
    if(task == NULL) return NULL;
    return *task;
}

void
task_free(struct Task *task)
{
    struct TaskBody *body = (struct TaskBody*)task;
    body->pd->references--;
    task->kbp -= 0xFFF;
    mem_slab_free(&s_kbp_cache, (void*)(task->kbp));
    klogf("Need impl for task_free\n");
}

void*
task_get_pd(struct Task *task)
{
    return ((struct TaskBody*)task)->pd;
}