#include "thread.h"
#include "alloc.h"
#include "assert.h"
#include "irq.h"
#include "klibc.h"
#include "list.h"
#include "mmuContext.h"
#include "process.h"
#include "time.h"
#include "types.h"
#include "vga.h"

static struct thread *currentThread;
static struct thread *threadWithTimeout;
static thread_id_t nextTid; // This is the TID for kernel thread ONLY

pid_t threadGetId(struct thread *th)
{
    return th->tid;
}

static void threadPrepareContext(struct thread *th);

void threadExit()
{
    uint32_t flags;

    disable_IRQs(flags);

    struct thread *current = currentThread;
    struct thread *next    = threadSelectNext();

    if (next == current)
        assert("cannot exit thread");

    currentThread->state = EXITING;
    currentThread        = next;
    currentThread->state = RUNNING;
    threadPrepareContext(next);
    cpu_context_exit_to(next->cpuState, (cpu_kstate_function_arg1_t *)threadDelete,
                        (uint32_t)current);

    restore_IRQs(flags);

    return;
}

int threadSetup(vaddr_t mainStack, size_t mainStackSize)
{
    struct thread *current = (struct thread *)malloc(sizeof(struct thread));
    if (current == NULL)
        return -ENOMEM;
    strzcpy(current->name, "[KINIT]", THREAD_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 thread *threadCreate(const char *name, cpu_kstate_function_arg1_t func, void *args)
{
    struct thread *thread = (struct thread *)malloc(sizeof(struct thread));
    if (!thread)
        return NULL;

    thread->stackAddr = (vaddr_t)malloc(THREAD_DEFAULT_STACK_SIZE);
    if (!thread->stackAddr){
        free(thread);
        return NULL;
    }
#ifdef DEBUG
    printf("Alloc stack at 0x%p struct at 0x%p\n", (void *)thread->stackAddr, thread);
#endif
    thread->stackSize = THREAD_DEFAULT_STACK_SIZE;


    if (name)
        strzcpy(thread->name, name, THREAD_NAME_MAX_LENGTH);
    else
        strzcpy(thread->name, "[UNKNOW]", THREAD_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 *)threadExit, 0))
        goto free_mem;

    thread->state  = READY;
    thread->tid    = nextTid++;
    thread->wqExit = NULL;
    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;
}

