// SPDX-License-Identifier: MIT

#include <zinnia/channel.h>
#include <zinnia/handle.h>
#include <zinnia/status.h>
#include <kernel/alloc.h>
#include <kernel/channel.h>
#include <kernel/namespace.h>
#include <kernel/percpu.h>
#include <kernel/sched.h>
#include <kernel/syscalls.h>
#include <kernel/usercopy.h>

zn_status_t syscall_channel_create(struct arch_context* ctx) {
    auto flags = (enum zn_channel_flags)ctx->a0;
    auto end0 = (__user zn_handle_t*)ctx->a1;
    auto end1 = (__user zn_handle_t*)ctx->a2;

    // Validate flags before doing anything.
    if (flags & !(ZN_CHANNEL_NONE | ZN_CHANNEL_NONBLOCK)) {
        return ZN_ERR_INVALID;
    }

    struct channel* ep0;
    struct channel* ep1;
    zn_status_t status = channel_create(&ep0, &ep1);
    if (status != ZN_OK)
        return status;

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

    struct namespace_desc desc0 = {
        .type = NAMESPACE_DESC_CHANNEL,
        .rights = ZN_RIGHTS_ALL,
        .value.channel = ep0,
    };
    struct namespace_desc desc1 = {
        .type = NAMESPACE_DESC_CHANNEL,
        .rights = ZN_RIGHTS_ALL,
        .value.channel = ep1,
    };

    zn_handle_t h0, h1;
    status = namespace_add_desc(current->namespace, desc0, &h0);
    if (status != ZN_OK)
        return status;

    status = namespace_add_desc(current->namespace, desc1, &h1);
    if (status != ZN_OK) {
        namespace_del_desc(current->namespace, h0);
        return status;
    }

    if (!usercopy_write(end0, &h0, sizeof(h0))) {
        namespace_del_desc(current->namespace, h0);
        namespace_del_desc(current->namespace, h1);
        return ZN_ERR_BAD_BUFFER;
    }

    if (!usercopy_write(end1, &h1, sizeof(h1))) {
        namespace_del_desc(current->namespace, h0);
        namespace_del_desc(current->namespace, h1);
        return ZN_ERR_BAD_BUFFER;
    }

    return ZN_OK;
}

zn_status_t syscall_channel_write(struct arch_context* ctx) {
    auto ch_handle = (zn_handle_t)ctx->a0;
    auto handles = (__user zn_handle_t*)ctx->a0;
    auto bytes = (__user void*)ctx->a2;
    size_t num_handles = ctx->a3;
    size_t num_bytes = ctx->a4;

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

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

    struct channel* ch = desc.value.channel;

    // Copy byte data from userspace.
    void* msg_data = nullptr;
    if (num_bytes > 0) {
        if (!bytes)
            return ZN_ERR_INVALID;

        msg_data = mem_alloc(num_bytes, ALLOC_NOZERO);
        if (!msg_data)
            return ZN_ERR_NO_MEMORY;

        if (!usercopy_read(msg_data, bytes, num_bytes)) {
            mem_free(msg_data);
            return ZN_ERR_BAD_BUFFER;
        }
    }

    // Copy and transfer handles.
    struct namespace_desc* descs = nullptr;
    if (num_handles > 0) {
        if (!handles) {
            mem_free(msg_data);
            return ZN_ERR_INVALID;
        }

        // Read handle values from userspace.
        size_t handles_size = num_handles * sizeof(zn_handle_t);
        zn_handle_t* src_handles = mem_alloc(handles_size, ALLOC_NOZERO);
        if (!src_handles) {
            mem_free(msg_data);
            return ZN_ERR_NO_MEMORY;
        }

        if (!usercopy_read(src_handles, handles, handles_size)) {
            mem_free(src_handles);
            mem_free(msg_data);
            return ZN_ERR_BAD_BUFFER;
        }

        // Resolve all handles to descriptors before removing any.
        descs = mem_alloc(num_handles * sizeof(struct namespace_desc), ALLOC_NOZERO);
        if (!descs) {
            mem_free(src_handles);
            mem_free(msg_data);
            return ZN_ERR_NO_MEMORY;
        }

        for (size_t i = 0; i < num_handles; i++) {
            status = namespace_get(current->namespace, src_handles[i], &descs[i]);
            if (status != ZN_OK) {
                mem_free(descs);
                mem_free(src_handles);
                mem_free(msg_data);
                return ZN_ERR_BAD_HANDLE;
            }
        }

        // Remove handles from sender's namespace.
        for (size_t i = 0; i < num_handles; i++) {
            namespace_del_desc(current->namespace, src_handles[i]);
        }

        mem_free(src_handles);
    }

    // Send the message through the channel.
    status = channel_write(ch, msg_data, num_bytes, descs, num_handles);
    if (status != ZN_OK) {
        mem_free(msg_data);
        mem_free(descs);
        return status;
    }

    return ZN_OK;
}

