// SPDX-License-Identifier: MIT

#include <zinnia/handle.h>
#include <zinnia/status.h>
#include <common/utils.h>
#include <kernel/alloc.h>
#include <kernel/arch.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 <kernel/vas.h>
#include <kernel/vmo.h>
#include <stdatomic.h>

zn_status_t syscall_vmo_create(struct arch_context* ctx) {
    size_t length = ctx->a0;
    __user zn_handle_t* out = (__user zn_handle_t*)ctx->a1;

    // Allocate new VMO.
    struct paged_vmo* vmo;
    zn_status_t status = vmo_new_phys(&vmo);
    if (status != ZN_OK)
        return status;

    struct namespace_desc desc = {
        .type = NAMESPACE_DESC_VMO,
        .rights = ZN_RIGHTS_COMMON | ZN_RIGHT_READ | ZN_RIGHT_WRITE | ZN_RIGHT_MAP,
        .value.vmo = &vmo->object,
    };

    struct task* current = percpu_get()->sched.current;
    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)))
        return ZN_ERR_BAD_BUFFER;

    return ZN_OK;
}

zn_status_t syscall_vmo_create_phys(struct arch_context* ctx) {
    uintptr_t addr = ctx->a0;
    size_t length = ctx->a1;
    __user zn_handle_t* out = (__user zn_handle_t*)ctx->a2;

    struct vmo* vmo;
    zn_status_t status = vmo_new_phys_range(addr, length, &vmo);
    if (status != ZN_OK)
        return status;

    struct namespace_desc desc = {
        .type = NAMESPACE_DESC_VMO,
        .rights = ZN_RIGHTS_COMMON | ZN_RIGHT_READ | ZN_RIGHT_WRITE | ZN_RIGHT_MAP,
        .value.vmo = vmo,
    };

    struct task* current = percpu_get()->sched.current;
    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)))
        return ZN_ERR_BAD_BUFFER;

    return ZN_OK;
}

zn_status_t syscall_vmo_map(struct arch_context* ctx) {
    zn_handle_t vmo_handle = ctx->a0;
    zn_handle_t vas_handle = ctx->a1;
    uintptr_t vmo_offset = ctx->a2;
    auto addr = (__user uintptr_t*)ctx->a3;
    size_t bytes = ctx->a4;
    enum zn_vm_flags flags = ctx->a5;

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

    // Get VAS.
    struct vas* vas;
    if (vas_handle == ZN_HANDLE_THIS_VAS) {
        vas = current->space;
    } else {
        struct namespace_desc vas_desc;
        zn_status_t status = namespace_get(current->namespace, vmo_handle, &vas_desc);
        if (status != ZN_OK)
            return status;
        if (vas_desc.type != NAMESPACE_DESC_VAS)
            return ZN_ERR_BAD_HANDLE;
        vas = vas_desc.value.vas;
    }

    // Get VMO.
    struct namespace_desc vmo_desc;
    zn_status_t status = namespace_get(current->namespace, vmo_handle, &vmo_desc);
    if (status != ZN_OK)
        return status;
    if (vmo_desc.type != NAMESPACE_DESC_VMO)
        return ZN_ERR_BAD_HANDLE;
    struct vmo* vmo = vmo_desc.value.vmo;

    // Determine mapping address.
    uintptr_t target_addr = 0;
    if (addr != nullptr) {
        if (!usercopy_read(&target_addr, addr, sizeof(target_addr)))
            return ZN_ERR_BAD_BUFFER;
    }

    // TODO: Use actual allocator for user task.
    if (target_addr == 0) {
        static _Atomic uintptr_t map_offset = 0xA000'0000;
        target_addr = atomic_fetch_add((volatile uintptr_t*)&map_offset, ALIGN_UP(bytes, arch_mem_page_size()));
    }

    if (addr != nullptr) {
        if (!usercopy_write(addr, &target_addr, sizeof(target_addr)))
            return ZN_ERR_BAD_BUFFER;
    }
    vas_map_vmo(vas, vmo, target_addr, bytes, flags, vmo_offset);

    return ZN_OK;
}