/* SPDX-License-Identifier: GPL-2.0 */
#ifndef KUBLK_INTERNAL_H
#define KUBLK_INTERNAL_H

#include <unistd.h>
#include <stdlib.h>
#include <assert.h>
#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <pthread.h>
#include <getopt.h>
#include <limits.h>
#include <poll.h>
#include <fcntl.h>
#include <sys/syscall.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/inotify.h>
#include <sys/wait.h>
#include <sys/eventfd.h>
#include <sys/ipc.h>
#include <sys/shm.h>
#include <linux/io_uring.h>
#include <liburing.h>
#include <semaphore.h>

/* allow ublk_dep.h to override ublk_cmd.h */
#include "ublk_dep.h"
#include <linux/ublk_cmd.h>

#include "utils.h"

#define MAX_BACK_FILES   4

/****************** part 1: libublk ********************/

#define CTRL_DEV		"/dev/ublk-control"
#define UBLKC_DEV		"/dev/ublkc"
#define UBLKB_DEV		"/dev/ublkb"
#define UBLK_CTRL_RING_DEPTH            32
#define ERROR_EVTFD_DEVID 	-2

#define UBLK_IO_MAX_BYTES               (1 << 20)
#define UBLK_MAX_QUEUES_SHIFT		5
#define UBLK_MAX_QUEUES                 (1 << UBLK_MAX_QUEUES_SHIFT)
#define UBLK_MAX_THREADS_SHIFT		5
#define UBLK_MAX_THREADS		(1 << UBLK_MAX_THREADS_SHIFT)
#define UBLK_QUEUE_DEPTH                1024

struct ublk_dev;
struct ublk_queue;
struct ublk_thread;

struct stripe_ctx {
	/* stripe */
	unsigned int    chunk_size;
};

struct fault_inject_ctx {
	/* fault_inject */
	unsigned long   delay_us;
};

struct dev_ctx {
	char tgt_type[16];
	unsigned long flags;
	unsigned nr_hw_queues;
	unsigned short nthreads;
	unsigned queue_depth;
	int dev_id;
	int nr_files;
	char *files[MAX_BACK_FILES];
	unsigned int	logging:1;
	unsigned int	all:1;
	unsigned int	fg:1;
	unsigned int	recovery:1;
	unsigned int	auto_zc_fallback:1;
	unsigned int	per_io_tasks:1;
	unsigned int	no_ublk_fixed_fd:1;
	unsigned int	safe_stop:1;
	unsigned int	no_auto_part_scan:1;
	__u32 integrity_flags;
	__u8 metadata_size;
	__u8 pi_offset;
	__u8 csum_type;
	__u8 tag_size;

	int _evtfd;
	int _shmid;

	/* built from shmem, only for ublk_dump_dev() */
	struct ublk_dev *shadow_dev;

	/* for 'update_size' command */
	unsigned long long size;

	union {
		struct stripe_ctx 	stripe;
		struct fault_inject_ctx fault_inject;
	};
};

struct ublk_ctrl_cmd_data {
	__u32 cmd_op;
#define CTRL_CMD_HAS_DATA	1
#define CTRL_CMD_HAS_BUF	2
	__u32 flags;

	__u64 data[2];
	__u64 addr;
	__u32 len;
};

struct ublk_io {
	char *buf_addr;
	void *integrity_buf;

#define UBLKS_IO_NEED_FETCH_RQ		(1UL << 0)
#define UBLKS_IO_NEED_COMMIT_RQ_COMP	(1UL << 1)
#define UBLKS_IO_FREE			(1UL << 2)
#define UBLKS_IO_NEED_GET_DATA           (1UL << 3)
#define UBLKS_IO_NEED_REG_BUF            (1UL << 4)
	unsigned short flags;
	unsigned short refs;		/* used by target code only */

	int tag;

	int result;

	unsigned short buf_index;
	unsigned short tgt_ios;
	void *private_data;
};

struct ublk_tgt_ops {
	const char *name;
	int (*init_tgt)(const struct dev_ctx *ctx, struct ublk_dev *);
	void (*deinit_tgt)(struct ublk_dev *);

