/* SPDX-License-Identifier: MIT */
/*
 * Description: UBLK_F_BATCH_IO buffer management
 */

#include "kublk.h"

static inline void *ublk_get_commit_buf(struct ublk_thread *t,
                                        unsigned short buf_idx)
{
        unsigned idx;

        if (buf_idx < t->commit_buf_start ||
                        buf_idx >= t->commit_buf_start + t->nr_commit_buf)
                return NULL;
        idx = buf_idx - t->commit_buf_start;
        return t->commit_buf + idx * t->commit_buf_size;
}

/*
 * Allocate one buffer for UBLK_U_IO_PREP_IO_CMDS or UBLK_U_IO_COMMIT_IO_CMDS
 *
 * Buffer index is returned.
 */
static inline unsigned short ublk_alloc_commit_buf(struct ublk_thread *t)
{
        int idx = allocator_get(&t->commit_buf_alloc);

        if (idx >= 0)
                return  idx + t->commit_buf_start;
        return UBLKS_T_COMMIT_BUF_INV_IDX;
}

/*
 * Free one commit buffer which is used by UBLK_U_IO_PREP_IO_CMDS or
 * UBLK_U_IO_COMMIT_IO_CMDS
 */
static inline void ublk_free_commit_buf(struct ublk_thread *t,
                                         unsigned short i)
{
        unsigned short idx = i - t->commit_buf_start;

        ublk_assert(idx < t->nr_commit_buf);
        ublk_assert(allocator_get_val(&t->commit_buf_alloc, idx) != 0);

        allocator_put(&t->commit_buf_alloc, idx);
}

static unsigned char ublk_commit_elem_buf_size(struct ublk_dev *dev)
{
        if (dev->dev_info.flags & (UBLK_F_SUPPORT_ZERO_COPY | UBLK_F_USER_COPY |
                                UBLK_F_AUTO_BUF_REG))
                return 8;

        /* one extra 8bytes for carrying buffer address */
        return 16;
}

static unsigned ublk_commit_buf_size(struct ublk_thread *t)
{
        struct ublk_dev *dev = t->dev;
        unsigned elem_size = ublk_commit_elem_buf_size(dev);
        unsigned int total = elem_size * dev->dev_info.queue_depth;
        unsigned int page_sz = getpagesize();

        return round_up(total, page_sz);
}

static void free_batch_commit_buf(struct ublk_thread *t)
{
        if (t->commit_buf) {
                unsigned buf_size = ublk_commit_buf_size(t);
                unsigned int total = buf_size * t->nr_commit_buf;

                munlock(t->commit_buf, total);
                free(t->commit_buf);
        }
        allocator_deinit(&t->commit_buf_alloc);
        free(t->commit);
}

static int alloc_batch_commit_buf(struct ublk_thread *t)
{
        unsigned buf_size = ublk_commit_buf_size(t);
        unsigned int total = buf_size * t->nr_commit_buf;
        unsigned int page_sz = getpagesize();
        void *buf = NULL;
        int i, ret, j = 0;

        t->commit = calloc(t->nr_queues, sizeof(*t->commit));
        for (i = 0; i < t->dev->dev_info.nr_hw_queues; i++) {
                if (t->q_map[i])
                        t->commit[j++].q_id = i;
        }

        allocator_init(&t->commit_buf_alloc, t->nr_commit_buf);

        t->commit_buf = NULL;
        ret = posix_memalign(&buf, page_sz, total);
        if (ret || !buf)
                goto fail;

        t->commit_buf = buf;

        /* lock commit buffer pages for fast access */
        if (mlock(t->commit_buf, total))
                ublk_err("%s: can't lock commit buffer %s\n", __func__,
                        strerror(errno));

        return 0;

fail:
        free_batch_commit_buf(t);
        return ret;
}

static unsigned int ublk_thread_nr_queues(const struct ublk_thread *t)
{
        int i;
        int ret = 0;

        for (i = 0; i < t->dev->dev_info.nr_hw_queues; i++)
                ret += !!t->q_map[i];

        return ret;
}

