// SPDX-License-Identifier: ISC

#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/atomic.h>
#include <zinnia/errno.h>
#include <zinnia/process.h>
#include <zinnia/sched.h>
#include <zinnia/spin.h>
#include <zinnia/stack.h>
#include <zinnia/tailq.h>
#include <zinnia/turnstile.h>
#include <zinnia/vma.h>
#include <arch/sched.h>
#include <stdint.h>
#include <string.h>
#include "internal.h"

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) {
    struct arch_context context = *(struct arch_context*)ctx;
    mem_free((void*)ctx);
    arch_sched_jump_to_context(&context);
}

static struct spinlock task_cache_lock;
static tailq_head(struct task) task_cache = {nullptr, &task_cache.tqh_first};

static struct task* task_alloc() {
    spin_lock(&task_cache_lock);
    struct task* task = tailq_first(&task_cache);
    if (task)
        tailq_remove(&task_cache, task, next);
    spin_unlock(&task_cache_lock);

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

    size_t keep_start = offsetof(struct task, on_cpu);
    size_t keep_end = keep_start + sizeof(task->on_cpu);
    memset(task, 0, keep_start);
    memset((char*)task + keep_end, 0, sizeof(struct task) - keep_end);
    return task;
}

static void task_release(struct task* task) {
    spin_lock(&task_cache_lock);
    tailq_insert_tail(&task_cache, task, next);
    spin_unlock(&task_cache_lock);
}

static void task_init_common(struct task* task, const char* name, struct process* proc) {
    task->process = proc;
    task->space = proc->space;
    process_ref(proc);
    if (task->space != &kernel_vma)
        vma_ref(task->space);
    spin_init(&task->lock);
    atomic_store(&task->lockp, &task->lock, ATOMIC_RELAXED);
    strncpy(task->name, name, sizeof(task->name));
    task->base_prio = process_priority(proc);
    task->time_slice = SCHED_TIME_SLICE_NS;
    task->penalty = 0;
    task->feedback_epoch = 0;
    atomic_store(&task->priority, task->base_prio, ATOMIC_RELAXED);
    tailq_init(&task->pi_lenders);
    process_attach_task(proc, task);
}

errno_t task_create_current(const char* name, struct process* proc, struct task** out) {
    struct task* task = task_alloc();
    if (task == nullptr)
        return ENOMEM;

    task->turnstile = turnstile_alloc();
    if (task->turnstile == nullptr) {
        task_release(task);
        return ENOMEM;
    }

    task_init_common(task, name, proc);
    *out = task;
    return 0;
}

constexpr size_t KERNEL_STACK_SIZE = 16 * 1024;

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

    struct task* new_task = task_alloc();
    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->turnstile = turnstile_alloc();
    if (new_task->turnstile == nullptr) {
        task_release(new_task);
        return ENOMEM;
    }

    // Allocate a kernel stack.
    errno_t s = stack_init(&new_task->kernel_stack, KERNEL_STACK_SIZE);
    if (s != 0) {
        turnstile_free(new_task->turnstile);
        task_release(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);
        turnstile_free(new_task->turnstile);
        task_release(new_task);
        return s;
    }

    task_init_common(new_task, name, proc);
    *out = new_task;
    return 0;
}

void task_destroy(struct task* task) {
    assert(tailq_empty(&task->pi_lenders), "Task %zu exited while owning contended locks", task->id);
    turnstile_free(task->turnstile);
    arch_task_fini(&task->context);
    stack_destroy(&task->kernel_stack);
    process_detach_task(task->process, task);
    if (task->space != &kernel_vma)
        vma_unref(task->space);
    process_unref(task->process);
    task_release(task);
}

void task_lock(struct task* task) {
    while (true) {
        struct spinlock* lock = atomic_load(&task->lockp, ATOMIC_ACQUIRE);
        spin_lock(lock);
        if (atomic_load(&task->lockp, ATOMIC_RELAXED) == lock)
            return;
        spin_unlock(lock);
    }
}

void task_unlock(struct task* task) {
    spin_unlock(atomic_load(&task->lockp, ATOMIC_RELAXED));
}