#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;
static struct kthread *threadWithTimeout;

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);
    list_init_named(threadWithTimeout, timePrev, timeNext);

    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);
#ifdef DEBUG
    printf("Alloc stack at 0x%x struct at 0x%x\n", thread->stackAddr, thread);
#endif
    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);

#ifdef DEBUG
    printf("Free stack at 0x%x struct at 0x%x\n", thread->stackAddr, thread);
#endif
    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;
    struct kthread *nextThread;

    disable_IRQs(flags);

    nextThread              = kthreadSelectNext();
    currentThread->cpuState = prevCpu;
    currentThread->state    = READY;
    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;
            }
        }
    }
    list_foreach_named(threadWithTimeout, nextThread, idx, timePrev, timeNext)
    {
        if (nextThread->state == WAITING && nextThread->jiffiesSleeping) {
            nextThread->jiffiesSleeping--;
            if (!nextThread->jiffiesSleeping) {
                nextThread->sleepHaveTimeouted = 1;
                list_delete_named(threadWithTimeout, nextThread, timePrev, timeNext);
                kthreadAddThread(nextThread);
            }
        }
    }
    restore_IRQs(flags);
    return 0;
}

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

    return 0;
}

// Must be called with IRQ disabled
int kthreadWait(struct kthread *current, struct kthread *next, unsigned long msec)
{
    if (current == next) {
        assertmsg(0, "Cannot yield from %s to %s\n", current->name, next->name);
        return 0;
    }

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

    current->jiffiesSleeping    = msecs_to_jiffies(msec);
    current->sleepHaveTimeouted = 0;

    if (current->jiffiesSleeping)
        list_add_tail_named(threadWithTimeout, current, timePrev, timeNext);

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

    return current->sleepHaveTimeouted;
}

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->sleepHaveTimeouted = 0;
    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 current->sleepHaveTimeouted == 1;
}

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