	int (*queue_io)(struct ublk_thread *, struct ublk_queue *, int tag);
	void (*tgt_io_done)(struct ublk_thread *, struct ublk_queue *,
			    const struct io_uring_cqe *);

	/*
	 * Target specific command line handling
	 *
	 * each option requires argument for target command line
	 */
	void (*parse_cmd_line)(struct dev_ctx *ctx, int argc, char *argv[]);
	void (*usage)(const struct ublk_tgt_ops *ops);

	/* return buffer index for UBLK_F_AUTO_BUF_REG */
	unsigned short (*buf_index)(const struct ublk_thread *t,
			const struct ublk_queue *, int tag);
};

struct ublk_tgt {
	unsigned long dev_size;
	unsigned int  sq_depth;
	unsigned int  cq_depth;
	const struct ublk_tgt_ops *ops;
	struct ublk_params params;

	int nr_backing_files;
	unsigned long backing_file_size[MAX_BACK_FILES];
	char backing_file[MAX_BACK_FILES][PATH_MAX];
};

struct ublk_queue {
	int q_id;
	int q_depth;
	struct ublk_dev *dev;
	const struct ublk_tgt_ops *tgt_ops;
	struct ublksrv_io_desc *io_cmd_buf;

/* borrow three bit of ublk uapi flags, which may never be used */
#define UBLKS_Q_AUTO_BUF_REG_FALLBACK	(1ULL << 63)
#define UBLKS_Q_NO_UBLK_FIXED_FD	(1ULL << 62)
#define UBLKS_Q_PREPARED	(1ULL << 61)
	__u64 flags;
	int ublk_fd;	/* cached ublk char device fd */
	__u8 metadata_size;
	struct ublk_io ios[UBLK_QUEUE_DEPTH];

	/* used for prep io commands */
	pthread_spinlock_t lock;
};

/* align with `ublk_elem_header` */
struct ublk_batch_elem {
	__u16 tag;
	__u16 buf_index;
	__s32 result;
	__u64 buf_addr;
};

struct batch_commit_buf {
	unsigned short q_id;
	unsigned short buf_idx;
	void *elem;
	unsigned short done;
	unsigned short count;
};

struct batch_fetch_buf {
	struct io_uring_buf_ring *br;
	void *fetch_buf;
	unsigned int fetch_buf_size;
	unsigned int fetch_buf_off;
};

struct ublk_thread {
	/* Thread-local copy of queue-to-thread mapping for this thread */
	unsigned char q_map[UBLK_MAX_QUEUES];

	struct ublk_dev *dev;
	unsigned short idx;
	unsigned short nr_queues;

#define UBLKS_T_STOPPING	(1U << 0)
#define UBLKS_T_IDLE	(1U << 1)
#define UBLKS_T_BATCH_IO	(1U << 31) 	/* readonly */
	unsigned state;
	unsigned int cmd_inflight;
	unsigned int io_inflight;

	unsigned short nr_bufs;

       /* followings are for BATCH_IO */
	unsigned short commit_buf_start;
	unsigned char  commit_buf_elem_size;
       /*
        * We just support single device, so pre-calculate commit/prep flags
        */
	unsigned short cmd_flags;
	unsigned int   nr_commit_buf;
	unsigned int   commit_buf_size;
	void *commit_buf;
#define UBLKS_T_COMMIT_BUF_INV_IDX  ((unsigned short)-1)
	struct allocator commit_buf_alloc;
	struct batch_commit_buf *commit;
	/* FETCH_IO_CMDS buffer */
	unsigned short nr_fetch_bufs;
	struct batch_fetch_buf *fetch;

	struct io_uring ring;
};

struct ublk_dev {
	struct ublk_tgt tgt;
	struct ublksrv_ctrl_dev_info  dev_info;
	struct ublk_queue q[UBLK_MAX_QUEUES];
	unsigned nthreads;
	unsigned per_io_tasks;

	int fds[MAX_BACK_FILES + 1];	/* fds[0] points to /dev/ublkcN */
	int nr_fds;
	int ctrl_fd;
	struct io_uring ring;

	void *private_data;
};

extern int ublk_queue_io_cmd(struct ublk_thread *t, struct ublk_io *io);

