#include "uv_writer.h"

#include <string.h>
#include <unistd.h>

#include "../include/raft.h"
#include "assert.h"
#include "heap.h"

/* Copy the error message from the request object to the writer object. */
static void uvWriterReqTransferErrMsg(struct UvWriterReq *req)
{
    ErrMsgPrintf(req->writer->errmsg, "%s", req->errmsg);
}

/* Set the request status according the given result code. */
static void uvWriterReqSetStatus(struct UvWriterReq *req, int result)
{
    if (result < 0) {
        ErrMsgPrintf(req->errmsg, "write failed: %d", result);
        req->status = RAFT_IOERR;
    } else if ((size_t)result < req->len) {
        ErrMsgPrintf(req->errmsg, "short write: %d bytes instead of %zu",
                     result, req->len);
        req->status = RAFT_NOSPACE;
    } else {
        req->status = 0;
    }
}

/* Remove the request from the queue of inflight writes and invoke the request
 * callback if set. */
static void uvWriterReqFinish(struct UvWriterReq *req)
{
    QUEUE_REMOVE(&req->queue);
    if (req->status != 0) {
        uvWriterReqTransferErrMsg(req);
    }
    req->cb(req, req->status);
}

/* Wrapper around the low-level OS syscall, providing a better error message. */
static int uvWriterIoSetup(unsigned n, aio_context_t *ctx, char *errmsg)
{
    int rv;
    rv = UvOsIoSetup(n, ctx);
    if (rv != 0) {
        switch (rv) {
            case UV_EAGAIN:
                ErrMsgPrintf(errmsg, "AIO events user limit exceeded");
                rv = RAFT_TOOMANY;
                break;
            default:
                UvOsErrMsg(errmsg, "io_setup", rv);
                rv = RAFT_IOERR;
                break;
        }
        return rv;
    }
    return 0;
}

/* Run blocking syscalls involved in a file write request.
 *
 * Perform a KAIO write request and synchronously wait for it to complete. */
static void uvWriterWorkCb(uv_work_t *work)
{
    struct UvWriterReq *req; /* Writer request object */
    struct UvWriter *w;      /* Writer object */
    aio_context_t ctx;       /* KAIO handle */
    struct iocb *iocbs;      /* Pointer to KAIO request object */
    struct io_event event;   /* KAIO response object */
    int n_events;
    int rv;

    req = work->data;
    w = req->writer;

    iocbs = &req->iocb;

    /* If more than one write in parallel is allowed, submit the AIO request
     * using a dedicated context, to avoid synchronization issues between
     * threads when multiple writes are submitted in parallel. This is
     * suboptimal but in real-world users should use file systems and kernels
     * with proper async write support. */
    if (w->n_events > 1) {
        ctx = 0;
        rv = uvWriterIoSetup(1 /* Maximum concurrent requests */, &ctx,
                             req->errmsg);
        if (rv != 0) {
            goto out;
        }
    } else {
        ctx = w->ctx;
    }

    /* Submit the request */
    rv = UvOsIoSubmit(ctx, 1, &iocbs);
    if (rv != 0) {
        /* UNTESTED: since we're not using NOWAIT and the parameters are valid,
         * this shouldn't fail. */
        UvOsErrMsg(req->errmsg, "io_submit", rv);
        rv = RAFT_IOERR;
        goto out_after_io_setup;
    }

    /* Wait for the request to complete */
    n_events = UvOsIoGetevents(ctx, 1, 1, &event, NULL);
    assert(n_events == 1);
    if (n_events != 1) {
        /* UNTESTED */
        rv = n_events >= 0 ? -1 : n_events;
    }

out_after_io_setup:
    if (w->n_events > 1) {
        UvOsIoDestroy(ctx);
    }

out:
    if (rv != 0) {
        req->status = rv;
    } else {
        uvWriterReqSetStatus(req, (int)event.res);
    }

    return;
}

/* Callback run after writeWorkCb has returned. It normally invokes the write
 * request callback. */
static void uvWriterAfterWorkCb(uv_work_t *work, int status)
{
    struct UvWriterReq *req = work->data; /* Write file request object */
    assert(status == 0);                  /* We don't cancel worker requests */
    uvWriterReqFinish(req);
}

/* Callback fired when the event fd associated with AIO write requests should be
 * ready for reading (i.e. when a write has completed). */
