#include "kthread.h"
#include "alloc.h"
#include "assert.h"
#include "irq.h"
#include "klibc.h"
#include "list.h"
#include "time.h"
#include "vga.h"

static struct kthread *currentThread;

void kthreadExit()
{
    uint32_t flags;
    disable_IRQs(flags);
    struct kthread *current = currentThread;
    struct kthread *next    = kthreadSelectNext();
    if (next == current)
        assert("cannot exit thread");
    currentThread = next;
    cpu_context_exit_to(next->cpuState, (cpu_kstate_function_arg1_t *)kthreadDelete,
                        (uint32_t)current);
    restore_IRQs(flags);
    return;
}

int kthreadSetup(vaddr_t mainStack, size_t mainStackSize)
{
    struct kthread *current = (struct kthread *)malloc(sizeof(struct kthread));
    strzcpy(current->name, "[KINIT]", KTHREAD_NAME_MAX_LENGTH);
    current->stackAddr = mainStack;
    current->stackSize = mainStackSize;

    list_singleton(currentThread, current);

    return 0;
}

struct kthread *kthreadCreate(const char *name, cpu_kstate_function_arg1_t func, void *args)
{
    struct kthread *thread = (struct kthread *)malloc(sizeof(struct kthread));
    if (!thread)
        return NULL;

    thread->stackAddr = (vaddr_t)malloc(KTHREAD_DEFAULT_STACK_SIZE);
    printf("Alloc stask at 0x%x strcut at 0x%x\n", thread->stackAddr, thread);
    thread->stackSize = KTHREAD_DEFAULT_STACK_SIZE;

    if (!thread->stackAddr)
        goto free_mem;

    if (name)
        strzcpy(thread->name, name, KTHREAD_NAME_MAX_LENGTH);
    else
        strzcpy(thread->name, "[UNKNOW]", KTHREAD_NAME_MAX_LENGTH);

    if (cpu_kstate_init(&thread->cpuState, (cpu_kstate_function_arg1_t *)func, (vaddr_t)args,
                        thread->stackAddr, thread->stackSize,
                        (cpu_kstate_function_arg1_t *)kthreadExit, 0))
        goto free_mem;

    list_add_tail(currentThread, thread);
    return thread;
free_mem:
    free((void *)thread->stackAddr);
    free((void *)thread);
    return NULL;
}

void kthreadDelete(struct kthread *thread)
{
    list_delete(currentThread, thread);

    free((void *)thread->stackAddr);
    free((void *)thread);
}

struct kthread *kthreadSelectNext()
{
    struct kthread *nextThread = currentThread->next;
    return nextThread;
}

struct cpu_state *kthreadSwitch(struct cpu_state *prevCpu)
{
    uint32_t flags;
    disable_IRQs(flags);
    currentThread->cpuState    = prevCpu;
    struct kthread *nextThread = kthreadSelectNext();
    printStringDetails(nextThread->name, RED, BLACK, 40, VGA_HEIGHT - 1);
    currentThread = nextThread;
    restore_IRQs(flags);
    return nextThread->cpuState;
}

int kthreadYield()
{
    uint32_t flags;
    disable_IRQs(flags);
    struct kthread *next    = kthreadSelectNext();
    struct kthread *current = currentThread;
    currentThread           = next;
    cpu_context_switch(&current->cpuState, next->cpuState);
    restore_IRQs(flags);
    return 0;
}