static inline int __ublk_use_batch_io(__u64 flags)
{
	return flags & UBLK_F_BATCH_IO;
}

static inline int ublk_queue_batch_io(const struct ublk_queue *q)
{
	return __ublk_use_batch_io(q->flags);
}

static inline int ublk_dev_batch_io(const struct ublk_dev *dev)
{
	return __ublk_use_batch_io(dev->dev_info.flags);
}

/* only work for handle single device in this pthread context */
static inline int ublk_thread_batch_io(const struct ublk_thread *t)
{
	return t->state & UBLKS_T_BATCH_IO;
}

static inline void ublk_set_integrity_params(const struct dev_ctx *ctx,
					     struct ublk_params *params)
{
	if (!ctx->metadata_size)
		return;

	params->types |= UBLK_PARAM_TYPE_INTEGRITY;
	params->integrity = (struct ublk_param_integrity) {
		.flags = ctx->integrity_flags,
		.interval_exp = params->basic.logical_bs_shift,
		.metadata_size = ctx->metadata_size,
		.pi_offset = ctx->pi_offset,
		.csum_type = ctx->csum_type,
		.tag_size = ctx->tag_size,
	};
}

static inline size_t ublk_integrity_len(const struct ublk_queue *q, size_t len)
{
	/* All targets currently use interval_exp = logical_bs_shift = 9 */
	return (len >> 9) * q->metadata_size;
}

static inline size_t
ublk_integrity_data_len(const struct ublk_queue *q, size_t integrity_len)
{
	return (integrity_len / q->metadata_size) << 9;
}

static inline int ublk_io_auto_zc_fallback(const struct ublksrv_io_desc *iod)
{
	return !!(iod->op_flags & UBLK_IO_F_NEED_REG_BUF);
}

static inline __u64 ublk_user_copy_offset(unsigned q_id, unsigned tag)
{
	return UBLKSRV_IO_BUF_OFFSET +
	       ((__u64)q_id << UBLK_QID_OFF | (__u64)tag << UBLK_TAG_OFF);
}

static inline int is_target_io(__u64 user_data)
{
	return (user_data & (1ULL << 63)) != 0;
}

static inline __u64 build_user_data(unsigned tag, unsigned op,
		unsigned tgt_data, unsigned q_id, unsigned is_target_io)
{
	/* we only have 7 bits to encode q_id */
	_Static_assert(UBLK_MAX_QUEUES_SHIFT <= 7, "UBLK_MAX_QUEUES_SHIFT must be <= 7");
	ublk_assert(!(tag >> 16) && !(op >> 8) && !(tgt_data >> 16) && !(q_id >> 7));

	return tag | ((__u64)op << 16) | ((__u64)tgt_data << 24) |
		(__u64)q_id << 56 | (__u64)is_target_io << 63;
}

static inline unsigned int user_data_to_tag(__u64 user_data)
{
	return user_data & 0xffff;
}

static inline unsigned int user_data_to_op(__u64 user_data)
{
	return (user_data >> 16) & 0xff;
}

static inline unsigned int user_data_to_tgt_data(__u64 user_data)
{
	return (user_data >> 24) & 0xffff;
}

static inline unsigned int user_data_to_q_id(__u64 user_data)
{
	return (user_data >> 56) & 0x7f;
}

static inline unsigned short ublk_cmd_op_nr(unsigned int op)
{
	return _IOC_NR(op);
}

static inline struct ublk_queue *ublk_io_to_queue(const struct ublk_io *io)
{
	return container_of(io, struct ublk_queue, ios[io->tag]);
}

static inline int ublk_io_alloc_sqes(struct ublk_thread *t,
		struct io_uring_sqe *sqes[], int nr_sqes)
{
	struct io_uring *ring = &t->ring;
	unsigned left = io_uring_sq_space_left(ring);
	int i;

	if (left < nr_sqes)
		io_uring_submit(ring);

	for (i = 0; i < nr_sqes; i++) {
		sqes[i] = io_uring_get_sqe(ring);
		if (!sqes[i])
			return i;
	}

	return nr_sqes;
}