void ublk_batch_prepare(struct ublk_thread *t)
{
        /*
         * We only handle single device in this thread context.
         *
         * All queues have same feature flags, so use queue 0's for
         * calculate uring_cmd flags.
         *
         * This way looks not elegant, but it works so far.
         */
        struct ublk_queue *q = &t->dev->q[0];

        /* cache nr_queues because we don't support dynamic load-balance yet */
        t->nr_queues = ublk_thread_nr_queues(t);

        t->commit_buf_elem_size = ublk_commit_elem_buf_size(t->dev);
        t->commit_buf_size = ublk_commit_buf_size(t);
        t->commit_buf_start = t->nr_bufs;
        t->nr_commit_buf = 2 * t->nr_queues;
        t->nr_bufs += t->nr_commit_buf;

        t->cmd_flags = 0;
        if (ublk_queue_use_auto_zc(q)) {
                if (ublk_queue_auto_zc_fallback(q))
                        t->cmd_flags |= UBLK_BATCH_F_AUTO_BUF_REG_FALLBACK;
        } else if (!ublk_queue_no_buf(q))
                t->cmd_flags |= UBLK_BATCH_F_HAS_BUF_ADDR;

        t->state |= UBLKS_T_BATCH_IO;

        ublk_log("%s: thread %d commit(nr_bufs %u, buf_size %u, start %u)\n",
                        __func__, t->idx,
                        t->nr_commit_buf, t->commit_buf_size,
                        t->nr_bufs);
}

static void free_batch_fetch_buf(struct ublk_thread *t)
{
        int i;

        for (i = 0; i < t->nr_fetch_bufs; i++) {
                io_uring_free_buf_ring(&t->ring, t->fetch[i].br, 1, i);
                munlock(t->fetch[i].fetch_buf, t->fetch[i].fetch_buf_size);
                free(t->fetch[i].fetch_buf);
        }
        free(t->fetch);
}

static int alloc_batch_fetch_buf(struct ublk_thread *t)
{
        /* page aligned fetch buffer, and it is mlocked for speedup delivery */
        unsigned pg_sz = getpagesize();
        unsigned buf_size = round_up(t->dev->dev_info.queue_depth * 2, pg_sz);
        int ret;
        int i = 0;

        /* double fetch buffer for each queue */
        t->nr_fetch_bufs = t->nr_queues * 2;
        t->fetch = calloc(t->nr_fetch_bufs, sizeof(*t->fetch));

        /* allocate one buffer for each queue */
        for (i = 0; i < t->nr_fetch_bufs; i++) {
                t->fetch[i].fetch_buf_size = buf_size;

                if (posix_memalign((void **)&t->fetch[i].fetch_buf, pg_sz,
                                        t->fetch[i].fetch_buf_size))
                        return -ENOMEM;

                /* lock fetch buffer page for fast fetching */
                if (mlock(t->fetch[i].fetch_buf, t->fetch[i].fetch_buf_size))
                        ublk_err("%s: can't lock fetch buffer %s\n", __func__,
                                strerror(errno));
                t->fetch[i].br = io_uring_setup_buf_ring(&t->ring, 1,
                        i, IOU_PBUF_RING_INC, &ret);
                if (!t->fetch[i].br) {
                        ublk_err("Buffer ring register failed %d\n", ret);
                        return ret;
                }
        }

        return 0;
}

int ublk_batch_alloc_buf(struct ublk_thread *t)
{
        int ret;

        ublk_assert(t->nr_commit_buf < 2 * UBLK_MAX_QUEUES);

        ret = alloc_batch_commit_buf(t);
        if (ret)
                return ret;
        return alloc_batch_fetch_buf(t);
}

void ublk_batch_free_buf(struct ublk_thread *t)
{
        free_batch_commit_buf(t);
        free_batch_fetch_buf(t);
}

static void ublk_init_batch_cmd(struct ublk_thread *t, __u16 q_id,
                                struct io_uring_sqe *sqe, unsigned op,
                                unsigned short elem_bytes,
                                unsigned short nr_elem,
                                unsigned short buf_idx)
{
        struct ublk_batch_io *cmd;
        __u64 user_data;

        cmd = (struct ublk_batch_io *)ublk_get_sqe_cmd(sqe);

        ublk_set_sqe_cmd_op(sqe, op);

        sqe->fd = 0;    /* dev->fds[0] */
        sqe->opcode     = IORING_OP_URING_CMD;
        sqe->flags      = IOSQE_FIXED_FILE;

        cmd->q_id       = q_id;
        cmd->flags      = 0;
        cmd->reserved   = 0;
        cmd->elem_bytes = elem_bytes;
        cmd->nr_elem    = nr_elem;

        user_data = build_user_data(buf_idx, _IOC_NR(op), nr_elem, q_id, 0);
        io_uring_sqe_set_data64(sqe, user_data);

        t->cmd_inflight += 1;

        ublk_dbg(UBLK_DBG_IO_CMD, "%s: thread %u qid %d cmd_op %x data %lx "
                        "nr_elem %u elem_bytes %u buf_size %u buf_idx %d "
                        "cmd_inflight %u\n",
                        __func__, t->idx, q_id, op, user_data,
                        cmd->nr_elem, cmd->elem_bytes,
                        nr_elem * elem_bytes, buf_idx, t->cmd_inflight);
}