static void uvWriterPollCb(uv_poll_t *poller, int status, int events)
{
    struct UvWriter *w = poller->data;
    uint64_t completed; /* True if the write is complete */
    unsigned i;
    int n_events;
    int rv;

    assert(w->event_fd >= 0);
    assert(status == 0);
    if (status != 0) {
        /* UNTESTED libuv docs: If an error happens while polling, status will be < 0 and
         * corresponds with one of the UV_E* error codes. */
        goto fail_requests;
    }

    assert(events & UV_READABLE);

    /* Read the event file descriptor */
    rv = (int)read(w->event_fd, &completed, sizeof completed);
    if (rv != sizeof completed) {
        /* UNTESTED: According to eventfd(2) this is the only possible failure
         * mode, meaning that epoll has indicated that the event FD is not yet
         * ready. */
        assert(errno == EAGAIN);
        return;
    }

    /* TODO: this assertion fails in unit tests */
    /* assert(completed == 1); */

    /* Try to fetch the write responses.
     *
     * If we got here at least one write should have completed and io_events
     * should return immediately without blocking. */
    n_events = UvOsIoGetevents(w->ctx, 1, (long int)w->n_events, w->events, NULL);
    assert(n_events >= 1);
    if (n_events < 1) {
        /* UNTESTED */
        status = n_events == 0 ? -1 : n_events;
        goto fail_requests;
    }

    for (i = 0; i < (unsigned)n_events; i++) {
        struct io_event *event = &w->events[i];
        struct UvWriterReq *req = *((void **)&event->data);

#if defined(RWF_NOWAIT)
        /* If we got EAGAIN, it means it was not possible to perform the write
         * asynchronously, so let's fall back to the threadpool. */
        if (event->res == -EAGAIN) {
            req->iocb.aio_flags &= (unsigned)~IOCB_FLAG_RESFD;
            req->iocb.aio_resfd = 0;
            req->iocb.aio_rw_flags &= ~RWF_NOWAIT;
            assert(req->work.data == NULL);
            req->work.data = req;
            rv = uv_queue_work(w->loop, &req->work, uvWriterWorkCb,
                               uvWriterAfterWorkCb);
            if (rv != 0) {
                /* UNTESTED: with the current libuv implementation this should
                 * never fail. */
                UvOsErrMsg(req->errmsg, "uv_queue_work", rv);
                req->status = RAFT_IOERR;
                goto finish;
            }
            return;
        }
#endif /* RWF_NOWAIT */

        uvWriterReqSetStatus(req, (int)event->res);

#if defined(RWF_NOWAIT)
    finish:
#endif /* RWF_NOWAIT */
        uvWriterReqFinish(req);
    }

    return;

fail_requests:
    while (!QUEUE_IS_EMPTY(&w->poll_queue)) {
        queue *head;
        struct UvWriterReq *req;
        head = QUEUE_HEAD(&w->poll_queue);
        req = QUEUE_DATA(head, struct UvWriterReq, queue);
        uvWriterReqSetStatus(req, status);
        uvWriterReqFinish(req);
    }
}

int UvWriterInit(struct UvWriter *w,
                 struct uv_loop_s *loop,
                 uv_file fd,
                 bool direct /* Whether to use direct I/O */,
                 bool async /* Whether async I/O is available */,
                 unsigned max_concurrent_writes,
                 char *errmsg)
{
    void *data = w->data;
    int rv = 0;
    memset(w, 0, sizeof *w);
    w->data = data;
    w->loop = loop;
    w->fd = fd;
    w->async = async;
    w->ctx = 0;
    w->events = NULL;
    w->n_events = max_concurrent_writes;
    w->event_fd = -1;
    w->event_poller.data = NULL;
    w->check.data = NULL;
    w->close_cb = NULL;
    QUEUE_INIT(&w->poll_queue);
    QUEUE_INIT(&w->work_queue);
    w->closing = false;
    w->errmsg = errmsg;

    /* Set direct I/O if available. */
    if (direct) {
        rv = UvOsSetDirectIo(w->fd);
        if (rv != 0) {
            UvOsErrMsg(errmsg, "fcntl", rv);
            goto err;
        }
    }

    /* Setup the AIO context. */
    rv = uvWriterIoSetup(w->n_events, &w->ctx, errmsg);
    if (rv != 0) {
        goto err;
    }

    /* Initialize the array of re-usable event objects. */
    w->events = RaftHeapCalloc(w->n_events, sizeof *w->events);
    if (w->events == NULL) {
        /* UNTESTED: todo */
        ErrMsgOom(errmsg);
        rv = RAFT_NOMEM;
        goto err_after_io_setup;
    }

    /* Create an event file descriptor to get notified when a write has
     * completed. */
    rv = UvOsEventfd(0, UV_FS_O_NONBLOCK);
    if (rv < 0) {
        /* UNTESTED: should fail only with ENOMEM */
        UvOsErrMsg(errmsg, "eventfd", rv);
        rv = RAFT_IOERR;
        goto err_after_events_alloc;
    }
    w->event_fd = rv;

    rv = uv_poll_init(loop, &w->event_poller, w->event_fd);
    if (rv != 0) {
        /* UNTESTED: with the current libuv implementation this should never
         * fail. */
        UvOsErrMsg(errmsg, "uv_poll_init", rv);
        rv = RAFT_IOERR;
        goto err_after_event_fd;
    }
    w->event_poller.data = w;

    rv = uv_check_init(loop, &w->check);
    if (rv != 0) {
        /* UNTESTED: with the current libuv implementation this should never
         * fail. */
        UvOsErrMsg(errmsg, "uv_check_init", rv);
        rv = RAFT_IOERR;
        goto err_after_event_fd;
    }
    w->check.data = w;

    rv = uv_poll_start(&w->event_poller, UV_READABLE, uvWriterPollCb);
    if (rv != 0) {
        /* UNTESTED: with the current libuv implementation this should never
         * fail. */
        UvOsErrMsg(errmsg, "uv_poll_start", rv);
        rv = RAFT_IOERR;
        goto err_after_event_fd;
    }

    return 0;

err_after_event_fd:
    UvOsClose(w->event_fd);
err_after_events_alloc:
    RaftHeapFree(w->events);
err_after_io_setup:
    UvOsIoDestroy(w->ctx);
err:
    assert(rv != 0);
    return rv;
}

