#include "kpl.h"

static void *kpl_io_uring_mmap(uint32_t size, int32_t fd, off_t offset) {
    return mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_POPULATE, fd, offset);
}

int32_t kpl_io_uring_init(kpl_io_uring *io_uring, int32_t queue_size_pow2) {
    io_uring->params.flags |= IORING_SETUP_NO_SQARRAY;
    io_uring->ring_queue_entries = 1 << queue_size_pow2;
    io_uring->ring_fd = syscall(__NR_io_uring_setup, io_uring->ring_queue_entries, &io_uring->params);
    if (io_uring->ring_fd < 0)
        return -errno;
    uint32_t ssize = io_uring->params.sq_off.array + io_uring->params.sq_entries * sizeof(uint32_t);
    uint32_t csize = io_uring->params.cq_off.cqes + io_uring->params.cq_entries * sizeof(struct io_uring_cqe);
    io_uring->sq_cq_size = MAX(ssize, csize);
    io_uring->sqes_size = io_uring->params.sq_entries * sizeof(struct io_uring_sqe);
    void *sring_ptr, *cring_ptr;
    sring_ptr = cring_ptr = kpl_io_uring_mmap(io_uring->sq_cq_size, io_uring->ring_fd, IORING_OFF_SQ_RING);
    io_uring->sqes = kpl_io_uring_mmap(io_uring->sqes_size, io_uring->ring_fd, IORING_OFF_SQES);
    io_uring->sring_head = sring_ptr + io_uring->params.sq_off.head;
    io_uring->sring_tail = sring_ptr + io_uring->params.sq_off.tail;
    io_uring->sring_mask = sring_ptr + io_uring->params.sq_off.ring_mask;
    io_uring->cring_head = cring_ptr + io_uring->params.cq_off.head;
    io_uring->cring_tail = cring_ptr + io_uring->params.cq_off.tail;
    io_uring->cring_mask = cring_ptr + io_uring->params.cq_off.ring_mask;
    io_uring->cqes = cring_ptr + io_uring->params.cq_off.cqes;
    io_uring->sring_ptr = sring_ptr;
    io_uring->cring_ptr = cring_ptr;
    return 0;
}

extern inline void kpl_io_uring_free(kpl_io_uring *io_uring);

static int32_t kpl_io_uring_enter(int32_t ring_fd, uint32_t to_submit, uint32_t to_wait, uint32_t flags) {
    int32_t ret = syscall(__NR_io_uring_enter, ring_fd, to_submit, to_wait, flags, NULL, 0);
    return ret < 0 ? -errno : ret;
}

kpl_object_task *kpl_io_uring_get_completed(kpl_io_uring *io_uring) {
    if (!io_uring->submitted_to_kernel)
        return NULL;
    uint32_t cring_head = __atomic_load_n(io_uring->cring_head, __ATOMIC_ACQUIRE);
    if (cring_head == *io_uring->cring_tail) {
        if (kpl_io_uring_enter(io_uring->ring_fd, 0, 1, IORING_ENTER_GETEVENTS) < 0)
            return NULL;
    }
    struct io_uring_cqe *cqe = &io_uring->cqes[cring_head & *io_uring->cring_mask];
    kpl_object_task *task = (kpl_object_task*) cqe->user_data;
    task->io_res = cqe->res;
    task->io_flags = cqe->flags;
    cring_head++;
    __atomic_store_n(io_uring->cring_head, cring_head, __ATOMIC_RELEASE);
    io_uring->submitted_to_kernel--;
    return task;
}

int32_t kpl_io_uring_submit(kpl_io_uring *io_uring) {
    if (!io_uring->added_to_queue)
        return -1;
    int32_t ret = kpl_io_uring_enter(io_uring->ring_fd, io_uring->added_to_queue, 0, 0);
    if (ret > 0) {
        io_uring->added_to_queue -= ret;
        io_uring->submitted_to_kernel += ret;
    }
    return ret;
}

struct io_uring_sqe *kpl_io_uring_op(kpl_io_uring *io_uring, kpl_object_task *task, int32_t op, int32_t fd,
    const void *addr, uint32_t len, uint64_t offset) {
    uint32_t sring_head = __atomic_load_n(io_uring->sring_head, __ATOMIC_ACQUIRE);
    uint32_t sring_tail = *io_uring->sring_tail;
    if (sring_tail - sring_head >= (uint32_t) io_uring->ring_queue_entries)
        return NULL;
    struct io_uring_sqe *sqe = &io_uring->sqes[sring_tail & *io_uring->sring_mask];
    sring_tail++;
    __atomic_store_n(io_uring->sring_tail, sring_tail, __ATOMIC_RELEASE);
    io_uring->added_to_queue++;
    memset(sqe, 0, sizeof(struct io_uring_sqe));
    sqe->user_data = (uint64_t) task;
    sqe->opcode = op;
    sqe->fd = fd;
    sqe->off = offset;
    sqe->addr = (uint64_t) addr;
    sqe->len = len;
    return sqe;
}

extern inline bool kpl_io_uring_openat(kpl_io_uring *io_uring, kpl_object_task *task, int32_t dir_fd,
    const char *path, int32_t flags, mode_t mode);

extern inline bool kpl_io_uring_close(kpl_io_uring *io_uring, kpl_object_task *task, int32_t fd);

extern inline bool kpl_io_uring_read(kpl_io_uring *io_uring, kpl_object_task *task, int32_t fd, void *buffer,
    uint32_t byte_size, uint64_t offset);

extern inline bool kpl_io_uring_write(kpl_io_uring *uring, kpl_object_task *task, int32_t fd, void *buffer,
    uint32_t byte_size, uint64_t offset);

extern inline bool kpl_io_uring_statx(kpl_io_uring *io_uring, kpl_object_task *task, int32_t dir_fd,
    const char *path, int32_t flags, uint32_t mask, struct statx *statx_object);