zn_status_t syscall_channel_read(struct arch_context* ctx) {
    zn_handle_t ch_handle = ctx->a0;
    __user zn_handle_t* handles = (__user zn_handle_t*)ctx->a1;
    __user void* bytes = (__user void*)ctx->a2;
    size_t num_handles = ctx->a3;
    size_t num_bytes = ctx->a4;
    __user size_t* read_handles = (__user size_t*)ctx->a5;
    __user size_t* read_bytes = (__user size_t*)ctx->a6;

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

    // Look up the channel.
    struct namespace_desc ch_desc;
    zn_status_t status = namespace_get(current->namespace, ch_handle, &ch_desc);
    if (status != ZN_OK)
        return status;
    if (ch_desc.type != NAMESPACE_DESC_CHANNEL)
        return ZN_ERR_BAD_HANDLE;

    struct channel* ch = ch_desc.value.channel;

    // Block until a message is available.
    while (true) {
        struct channel_msg* msg;
        status = channel_read(ch, &msg);

        if (status == ZN_OK) {
            // Copy byte data to userspace.
            size_t bytes_to_copy = MIN(msg->num_bytes, num_bytes);
            if (bytes_to_copy > 0 && bytes) {
                if (!usercopy_write(bytes, msg->data, bytes_to_copy)) {
                    channel_msg_free(msg);
                    return ZN_ERR_BAD_BUFFER;
                }
            }

            // Transfer handle descriptors into the reader's namespace.
            size_t handles_to_copy = MIN(msg->num_handles, num_handles);
            if (handles_to_copy > 0 && handles) {
                for (size_t i = 0; i < handles_to_copy; i++) {
                    zn_handle_t new_handle;
                    status = namespace_add_desc(current->namespace, msg->descs[i], &new_handle);
                    if (status != ZN_OK) {
                        channel_msg_free(msg);
                        return status;
                    }

                    if (!usercopy_write(&handles[i], &new_handle, sizeof(new_handle))) {
                        channel_msg_free(msg);
                        return ZN_ERR_BAD_BUFFER;
                    }
                }
            }

            // Write actual counts to userspace.
            if (read_bytes) {
                if (!usercopy_write(read_bytes, &msg->num_bytes, sizeof(msg->num_bytes))) {
                    channel_msg_free(msg);
                    return ZN_ERR_BAD_BUFFER;
                }
            }

            if (read_handles) {
                if (!usercopy_write(read_handles, &msg->num_handles, sizeof(msg->num_handles))) {
                    channel_msg_free(msg);
                    return ZN_ERR_BAD_BUFFER;
                }
            }

            channel_msg_free(msg);
            return ZN_OK;
        }

        if (status == ZN_ERR_BAD_HANDLE)
            return status;

        // Block until a message arrives or peer closes.
        sched_block(sched, &ch->waiters);
    }
}

zn_status_t syscall_channel_wait(struct arch_context* ctx) {
    zn_handle_t ch_handle = ctx->a0;

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

    struct namespace_desc desc;
    zn_status_t status = namespace_get(current->namespace, ch_handle, &desc);
    if (status != ZN_OK)
        return status;
    if (desc.type != NAMESPACE_DESC_CHANNEL)
        return ZN_ERR_BAD_HANDLE;

    struct channel* ch = desc.value.channel;

    status = channel_poll(ch);
    if (status == ZN_ERR_INVALID) {
        sched_block(sched, &ch->waiters);
        return ZN_OK;
    }

    return status;
}