static void uvWriterCleanUpAndFireCloseCb(struct UvWriter *w)
{
    assert(w->closing);

    UvOsClose(w->fd);
    RaftHeapFree(w->events);
    UvOsIoDestroy(w->ctx);

    if (w->close_cb != NULL) {
        w->close_cb(w);
    }
}

static void uvWriterPollerCloseCb(struct uv_handle_s *handle)
{
    struct UvWriter *w = handle->data;
    w->event_poller.data = NULL;

    /* Cancel all pending requests. */
    while (!QUEUE_IS_EMPTY(&w->poll_queue)) {
        queue *head;
        struct UvWriterReq *req;
        head = QUEUE_HEAD(&w->poll_queue);
        req = QUEUE_DATA(head, struct UvWriterReq, queue);
        assert(req->work.data == NULL);
        req->status = RAFT_CANCELED;
        uvWriterReqFinish(req);
    }

    if (w->check.data != NULL) {
        return;
    }

    uvWriterCleanUpAndFireCloseCb(w);
}

static void uvWriterCheckCloseCb(struct uv_handle_s *handle)
{
    struct UvWriter *w = handle->data;
    w->check.data = NULL;
    if (w->event_poller.data != NULL) {
        return;
    }
    uvWriterCleanUpAndFireCloseCb(w);
}

static void uvWriterCheckCb(struct uv_check_s *check)
{
    struct UvWriter *w = check->data;
    if (!QUEUE_IS_EMPTY(&w->work_queue)) {
        return;
    }
    uv_close((struct uv_handle_s *)&w->check, uvWriterCheckCloseCb);
}

void UvWriterClose(struct UvWriter *w, UvWriterCloseCb cb)
{
    int rv;
    assert(!w->closing);
    w->closing = true;
    w->close_cb = cb;

    /* We can close the event file descriptor right away, but we shouldn't close
     * the main file descriptor or destroy the AIO context since there might be
     * threadpool requests in flight. */
    UvOsClose(w->event_fd);

    rv = uv_poll_stop(&w->event_poller);
    assert(rv == 0); /* Can this ever fail? */

    uv_close((struct uv_handle_s *)&w->event_poller, uvWriterPollerCloseCb);

    /* If we have requests executing in the threadpool, we need to wait for
     * them. That's done in the check callback. */
    if (!QUEUE_IS_EMPTY(&w->work_queue)) {
        uv_check_start(&w->check, uvWriterCheckCb);
    } else {
        uv_close((struct uv_handle_s *)&w->check, uvWriterCheckCloseCb);
    }
}

/* Return the total lengths of the given buffers. */
static size_t lenOfBufs(const uv_buf_t bufs[], unsigned n)
{
    size_t len = 0;
    unsigned i;
    for (i = 0; i < n; i++) {
        len += bufs[i].len;
    }
    return len;
}

