#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->state = EXITING;
    currentThread        = next;
    currentThread->state = RUNNING;
    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;

    current->state = RUNNING;

    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 struct 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;

    thread->state = READY;
    uint32_t flags;
    disable_IRQs(flags);
    list_add_tail(currentThread, thread);
    restore_IRQs(flags);
    return thread;
free_mem:
    free((void *)thread->stackAddr);
    free((void *)thread);
    return NULL;
}

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

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

struct kthread *kthreadSelectNext()
{
    struct kthread *nextThread;
    int idx;
    list_foreach(currentThread->next, nextThread, idx)
    {
        if (nextThread->state == READY) {
            return nextThread;
        }
    }
    return currentThread;
}

struct cpu_state *kthreadSwitch(struct cpu_state *prevCpu)
{
    uint32_t flags;
    disable_IRQs(flags);
    currentThread->cpuState    = prevCpu;
    currentThread->state       = READY;
    struct kthread *nextThread = kthreadSelectNext();
    currentThread        = nextThread;
    currentThread->state = RUNNING;
    restore_IRQs(flags);
    return nextThread->cpuState;
}

int kthreadOnJieffiesTick()
{
    struct kthread *nextThread;
    int idx;
    uint32_t flags;
    disable_IRQs(flags);
    list_foreach(currentThread, nextThread, idx)
    {
        if (nextThread->state == SLEEPING && nextThread->jiffiesSleeping) {
            nextThread->jiffiesSleeping--;
            if (!nextThread->jiffiesSleeping) {
                nextThread->state = READY;
            }
        }
    }
    restore_IRQs(flags);
    return 0;
}

int kthreadUnsched(struct kthread *th)
{
    list_delete(currentThread, th);

    return 0;
}

int kthreadSaveAndYield(struct kthread *current)
{
    uint32_t flags;

    disable_IRQs(flags);
    struct kthread *next = kthreadSelectNext();

    if (current == next) {
        restore_IRQs(flags);
        return 0;
    }

    assertmsg(next->state == READY, "thread %s is in state %d\n", next->name, next->state);

    if (current->state == RUNNING)
        current->state = READY;

    currentThread        = next;
    currentThread->state = RUNNING;
    cpu_context_switch(&current->cpuState, next->cpuState);
    restore_IRQs(flags);

    return 0;
}

int kthreadYield()
{
    uint32_t flags;

    disable_IRQs(flags);
    struct kthread *next    = kthreadSelectNext();
    struct kthread *current = currentThread;

    if (current == next) {
        restore_IRQs(flags);
        return 0;
    }

    assert(current->state == RUNNING);
    assertmsg(next->state == READY, "thread %s is in state %d\n", next->name, next->state);

    if (current->state == RUNNING)
        current->state = READY;

    currentThread        = next;
    currentThread->state = RUNNING;
    cpu_context_switch(&current->cpuState, next->cpuState);
    restore_IRQs(flags);

    return 0;
}

int kthreadMsleep(unsigned long msec)
{
    uint32_t flags;
    struct kthread *next, *current;

    disable_IRQs(flags);

    current = currentThread;
    assertmsg(current->state == RUNNING, "thread %s is in state %d for %d\n", current->name,
              current->state, msec);

    current->state           = SLEEPING;
    current->jiffiesSleeping = msecs_to_jiffies(msec);
    next                     = kthreadSelectNext();

    assert(next != current);
    assert(next->state == READY);

    currentThread        = next;
    currentThread->state = RUNNING;
    cpu_context_switch(&current->cpuState, next->cpuState);

    restore_IRQs(flags);
    return 0;
}

struct kthread *getCurrentThread()
{
    return currentThread;
}

int kthreadAddThread(struct kthread *th)
{
    if (th->state == READY)
        return 0;

    th->state = READY;
    list_add_tail(currentThread, th);

    return 0;
}