#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/tailq.h>
#include <kernel/wait.h>
#include <stdatomic.h>

zn_status_t channel_create(struct channel** ep0, struct channel** ep1) {
    struct channel* a = mem_alloc(sizeof(struct channel), 0);
    if (!a)
        return ZN_ERR_NO_MEMORY;

    struct channel* b = mem_alloc(sizeof(struct channel), 0);
    if (!b) {
        mem_free(a);
        return ZN_ERR_NO_MEMORY;
    }

    a->peer = b;
    a->lock = (struct spinlock){0};
    tailq_init(&a->messages);
    wait_queue_init(&a->waiters);

    b->peer = a;
    b->lock = (struct spinlock){0};
    tailq_init(&b->messages);
    wait_queue_init(&b->waiters);

    *ep0 = a;
    *ep1 = b;
    return ZN_OK;
}

void channel_msg_free(struct channel_msg* msg) {
    mem_free(msg->data);
    mem_free(msg->descs);
    mem_free(msg);
}

zn_status_t channel_write(
    struct channel* ch,
    void* data,
    size_t num_bytes,
    struct namespace_desc* descs,
    size_t num_handles
) {
    // Check that the peer is still alive.
    struct channel* peer = atomic_load_explicit(&ch->peer, memory_order_acquire);
    if (!peer)
        return ZN_ERR_BAD_HANDLE;

    // Allocate the message.
    struct channel_msg* msg = mem_alloc(sizeof(struct channel_msg), 0);
    if (!msg)
        return ZN_ERR_NO_MEMORY;

    msg->data = data;
    msg->num_bytes = num_bytes;
    msg->descs = descs;
    msg->num_handles = num_handles;

    // Enqueue on the peer's message queue and wake a waiter.
    spin_lock(&peer->lock);
    tailq_insert_tail(&peer->messages, msg, next);
    spin_unlock(&peer->lock);

    struct sched_percpu* sched = &percpu_get()->sched;
    sched_wake(sched, &peer->waiters);

    return ZN_OK;
}

zn_status_t channel_read(struct channel* ch, struct channel_msg** msg_out) {
    spin_lock(&ch->lock);
    struct channel_msg* msg = tailq_first(&ch->messages);

    if (msg != nullptr) {
        tailq_remove(&ch->messages, msg, next);
        spin_unlock(&ch->lock);
        *msg_out = msg;
        return ZN_OK;
    }

    // No message and peer is gone.
    if (!ch->peer) {
        spin_unlock(&ch->lock);
        return ZN_ERR_BAD_HANDLE;
    }

    spin_unlock(&ch->lock);
    return ZN_ERR_INVALID;
}

zn_status_t channel_poll(struct channel* ch) {
    spin_lock(&ch->lock);

    if (!tailq_empty(&ch->messages)) {
        spin_unlock(&ch->lock);
        return ZN_OK;
    }

    if (!ch->peer) {
        spin_unlock(&ch->lock);
        return ZN_ERR_BAD_HANDLE;
    }

    spin_unlock(&ch->lock);
    return ZN_ERR_INVALID;
}