static void ublk_setup_commit_sqe(struct ublk_thread *t,
                                  struct io_uring_sqe *sqe,
                                  unsigned short buf_idx)
{
        struct ublk_batch_io *cmd;

        cmd = (struct ublk_batch_io *)ublk_get_sqe_cmd(sqe);

        /* Use plain user buffer instead of fixed buffer */
        cmd->flags |= t->cmd_flags;
}

static void ublk_batch_queue_fetch(struct ublk_thread *t,
                                   struct ublk_queue *q,
                                   unsigned short buf_idx)
{
        unsigned short nr_elem = t->fetch[buf_idx].fetch_buf_size / 2;
        struct io_uring_sqe *sqe;

        io_uring_buf_ring_add(t->fetch[buf_idx].br, t->fetch[buf_idx].fetch_buf,
                        t->fetch[buf_idx].fetch_buf_size,
                        0, 0, 0);
        io_uring_buf_ring_advance(t->fetch[buf_idx].br, 1);

        ublk_io_alloc_sqes(t, &sqe, 1);

        ublk_init_batch_cmd(t, q->q_id, sqe, UBLK_U_IO_FETCH_IO_CMDS, 2, nr_elem,
                        buf_idx);

        sqe->rw_flags= IORING_URING_CMD_MULTISHOT;
        sqe->buf_group = buf_idx;
        sqe->flags |= IOSQE_BUFFER_SELECT;

        t->fetch[buf_idx].fetch_buf_off = 0;
}

void ublk_batch_start_fetch(struct ublk_thread *t)
{
        int i;
        int j = 0;

        for (i = 0; i < t->dev->dev_info.nr_hw_queues; i++) {
                if (t->q_map[i]) {
                        struct ublk_queue *q = &t->dev->q[i];

                        /* submit two fetch commands for each queue */
                        ublk_batch_queue_fetch(t, q, j++);
                        ublk_batch_queue_fetch(t, q, j++);
                }
        }
}

static unsigned short ublk_compl_batch_fetch(struct ublk_thread *t,
                                   struct ublk_queue *q,
                                   const struct io_uring_cqe *cqe)
{
        unsigned short buf_idx = user_data_to_tag(cqe->user_data);
        unsigned start = t->fetch[buf_idx].fetch_buf_off;
        unsigned end = start + cqe->res;
        void *buf = t->fetch[buf_idx].fetch_buf;
        int i;

        if (cqe->res < 0)
                return buf_idx;

       if ((end - start) / 2 > q->q_depth) {
               ublk_err("%s: fetch duplicated ios offset %u count %u\n", __func__, start, cqe->res);

               for (i = start; i < end; i += 2) {
                       unsigned short tag = *(unsigned short *)(buf + i);

                       ublk_err("%u ", tag);
               }
               ublk_err("\n");
       }

        for (i = start; i < end; i += 2) {
                unsigned short tag = *(unsigned short *)(buf + i);

                if (tag >= q->q_depth)
                        ublk_err("%s: bad tag %u\n", __func__, tag);

                if (q->tgt_ops->queue_io)
                        q->tgt_ops->queue_io(t, q, tag);
        }
        t->fetch[buf_idx].fetch_buf_off = end;
        return buf_idx;
}

static int __ublk_batch_queue_prep_io_cmds(struct ublk_thread *t, struct ublk_queue *q)
{
        unsigned short nr_elem = q->q_depth;
        unsigned short buf_idx = ublk_alloc_commit_buf(t);
        struct io_uring_sqe *sqe;
        void *buf;
        int i;

        ublk_assert(buf_idx != UBLKS_T_COMMIT_BUF_INV_IDX);

        ublk_io_alloc_sqes(t, &sqe, 1);

        ublk_assert(nr_elem == q->q_depth);
        buf = ublk_get_commit_buf(t, buf_idx);
        for (i = 0; i < nr_elem; i++) {
                struct ublk_batch_elem *elem = (struct ublk_batch_elem *)(
                                buf + i * t->commit_buf_elem_size);
                struct ublk_io *io = &q->ios[i];

                elem->tag = i;
                elem->result = 0;

                if (ublk_queue_use_auto_zc(q))
                        elem->buf_index = ublk_batch_io_buf_idx(t, q, i);
                else if (!ublk_queue_no_buf(q))
                        elem->buf_addr = (__u64)io->buf_addr;
        }

        sqe->addr = (__u64)buf;
        sqe->len = t->commit_buf_elem_size * nr_elem;

        ublk_init_batch_cmd(t, q->q_id, sqe, UBLK_U_IO_PREP_IO_CMDS,
                        t->commit_buf_elem_size, nr_elem, buf_idx);
        ublk_setup_commit_sqe(t, sqe, buf_idx);
        return 0;
}