static inline int ublk_get_registered_fd(struct ublk_queue *q, int fd_index)
{
	if (q->flags & UBLKS_Q_NO_UBLK_FIXED_FD) {
		if (fd_index == 0)
			/* Return the raw ublk FD for index 0 */
			return q->ublk_fd;
		/* Adjust index for backing files (index 1 becomes 0, etc.) */
		return fd_index - 1;
	}
	return fd_index;
}

static inline void __io_uring_prep_buf_reg_unreg(struct io_uring_sqe *sqe,
		struct ublk_queue *q, int tag, int q_id, __u64 index)
{
	struct ublksrv_io_cmd *cmd = (struct ublksrv_io_cmd *)sqe->cmd;
	int dev_fd = ublk_get_registered_fd(q, 0);

	io_uring_prep_read(sqe, dev_fd, 0, 0, 0);
	sqe->opcode		= IORING_OP_URING_CMD;
	if (q->flags & UBLKS_Q_NO_UBLK_FIXED_FD)
		sqe->flags	&= ~IOSQE_FIXED_FILE;
	else
		sqe->flags	|= IOSQE_FIXED_FILE;

	cmd->tag		= tag;
	cmd->addr		= index;
	cmd->q_id		= q_id;
}

static inline void io_uring_prep_buf_register(struct io_uring_sqe *sqe,
		struct ublk_queue *q, int tag, int q_id, __u64 index)
{
	__io_uring_prep_buf_reg_unreg(sqe, q, tag, q_id, index);
	sqe->cmd_op		= UBLK_U_IO_REGISTER_IO_BUF;
}

static inline void io_uring_prep_buf_unregister(struct io_uring_sqe *sqe,
		struct ublk_queue *q, int tag, int q_id, __u64 index)
{
	__io_uring_prep_buf_reg_unreg(sqe, q, tag, q_id, index);
	sqe->cmd_op		= UBLK_U_IO_UNREGISTER_IO_BUF;
}

static inline void *ublk_get_sqe_cmd(const struct io_uring_sqe *sqe)
{
	return (void *)&sqe->cmd;
}

static inline void ublk_set_io_res(struct ublk_queue *q, int tag, int res)
{
	q->ios[tag].result = res;
}

static inline int ublk_get_io_res(const struct ublk_queue *q, unsigned tag)
{
	return q->ios[tag].result;
}

static inline void ublk_mark_io_done(struct ublk_io *io, int res)
{
	io->flags |= (UBLKS_IO_NEED_COMMIT_RQ_COMP | UBLKS_IO_FREE);
	io->result = res;
}

static inline const struct ublksrv_io_desc *ublk_get_iod(const struct ublk_queue *q, int tag)
{
	return &q->io_cmd_buf[tag];
}

static inline void ublk_set_sqe_cmd_op(struct io_uring_sqe *sqe, __u32 cmd_op)
{
	__u32 *addr = (__u32 *)&sqe->off;

	addr[0] = cmd_op;
	addr[1] = 0;
}

static inline unsigned short ublk_batch_io_buf_idx(
		const struct ublk_thread *t, const struct ublk_queue *q,
		unsigned tag);

static inline unsigned short ublk_io_buf_idx(const struct ublk_thread *t,
					     const struct ublk_queue *q,
					     unsigned tag)
{
	if (ublk_queue_batch_io(q))
		return ublk_batch_io_buf_idx(t, q, tag);
	return q->ios[tag].buf_index;
}

static inline struct ublk_io *ublk_get_io(struct ublk_queue *q, unsigned tag)
{
	return &q->ios[tag];
}

static inline int ublk_completed_tgt_io(struct ublk_thread *t,
					struct ublk_queue *q, unsigned tag)
{
	struct ublk_io *io = ublk_get_io(q, tag);

	t->io_inflight--;

	return --io->tgt_ios == 0;
}

static inline bool ublk_queue_use_zc(const struct ublk_queue *q)
{
	return !!(q->flags & UBLK_F_SUPPORT_ZERO_COPY);
}

static inline bool ublk_queue_use_auto_zc(const struct ublk_queue *q)
{
	return !!(q->flags & UBLK_F_AUTO_BUF_REG);
}

