// SPDX-License-Identifier: ISC

#include <zinnia/alloc.h>
#include <zinnia/assert.h>
#include <zinnia/atomic.h>
#include <zinnia/errno.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) {
    arch_sched_jump_to_context((struct arch_context*)ctx);
}

static void task_init_common(struct task* task, const char* name, struct vma* space) {
    task->space = space;
    spin_init(&task->lock);
    strncpy(task->name, name, sizeof(task->name));
    task->base_prio = SCHED_PRIO_DEFAULT;
    task->time_slice = SCHED_TIME_SLICE_NS;
    task->penalty = 0;
    task->feedback_epoch = 0;
    atomic_store(&task->priority, SCHED_PRIO_DEFAULT, ATOMIC_RELAXED);
    tailq_init(&task->pi_lenders);
}

errno_t task_create_current(const char* name, struct vma* space, struct task** out) {
    struct task* task = mem_alloc(sizeof(struct task), 0);
    if (task == nullptr)
        return ENOMEM;

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

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

constexpr size_t KERNEL_STACK_SIZE = 16 * 1024;

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->turnstile = turnstile_alloc();
    if (new_task->turnstile == nullptr) {
        mem_free(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);
        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);
        turnstile_free(new_task->turnstile);
        mem_free(new_task);
        return s;
    }

    task_init_common(new_task, name, space);
    *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);
    mem_free(task);
}