struct thread *threadCreateUser(const char *name, struct process *proc, uaddr_t startPc,
                                uint32_t arg1, uint32_t arg2, uaddr_t startSP)
{
    struct thread *thread = malloc(sizeof(struct thread));

    if (thread == NULL)
        return NULL;

    thread->stackAddr = (vaddr_t)malloc(THREAD_DEFAULT_STACK_SIZE);
    if (!thread->stackAddr) {
        free(thread);
        return NULL;
    }
    thread->stackSize = THREAD_DEFAULT_STACK_SIZE;

    thread->wqExit = (struct wait_queue * )malloc(sizeof(struct wait_queue));
    if (!thread->wqExit) {
        free((void *)thread->stackAddr);
        free(thread);
        return NULL;
    }
    waitQueueInit(thread->wqExit);

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

    if (cpu_ustate_init(&thread->cpuState, startPc, arg1, arg2, startSP, thread->stackAddr,
                        thread->stackSize)) {
        goto free_mem;
    }

    if(processAddThread(proc, thread))
        goto free_mem;

    thread->state = READY;
    thread->tid = processGetNextTid(proc);
    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 threadDelete(struct thread *thread)
{
    uint32_t flags;
    disable_IRQs(flags);
    list_delete(currentThread, thread);
    restore_IRQs(flags);
    assert(thread->state == EXITING);

    if (thread->wqExit) {
        waitUp(thread->wqExit);
    }

    if (thread->squattedContext) {
        mmuContextUnref(thread->squattedContext);
    }

    if (thread->process)
        processRemoveThread(thread);

#ifdef DEBUG
    printf("Free stack at 0x%p struct at 0x%p\n", (void *)thread->stackAddr, thread);
#endif
    free((void *)thread->stackAddr);
    free((void *)thread);
}

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

struct cpu_state *threadSwitch(struct cpu_state *prevCpu)
{
    uint32_t flags;
    struct thread *nextThread;

    disable_IRQs(flags);

    nextThread              = threadSelectNext();
    currentThread->cpuState = prevCpu;
    if (nextThread != currentThread) {
        currentThread->state = READY;
//        printf("    Switch from %s to %s\n", currentThread->name, nextThread->name);
        currentThread        = nextThread;
        currentThread->state = RUNNING;
        threadPrepareContext(nextThread);
    }

    restore_IRQs(flags);

    return nextThread->cpuState;
}

int threadCount()
{
    struct thread *nextThread;
    int idx;
    uint32_t flags;
    disable_IRQs(flags);
    list_foreach(currentThread, nextThread, idx){
    }
    return idx;
}

int threadOnJieffiesTick()
{
    struct thread *nextThread;
    int idx;
    uint32_t flags;
    disable_IRQs(flags);
    list_foreach(currentThread, nextThread, idx)
    {
        if (nextThread->state == SLEEPING) {
            if (nextThread->jiffiesSleeping)
                nextThread->jiffiesSleeping--;
            if (!nextThread->jiffiesSleeping)
                nextThread->state = READY;
        }
    }
    list_foreach_named(threadWithTimeout, nextThread, idx, timePrev, timeNext)
    {
        if (nextThread->state == WAITING) {
            if (nextThread->jiffiesSleeping)
                nextThread->jiffiesSleeping--;
            if (!nextThread->jiffiesSleeping) {
                nextThread->sleepHaveTimeouted = 1;
                list_delete_named(threadWithTimeout, nextThread, timePrev, timeNext);
                threadAddThread(nextThread);
            }
        }
    }
    restore_IRQs(flags);
    return 0;
}

int threadUnsched(struct thread *th)
{
    list_delete(currentThread, th);

    return 0;
}

// Must be called with IRQ disabled
int threadWait(struct thread *current, struct thread *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;
    threadPrepareContext(next);
    cpu_context_switch(&current->cpuState, next->cpuState);

    return current->sleepHaveTimeouted;
}

int threadYield()
{
    uint32_t flags;

    disable_IRQs(flags);
    struct thread *next    = threadSelectNext();
    struct thread *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;
//    printf("    Yield from %s to %s\n", currentThread->name, next->name);
    currentThread        = next;
    currentThread->state = RUNNING;
    threadPrepareContext(next);
    cpu_context_switch(&current->cpuState, next->cpuState);
    restore_IRQs(flags);

    return 0;
}

int threadMsleep(unsigned long msec)
{
    return threadUsleep(msec*1000);
}

int threadUsleep(unsigned long usec)
{
    uint32_t flags;
    struct thread *next, *current;

    disable_IRQs(flags);

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

    current->state              = SLEEPING;
    current->sleepHaveTimeouted = 0;
    current->jiffiesSleeping    = usecs_to_jiffies(usec);
    next                        = threadSelectNext();

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

    currentThread        = next;
    currentThread->state = RUNNING;
    threadPrepareContext(next);
    cpu_context_switch(&current->cpuState, next->cpuState);
    restore_IRQs(flags);
    return current->sleepHaveTimeouted == 1;
}

struct thread *getCurrentThread()
{
    return currentThread;
}

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

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

    return 0;
}

static void threadPrepareContext(struct thread *th)
{
    if (cpu_context_is_in_user_mode(th->cpuState)) {
        assert(th->process != NULL);
        assert(th->squattedContext == NULL);
        mmuContextSwitch(processGetMMUContext(th->process));
    } else if (th->squattedContext) {
        mmuContextSwitch(th->squattedContext);
    }
}

int threadChangeCurrentContext(struct mmu_context *ctx)
{
    uint32_t flags;
    struct mmu_context *prev;

    assert(currentThread != NULL);

    prev = currentThread->squattedContext;

    if (ctx != NULL) {
        assert(prev == NULL);
    } else {
        assert(prev != NULL);
    }

    disable_IRQs(flags);
    currentThread->squattedContext = ctx;

    if (ctx != NULL) {
        mmuContextRef(ctx);
        mmuContextSwitch(ctx);
    } else {
        mmuContextUnref(prev);
    }
    restore_IRQs(flags);

    return 0;
}

void threadPrepareSyscallSwitchBack(struct cpu_state *cpuState)
{
    currentThread->cpuState = cpuState;
    threadPrepareContext(currentThread);
}

void threadPrepareExceptionSwitchBack(struct cpu_state *cpuState)
{
    currentThread->cpuState = cpuState;
    threadPrepareContext(currentThread);
}

void threadPrepareIrqServicing(struct cpu_state *cpuState)
{
    currentThread->cpuState = cpuState;
}

void threadPrepareIrqSwitchBack(struct cpu_state *cpuState)
{
    currentThread->cpuState = cpuState;
    threadPrepareContext(currentThread);
}