int ublk_batch_queue_prep_io_cmds(struct ublk_thread *t, struct ublk_queue *q)
{
        int ret = 0;

        pthread_spin_lock(&q->lock);
        if (q->flags & UBLKS_Q_PREPARED)
                goto unlock;
        ret = __ublk_batch_queue_prep_io_cmds(t, q);
        if (!ret)
                q->flags |= UBLKS_Q_PREPARED;
unlock:
        pthread_spin_unlock(&q->lock);

        return ret;
}

static void ublk_batch_compl_commit_cmd(struct ublk_thread *t,
                                        const struct io_uring_cqe *cqe,
                                        unsigned op)
{
        unsigned short buf_idx = user_data_to_tag(cqe->user_data);

        if (op == _IOC_NR(UBLK_U_IO_PREP_IO_CMDS))
                ublk_assert(cqe->res == 0);
        else if (op == _IOC_NR(UBLK_U_IO_COMMIT_IO_CMDS)) {
                int nr_elem = user_data_to_tgt_data(cqe->user_data);

                ublk_assert(cqe->res == t->commit_buf_elem_size * nr_elem);
        } else
                ublk_assert(0);

        ublk_free_commit_buf(t, buf_idx);
}

void ublk_batch_compl_cmd(struct ublk_thread *t,
                          const struct io_uring_cqe *cqe)
{
        unsigned op = user_data_to_op(cqe->user_data);
        struct ublk_queue *q;
        unsigned buf_idx;
        unsigned q_id;

        if (op == _IOC_NR(UBLK_U_IO_PREP_IO_CMDS) ||
                        op == _IOC_NR(UBLK_U_IO_COMMIT_IO_CMDS)) {
                t->cmd_inflight--;
                ublk_batch_compl_commit_cmd(t, cqe, op);
                return;
        }

        /* FETCH command is per queue */
        q_id = user_data_to_q_id(cqe->user_data);
        q = &t->dev->q[q_id];
        buf_idx = ublk_compl_batch_fetch(t, q, cqe);

        if (cqe->res < 0 && cqe->res != -ENOBUFS) {
                t->cmd_inflight--;
                t->state |= UBLKS_T_STOPPING;
        } else if (!(cqe->flags & IORING_CQE_F_MORE) || cqe->res == -ENOBUFS) {
                t->cmd_inflight--;
                ublk_batch_queue_fetch(t, q, buf_idx);
        }
}

static void __ublk_batch_commit_io_cmds(struct ublk_thread *t,
                                        struct batch_commit_buf *cb)
{
        struct io_uring_sqe *sqe;
        unsigned short buf_idx;
        unsigned short nr_elem = cb->done;

        /* nothing to commit */
        if (!nr_elem) {
                ublk_free_commit_buf(t, cb->buf_idx);
                return;
        }

        ublk_io_alloc_sqes(t, &sqe, 1);
        buf_idx = cb->buf_idx;
        sqe->addr = (__u64)cb->elem;
        sqe->len = nr_elem * t->commit_buf_elem_size;

        /* commit isn't per-queue command */
        ublk_init_batch_cmd(t, cb->q_id, sqe, UBLK_U_IO_COMMIT_IO_CMDS,
                        t->commit_buf_elem_size, nr_elem, buf_idx);
        ublk_setup_commit_sqe(t, sqe, buf_idx);
}

void ublk_batch_commit_io_cmds(struct ublk_thread *t)
{
        int i;

        for (i = 0; i < t->nr_queues; i++) {
                struct batch_commit_buf *cb = &t->commit[i];

                if (cb->buf_idx != UBLKS_T_COMMIT_BUF_INV_IDX)
                        __ublk_batch_commit_io_cmds(t, cb);
        }

}

