#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/atomic.h>
#include <zinnia/errno.h>
#include <zinnia/irq.h>
#include <zinnia/list.h>
#include <zinnia/percpu.h>
#include <zinnia/pmap.h>
#include <zinnia/print.h>
#include <zinnia/sched.h>
#include <zinnia/spin.h>
#include <zinnia/stack.h>
#include <zinnia/tailq.h>
#include <zinnia/vma.h>
#include <arch/sched.h>
#include <stdint.h>

static void idle_func(uintptr_t) {
    while (1) {
        irq_set_state(true);
        irq_wait();
    }
}

void sched_to_user(uintptr_t info_ptr) {
    struct task_start_info* info = (struct task_start_info*)info_ptr;
    uintptr_t ip = info->ip;
    uintptr_t sp = info->sp;
    uintptr_t arg = info->arg;
    mem_free(info);
    arch_sched_jump_to_user(ip, sp, arg);
}

void sched_to_user_context(uintptr_t ctx) {
    arch_sched_jump_to_context((struct arch_context*)ctx);
}

void sched_init(struct scheduler* sched) {
    struct task* idle_task;
    assert(task_create("idle", &kernel_vma, idle_func, 0, &idle_task) == 0, "Unable to create idle task!");
    sched->idle_task = idle_task;

    struct task* boot_task = mem_alloc(sizeof(struct task), 0);
    assert(boot_task != nullptr, "Unable to create boot task!");
    strncpy(boot_task->name, "boot", sizeof(boot_task->name));
    boot_task->space = &kernel_vma;
    sched->current = boot_task;

    tailq_init(&sched->run_queue);

    log("sched: Scheduler initialized\n");
}

void sched_add_task(struct scheduler* sched, struct task* task) {
    assert(task != nullptr, "Tried to add NULL to the run queue");
    task->state = TASK_STATE_READY;
    tailq_insert_tail(&sched->run_queue, task, next);
}

static struct task* next_task(struct scheduler* sched) {
    struct task* t = tailq_first(&sched->run_queue);

    if (t == nullptr)
        return sched->idle_task;

    tailq_remove(&sched->run_queue, t, next);
    t->state = TASK_STATE_RUNNING;

    return t;
}

static void finish_switch(struct task* prev) {
    if (prev->state == TASK_STATE_DEAD)
        task_destroy(prev);
}

static void do_reschedule(struct scheduler* sched) {
    struct task* from = sched->current;
    struct task* to = next_task(sched);

    // We'd end up on the same task, skip switching altogether.
    if (from == to) {
        irq_unlock();
        return;
    }

    sched->current = to;
    pmap_set(&to->space->pmap);

    struct percpu* cpu = percpu_get();
    from->user_stack = cpu->user_stack;
    cpu->kernel_stack = kernel_stack_top(&to->kernel_stack);
    cpu->user_stack = to->user_stack;

    struct irq_percpu irq_state = *percpu_ref(irq_data);
    struct task* prev = arch_sched_switch(from, to);
    *percpu_ref(irq_data) = irq_state;
    finish_switch(prev);
    irq_unlock();
}

void sched_reschedule(struct scheduler* sched) {
    irq_lock();
    if (sched->current != sched->idle_task && sched->current->state != TASK_STATE_DEAD)
        sched_add_task(sched, sched->current);
    do_reschedule(sched);
}

void sched_yield(struct scheduler* sched) {
    irq_lock();
    do_reschedule(sched);
}

void sched_block(struct scheduler* sched) {
    irq_lock();
    if (atomic_exchange(&sched->current->wake_pending, false, ATOMIC_ACQ_REL)) {
        irq_unlock();
        return;
    }
    sched->current->state = TASK_STATE_BLOCKED;
    do_reschedule(sched);
}

void sched_wake(struct task* task) {
    irq_lock();
    if (task->state == TASK_STATE_BLOCKED)
        sched_add_task(&percpu_get()->sched, task);
    else
        atomic_store(&task->wake_pending, true, ATOMIC_RELEASE);
    irq_unlock();
}

errno_t task_create(const char* name, struct vma* space, task_fn_t entry, uintptr_t arg0, struct task** out) {
    if (!out)
        return EINVAL;

    struct task* new_task = mem_alloc(sizeof(struct task), 0);
    if (new_task == nullptr)
        return ENOMEM;

    static atomic_size_t last_tid = 0;
    new_task->id = atomic_fetch_add(&last_tid, 1, ATOMIC_ACQUIRE);
    new_task->space = space;
    strncpy(new_task->name, name, sizeof(new_task->name));

    // Allocate a kernel stack.
    errno_t s = stack_init(&new_task->kernel_stack, KERNEL_STACK_SIZE);
    if (s != 0) {
        mem_free(new_task);
        return s;
    }
    new_task->user_stack = 0;

    // Initialize the arch-dependent parts.
    s = arch_task_init(&new_task->context, entry, arg0, kernel_stack_top(&new_task->kernel_stack), true);
    if (s != 0) {
        stack_destroy(&new_task->kernel_stack);
        mem_free(new_task);
        return s;
    }

    *out = new_task;
    return 0;
}

void task_destroy(struct task* task) {
    arch_task_fini(&task->context);
    stack_destroy(&task->kernel_stack);
    mem_free(task);
}

void task_entry(task_fn_t entry, uintptr_t arg0, struct task* prev) {
    finish_switch(prev);

    *percpu_ref(irq_data) = (struct irq_percpu){
        .depth = 1,
        .saved_state = true,
    };
    irq_unlock();

    entry(arg0);

    struct scheduler* sched = &percpu_get()->sched;
    sched->current->state = TASK_STATE_DEAD;
    sched_yield(sched);
    panic("unreachable");
}

[[__init]]
static void sched_stage_init(struct boot_info*) {
    sched_init(&percpu_get()->sched);
}
INIT_TASK(sched_stage, "core.sched", sched_stage_init);
INIT_DEPENDS(sched_stage, percpu_stage);