#include <zinnia/handle.h>
#include <zinnia/rights.h>
#include <zinnia/status.h>
#include <zinnia/thread.h>
#include <common/compiler.h>
#include <kernel/alloc.h>
#include <kernel/namespace.h>
#include <kernel/percpu.h>
#include <kernel/print.h>
#include <kernel/sched.h>
#include <kernel/syscalls.h>
#include <kernel/usercopy.h>
#include <stddef.h>
#include <string.h>

zn_status_t syscall_thread_create(struct arch_context* ctx) {
    zn_handle_t ns_handle = ctx->ARCH_CTX_A0;
    __user const char* name = (__user const char*)ctx->ARCH_CTX_A1;
    size_t name_len = ctx->ARCH_CTX_A2;
    enum zn_thread_flags flags = ctx->ARCH_CTX_A3;
    __user zn_handle_t* out = (__user zn_handle_t*)ctx->ARCH_CTX_A4;

    struct task* current = percpu_get()->sched.current;
    // Resolve namespace.
    struct namespace* ns;
    if (ns_handle == ZN_HANDLE_THIS_NS) {
        ns = current->namespace;
    } else {
        struct namespace_desc ns_desc;
        zn_status_t status = namespace_get(current->namespace, ns_handle, &ns_desc);
        if (status != ZN_OK)
            return status;
        if (ns_desc.type != NAMESPACE_DESC_NS)
            return ZN_ERR_BAD_HANDLE;
        ns = ns_desc.value.namespace;
    }

    // Copy name from userspace.
    if (name_len == 0 || name_len >= TASK_NAME_MAX)
        return ZN_ERR_INVALID;

    char task_name[TASK_NAME_MAX] = {};
    if (!usercopy_read(task_name, name, name_len))
        return ZN_ERR_BAD_BUFFER;

    // Create the task sharing the caller's address space.
    // Uses sched_to_user as a placeholder entry — the real IP/SP come from thread_start.
    struct task* new_task;
    zn_status_t status = task_create(task_name, current->space, ns, sched_to_user, 0, &new_task);
    if (status != ZN_OK)
        return status;

    // Register the task as a handle in the caller's namespace.
    struct namespace_desc desc = {
        .type = NAMESPACE_DESC_TASK,
        .rights = ZN_RIGHTS_ALL,
        .value.task = new_task,
    };

    zn_handle_t handle;
    status = namespace_add_desc(current->namespace, desc, &handle);
    if (status != ZN_OK)
        return status;

    if (!usercopy_write(out, &handle, sizeof(handle))) {
        namespace_del_desc(current->namespace, handle);
        return ZN_ERR_BAD_BUFFER;
    }

    return ZN_OK;
}

zn_status_t syscall_thread_start(struct arch_context* ctx) {
    zn_handle_t thread_handle = ctx->ARCH_CTX_A0;
    uintptr_t ip = ctx->ARCH_CTX_A1;
    uintptr_t sp = ctx->ARCH_CTX_A2;
    uintptr_t arg = ctx->ARCH_CTX_A3;

    struct task* current = percpu_get()->sched.current;
    struct sched_percpu* sched = &percpu_get()->sched;

    // Look up the task.
    struct namespace_desc desc;
    zn_status_t status = namespace_get(current->namespace, thread_handle, &desc);
    if (status != ZN_OK)
        return status;
    if (desc.type != NAMESPACE_DESC_TASK)
        return ZN_ERR_BAD_HANDLE;

    struct task* task = desc.value.task;

    // Allocate a start info block for the thread entry.
    struct task_start_info* info = mem_alloc(sizeof(struct task_start_info), 0);
    if (!info)
        return ZN_ERR_NO_MEMORY;

    info->ip = ip;
    info->sp = sp;
    info->arg = arg;

    // Re-initialize the task's arch context to enter user mode via sched_to_user_arg.
    status =
        arch_task_init(&task->context, sched_to_user, (uintptr_t)info, task->kernel_stack + KERNEL_STACK_SIZE, true);
    if (status != ZN_OK) {
        mem_free(info);
        return status;
    }

    // Add to run queue.
    sched_add_task(sched, task);
    return ZN_OK;
}

zn_status_t syscall_thread_exit(struct arch_context* ctx) {
    struct sched_percpu* sched = &percpu_get()->sched;
    sched->current->state = TASK_STATE_DEAD;
    sched_yield(sched);
    unreachable();
}