static inline bool ublk_queue_auto_zc_fallback(const struct ublk_queue *q)
{
	return !!(q->flags & UBLKS_Q_AUTO_BUF_REG_FALLBACK);
}

static inline bool ublk_queue_use_user_copy(const struct ublk_queue *q)
{
	return !!(q->flags & UBLK_F_USER_COPY);
}

static inline int ublk_queue_no_buf(const struct ublk_queue *q)
{
	return ublk_queue_use_zc(q) || ublk_queue_use_auto_zc(q);
}

static inline int ublk_batch_commit_prepared(struct batch_commit_buf *cb)
{
	return cb->buf_idx != UBLKS_T_COMMIT_BUF_INV_IDX;
}

static inline unsigned ublk_queue_idx_in_thread(const struct ublk_thread *t,
						const struct ublk_queue *q)
{
	unsigned char idx;

	idx = t->q_map[q->q_id];
	ublk_assert(idx != 0);
	return idx - 1;
}

/*
 * Each IO's buffer index has to be calculated by this helper for
 * UBLKS_T_BATCH_IO
 */
static inline unsigned short ublk_batch_io_buf_idx(
		const struct ublk_thread *t, const struct ublk_queue *q,
		unsigned tag)
{
	return ublk_queue_idx_in_thread(t, q) * q->q_depth + tag;
}

/* Queue UBLK_U_IO_PREP_IO_CMDS for a specific queue with batch elements */
int ublk_batch_queue_prep_io_cmds(struct ublk_thread *t, struct ublk_queue *q);
/* Start fetching I/O commands using multishot UBLK_U_IO_FETCH_IO_CMDS */
void ublk_batch_start_fetch(struct ublk_thread *t);
/* Handle completion of batch I/O commands (prep/commit) */
void ublk_batch_compl_cmd(struct ublk_thread *t,
			  const struct io_uring_cqe *cqe);
/* Initialize batch I/O state and calculate buffer parameters */
void ublk_batch_prepare(struct ublk_thread *t);
/* Allocate and register commit buffers for batch operations */
int ublk_batch_alloc_buf(struct ublk_thread *t);
/* Free commit buffers and cleanup batch allocator */
void ublk_batch_free_buf(struct ublk_thread *t);

/* Prepare a new commit buffer for batching completed I/O operations */
void ublk_batch_prep_commit(struct ublk_thread *t);
/* Submit UBLK_U_IO_COMMIT_IO_CMDS with batched completed I/O operations */
void ublk_batch_commit_io_cmds(struct ublk_thread *t);
/* Add a completed I/O operation to the current batch commit buffer */
void ublk_batch_complete_io(struct ublk_thread *t, struct ublk_queue *q,
			    unsigned tag, int res);
void ublk_batch_setup_map(unsigned char (*q_thread_map)[UBLK_MAX_QUEUES],
			   int nthreads, int queues);

static inline int ublk_complete_io(struct ublk_thread *t, struct ublk_queue *q,
				   unsigned tag, int res)
{
	if (ublk_queue_batch_io(q)) {
		ublk_batch_complete_io(t, q, tag, res);
		return 0;
	} else {
		struct ublk_io *io = &q->ios[tag];

		ublk_mark_io_done(io, res);
		return ublk_queue_io_cmd(t, io);
	}
}

static inline void ublk_queued_tgt_io(struct ublk_thread *t, struct ublk_queue *q,
				      unsigned tag, int queued)
{
	if (queued < 0)
		ublk_complete_io(t, q, tag, queued);
	else {
		struct ublk_io *io = ublk_get_io(q, tag);

		t->io_inflight += queued;
		io->tgt_ios = queued;
		io->result = 0;
	}
}

extern const struct ublk_tgt_ops null_tgt_ops;
extern const struct ublk_tgt_ops loop_tgt_ops;
extern const struct ublk_tgt_ops stripe_tgt_ops;
extern const struct ublk_tgt_ops fault_inject_tgt_ops;

void backing_file_tgt_deinit(struct ublk_dev *dev);
int backing_file_tgt_init(struct ublk_dev *dev, unsigned int nr_direct);

#endif