int UvWriterSubmit(struct UvWriter *w,
                   struct UvWriterReq *req,
                   const uv_buf_t bufs[],
                   unsigned n,
                   size_t offset,
                   UvWriterReqCb cb)
{
    int rv = 0;
#if defined(RWF_NOWAIT)
    struct iocb *iocbs = &req->iocb;
#endif /* RWF_NOWAIT */
    assert(!w->closing);

    /* TODO: at the moment we are not leveraging the support for concurrent
     *       writes, so ensure that we're getting write requests
     *       sequentially. */
    if (w->n_events == 1) {
        assert(QUEUE_IS_EMPTY(&w->poll_queue));
        assert(QUEUE_IS_EMPTY(&w->work_queue));
    }

    assert(w->fd >= 0);
    assert(w->event_fd >= 0);
    assert(w->ctx != 0);
    assert(req != NULL);
    assert(bufs != NULL);
    assert(n > 0);

    req->writer = w;
    req->len = lenOfBufs(bufs, n);
    req->status = -1;
    req->work.data = NULL;
    req->cb = cb;
    memset(&req->iocb, 0, sizeof req->iocb);
    memset(req->errmsg, 0, sizeof req->errmsg);

    req->iocb.aio_fildes = (uint32_t)w->fd;
    req->iocb.aio_lio_opcode = IOCB_CMD_PWRITEV;
    req->iocb.aio_reqprio = 0;
    *((void **)(&req->iocb.aio_buf)) = (void *)bufs;
    req->iocb.aio_nbytes = n;
    req->iocb.aio_offset = (int64_t)offset;
    *((void **)(&req->iocb.aio_data)) = (void *)req;

#if defined(RWF_HIPRI)
    /* High priority request, if possible */
    /* TODO: do proper kernel feature detection for this one. */
    /* req->iocb.aio_rw_flags |= RWF_HIPRI; */
#endif

#if defined(RWF_DSYNC)
    /* Use per-request synchronous I/O if available. Otherwise, we have opened
     * the file with O_DSYNC. */
    /* TODO: do proper kernel feature detection for this one. */
    /* req->iocb.aio_rw_flags |= RWF_DSYNC; */
#endif

#if defined(RWF_NOWAIT)
    /* If io_submit can be run in a 100% non-blocking way, we'll try to write
     * without using the threadpool. */
    if (w->async) {
        req->iocb.aio_flags |= IOCB_FLAG_RESFD;
        req->iocb.aio_resfd = (uint32_t)w->event_fd;
        req->iocb.aio_rw_flags |= RWF_NOWAIT;
    }
#else
    /* Since there's no support for NOWAIT, io_submit might occasionally block
     * and we need to run it in the threadpool. */
    assert(w->async == false);
#endif /* RWF_NOWAIT */

#if defined(RWF_NOWAIT)
    /* Try to submit the write request asynchronously */
    if (w->async) {
        QUEUE_PUSH(&w->poll_queue, &req->queue);
        rv = UvOsIoSubmit(w->ctx, 1, &iocbs);

        /* If no error occurred, we're done, the write request was
         * submitted. */
        if (rv == 0) {
            goto done;
        }

        QUEUE_REMOVE(&req->queue);

        /* Check the reason of the error. */
        switch (rv) {
            case UV_EAGAIN:
                break;
            default:
                /* Unexpected error */
                UvOsErrMsg(w->errmsg, "io_submit", rv);
                rv = RAFT_IOERR;
                goto err;
        }

        /* Submitting the write would block, or NOWAIT is not
         * supported. Let's run this request in the threadpool. */
        req->iocb.aio_flags &= (unsigned)~IOCB_FLAG_RESFD;
        req->iocb.aio_resfd = 0;
        req->iocb.aio_rw_flags &= ~RWF_NOWAIT;
    }
#endif /* RWF_NOWAIT */

    /* If we got here it means we need to run io_submit in the threadpool. */
    QUEUE_PUSH(&w->work_queue, &req->queue);
    req->work.data = req;
    rv =
        uv_queue_work(w->loop, &req->work, uvWriterWorkCb, uvWriterAfterWorkCb);
    if (rv != 0) {
        /* UNTESTED: with the current libuv implementation this can't fail. */
        req->work.data = NULL;
        QUEUE_REMOVE(&req->queue);
        UvOsErrMsg(w->errmsg, "uv_queue_work", rv);
        rv = RAFT_IOERR;
        goto err;
    }

#if defined(RWF_NOWAIT)
done:
#endif /* RWF_NOWAIT */
    return 0;

err:
    assert(rv != 0);
    return rv;
}