static void __ublk_batch_init_commit(struct ublk_thread *t,
                                     struct batch_commit_buf *cb,
                                     unsigned short buf_idx)
{
        /* so far only support 1:1 queue/thread mapping */
        cb->buf_idx = buf_idx;
        cb->elem = ublk_get_commit_buf(t, buf_idx);
        cb->done = 0;
        cb->count = t->commit_buf_size /
                t->commit_buf_elem_size;
}

/* COMMIT_IO_CMDS is per-queue command, so use its own commit buffer */
static void ublk_batch_init_commit(struct ublk_thread *t,
                                   struct batch_commit_buf *cb)
{
        unsigned short buf_idx = ublk_alloc_commit_buf(t);

        ublk_assert(buf_idx != UBLKS_T_COMMIT_BUF_INV_IDX);
        ublk_assert(!ublk_batch_commit_prepared(cb));

        __ublk_batch_init_commit(t, cb, buf_idx);
}

void ublk_batch_prep_commit(struct ublk_thread *t)
{
        int i;

        for (i = 0; i < t->nr_queues; i++)
                t->commit[i].buf_idx = UBLKS_T_COMMIT_BUF_INV_IDX;
}

void ublk_batch_complete_io(struct ublk_thread *t, struct ublk_queue *q,
                            unsigned tag, int res)
{
        unsigned q_t_idx = ublk_queue_idx_in_thread(t, q);
        struct batch_commit_buf *cb = &t->commit[q_t_idx];
        struct ublk_batch_elem *elem;
        struct ublk_io *io = &q->ios[tag];

        if (!ublk_batch_commit_prepared(cb))
                ublk_batch_init_commit(t, cb);

        ublk_assert(q->q_id == cb->q_id);

        elem = (struct ublk_batch_elem *)(cb->elem + cb->done * t->commit_buf_elem_size);
        elem->tag = tag;
        elem->buf_index = ublk_batch_io_buf_idx(t, q, tag);
        elem->result = res;

        if (!ublk_queue_no_buf(q))
                elem->buf_addr  = (__u64) (uintptr_t) io->buf_addr;

        cb->done += 1;
        ublk_assert(cb->done <= cb->count);
}

void ublk_batch_setup_map(unsigned char (*q_thread_map)[UBLK_MAX_QUEUES],
                           int nthreads, int queues)
{
        int i, j;

        /*
         * Setup round-robin queue-to-thread mapping for arbitrary N:M combinations.
         *
         * This algorithm distributes queues across threads (and threads across queues)
         * in a balanced round-robin fashion to ensure even load distribution.
         *
         * Examples:
         * - 2 threads, 4 queues: T0=[Q0,Q2], T1=[Q1,Q3]
         * - 4 threads, 2 queues: T0=[Q0], T1=[Q1], T2=[Q0], T3=[Q1]
         * - 3 threads, 3 queues: T0=[Q0], T1=[Q1], T2=[Q2] (1:1 mapping)
         *
         * Phase 1: Mark which queues each thread handles (boolean mapping)
         */
        for (i = 0, j = 0; i < queues || j < nthreads; i++, j++) {
                q_thread_map[j % nthreads][i % queues] = 1;
        }

        /*
         * Phase 2: Convert boolean mapping to sequential indices within each thread.
         *
         * Transform from: q_thread_map[thread][queue] = 1 (handles queue)
         * To:             q_thread_map[thread][queue] = N (queue index within thread)
         *
         * This allows each thread to know the local index of each queue it handles,
         * which is essential for buffer allocation and management. For example:
         * - Thread 0 handling queues [0,2] becomes: q_thread_map[0][0]=1, q_thread_map[0][2]=2
         * - Thread 1 handling queues [1,3] becomes: q_thread_map[1][1]=1, q_thread_map[1][3]=2
         */
        for (j = 0; j < nthreads; j++) {
                unsigned char seq = 1;

                for (i = 0; i < queues; i++) {
                        if (q_thread_map[j][i])
                                q_thread_map[j][i] = seq++;
                }
        }

#if 0
        for (j = 0; j < nthreads; j++) {
                printf("thread %0d: ", j);
                for (i = 0; i < queues; i++) {
                        if (q_thread_map[j][i])
                                printf("%03u ", i);
                }
                printf("\n");
        }
        printf("\n");
        for (j = 0; j < nthreads; j++) {
                for (i = 0; i < queues; i++) {
                        printf("%03u ", q_thread_map[j][i]);
                }
                printf("\n");
        }
#endif
}