//
// Copyright 2024 Staysail Systems, Inc. <info@staysail.tech>
// Copyright 2018 Capitar IT Group BV <info@capitar.com>
//
// This software is supplied under the terms of the MIT License, a
// copy of which should be located in the distribution where this
// file was obtained (LICENSE.txt).  A copy of the license may also be
// found online at https://opensource.org/licenses/MIT.
//

#include "core/nng_impl.h"

#include <limits.h>
#include <string.h>

struct nni_aio_expire_q {
	nni_mtx  eq_mtx;
	nni_cv   eq_cv;
	nni_list eq_list;
	nni_thr  eq_thr;
	nni_time eq_next;
	bool     eq_exit;
};

static nni_aio_expire_q **nni_aio_expire_q_list;
static int                nni_aio_expire_q_cnt;

static nni_reap_list aio_reap_list = {
	.rl_offset = offsetof(nni_aio, a_reap_node),
	.rl_func   = nni_aio_free_cb,
};

static void nni_aio_expire_add(nni_aio *);
static void nni_aio_expire_rm(nni_aio *);

void
nni_aio_init(nni_aio *aio, nni_cb cb, void *arg)
{
	memset(aio, 0, sizeof(*aio));
	nni_task_init(&aio->a_task, NULL, cb, arg);
	aio->a_expire  = NNI_TIME_NEVER;
	aio->a_timeout = NNG_DURATION_INFINITE;
	aio->a_expire_q =
	    nni_aio_expire_q_list[nni_random() % nni_aio_expire_q_cnt];
}

void
nni_aio_fini(nni_aio *aio)
{
	nni_aio_cancel_fn fn;
	void             *arg;
	nni_aio_expire_q *eq = aio->a_expire_q;

	nni_mtx_lock(&eq->eq_mtx);
	aio->a_stop = true;
	while (aio->a_expiring) {
		nni_cv_wait(&eq->eq_cv);
	}
	nni_aio_expire_rm(aio);
	fn                = aio->a_cancel_fn;
	arg               = aio->a_cancel_arg;
	aio->a_cancel_fn  = NULL;
	aio->a_cancel_arg = NULL;
	nni_mtx_unlock(&eq->eq_mtx);

	if (fn != NULL) {
		fn(aio, arg, NNG_ECLOSED);
	} else {
		nni_task_abort(&aio->a_task);
	}

	nni_task_fini(&aio->a_task);
}

int
nni_aio_alloc(nni_aio **aio_p, nni_cb cb, void *arg)
{
	nni_aio *aio;

	if ((aio = NNI_ALLOC_STRUCT(aio)) == NULL) {
		return (NNG_ENOMEM);
	}
	nni_aio_init(aio, cb, arg);
	*aio_p = aio;
	return (0);
}

void
nni_aio_free(nni_aio *aio)
{
	if (aio != NULL) {
		nni_aio_fini(aio);
		NNI_FREE_STRUCT(aio);
	}
}

void
nni_aio_free_cb(void *aio)
{
  nni_aio_free((nni_aio *) aio);
}

void
nni_aio_reap(nni_aio *aio)
{
	if (aio != NULL) {
		nni_reap(&aio_reap_list, aio);
	}
}

int
nni_aio_set_iov(nni_aio *aio, unsigned nio, const nni_iov *iov)
{

	if (nio > NNI_NUM_ELEMENTS((aio->a_iov))) {
		return (NNG_EINVAL);
	}

	if (iov != &aio->a_iov[0]) {
		for (unsigned i = 0; i < nio; i++) {
			aio->a_iov[i] = iov[i];
		}
	}
	aio->a_nio = nio;
	return (0);
}

void
nni_aio_stop(nni_aio *aio)
{
	if (aio != NULL) {
		nni_aio_cancel_fn fn;
		void             *arg;
		nni_aio_expire_q *eq = aio->a_expire_q;

		nni_mtx_lock(&eq->eq_mtx);
		nni_aio_expire_rm(aio);
		fn                = aio->a_cancel_fn;
		arg               = aio->a_cancel_arg;
		aio->a_cancel_fn  = NULL;
		aio->a_cancel_arg = NULL;
		aio->a_stop       = true;
		nni_mtx_unlock(&eq->eq_mtx);

		if (fn != NULL) {
			fn(aio, arg, NNG_ECANCELED);
		} else {
			nni_task_abort(&aio->a_task);
		}

		nni_aio_wait(aio);
	}
}

void
nni_aio_close(nni_aio *aio)
{
	if (aio != NULL) {
		nni_aio_cancel_fn fn;
		void             *arg;
		nni_aio_expire_q *eq = aio->a_expire_q;

		nni_mtx_lock(&eq->eq_mtx);
		nni_aio_expire_rm(aio);
		fn                = aio->a_cancel_fn;
		arg               = aio->a_cancel_arg;
		aio->a_cancel_fn  = NULL;
		aio->a_cancel_arg = NULL;
		aio->a_stop       = true;
		nni_mtx_unlock(&eq->eq_mtx);

		if (fn != NULL) {
			fn(aio, arg, NNG_ECLOSED);
		} else {
			nni_task_abort(&aio->a_task);
		}
	}
}

void
nni_aio_set_timeout(nni_aio *aio, nni_duration when)
{
	aio->a_timeout    = when;
	aio->a_use_expire = false;
}

void
nni_aio_set_expire(nni_aio *aio, nni_time expire)
{
	aio->a_expire     = expire;
	aio->a_use_expire = true;
}

nng_duration
nni_aio_get_timeout(nni_aio *aio)
{
	return (aio->a_timeout);
}

void
nni_aio_set_msg(nni_aio *aio, nni_msg *msg)
{
	aio->a_msg = msg;
}

nni_msg *
nni_aio_get_msg(nni_aio *aio)
{
	return (aio->a_msg);
}

void
nni_aio_set_input(nni_aio *aio, unsigned index, void *data)
{
	if (index < NNI_NUM_ELEMENTS(aio->a_inputs)) {
		aio->a_inputs[index] = data;
	}
}

void *
nni_aio_get_input(nni_aio *aio, unsigned index)
{
	if (index < NNI_NUM_ELEMENTS(aio->a_inputs)) {
		return (aio->a_inputs[index]);
	}
	return (NULL);
}

void
nni_aio_set_output(nni_aio *aio, unsigned index, void *data)
{
	if (index < NNI_NUM_ELEMENTS(aio->a_outputs)) {
		aio->a_outputs[index] = data;
	}
}

void *
nni_aio_get_output(nni_aio *aio, unsigned index)
{
	if (index < NNI_NUM_ELEMENTS(aio->a_outputs)) {
		return (aio->a_outputs[index]);
	}
	return (NULL);
}

int
nni_aio_result(nni_aio *aio)
{
	return (aio->a_result);
}

size_t
nni_aio_count(nni_aio *aio)
{
	return (aio->a_count);
}

void
nni_aio_wait(nni_aio *aio)
{
	nni_task_wait(&aio->a_task);
}

bool
nni_aio_busy(nni_aio *aio)
{
	return (nni_task_busy(&aio->a_task));
}

int
nni_aio_begin(nni_aio *aio)
{
	nni_aio_expire_q *eq = aio->a_expire_q;

  nni_mtx_lock(&eq->eq_mtx);
	NNI_ASSERT(!nni_aio_list_active(aio));
	NNI_ASSERT(aio->a_cancel_fn == NULL);
	NNI_ASSERT(!nni_list_node_active(&aio->a_expire_node));

	for (unsigned i = 0; i < NNI_NUM_ELEMENTS(aio->a_outputs); i++) {
		aio->a_outputs[i] = NULL;
	}
	aio->a_result    = 0;
	aio->a_count     = 0;
	aio->a_cancel_fn = NULL;

	if (aio->a_stop) {
		aio->a_result    = NNG_ECANCELED;
		aio->a_cancel_fn = NULL;
		aio->a_expire    = NNI_TIME_NEVER;
		aio->a_sleep     = false;
		aio->a_expire_ok = false;
		nni_mtx_unlock(&eq->eq_mtx);

		return (NNG_ECANCELED);
	}
	nni_task_prep(&aio->a_task);
	nni_mtx_unlock(&eq->eq_mtx);
	return (0);
}

int
nni_aio_schedule(nni_aio *aio, nni_aio_cancel_fn cancel, void *data)
{
	nni_aio_expire_q *eq = aio->a_expire_q;

	if ((!aio->a_sleep) && (!aio->a_use_expire)) {
		switch (aio->a_timeout) {
		case NNG_DURATION_ZERO:
			nni_task_abort(&aio->a_task);
			return (NNG_ETIMEDOUT);
		case NNG_DURATION_INFINITE:
		case NNG_DURATION_DEFAULT:
			aio->a_expire = NNI_TIME_NEVER;
			break;
		default:
			aio->a_expire = nni_clock() + aio->a_timeout;
			break;
		}
	}

	nni_mtx_lock(&eq->eq_mtx);
	if (aio->a_stop) {
		nni_task_abort(&aio->a_task);
		nni_mtx_unlock(&eq->eq_mtx);
		return (NNG_ECLOSED);
	}

	NNI_ASSERT(aio->a_cancel_fn == NULL);
	aio->a_cancel_fn  = cancel;
	aio->a_cancel_arg = data;

	if ((aio->a_expire != NNI_TIME_NEVER) && (cancel != NULL)) {
		nni_aio_expire_add(aio);
	}
	nni_mtx_unlock(&eq->eq_mtx);
	return (0);
}

void
nni_aio_abort(nni_aio *aio, int rv)
{
	nni_aio_cancel_fn fn;
	void             *arg;
	nni_aio_expire_q *eq = aio->a_expire_q;

	nni_mtx_lock(&eq->eq_mtx);
	nni_aio_expire_rm(aio);
	fn                = aio->a_cancel_fn;
	arg               = aio->a_cancel_arg;
	aio->a_cancel_fn  = NULL;
	aio->a_cancel_arg = NULL;
	nni_mtx_unlock(&eq->eq_mtx);

	if (fn != NULL) {
		fn(aio, arg, rv);
	} else {
		nni_task_abort(&aio->a_task);
	}
}

static void
nni_aio_finish_impl(
    nni_aio *aio, int rv, size_t count, nni_msg *msg, bool sync)
{
	nni_aio_expire_q *eq = aio->a_expire_q;

	nni_mtx_lock(&eq->eq_mtx);

	nni_aio_expire_rm(aio);
	aio->a_result     = rv;
	aio->a_count      = count;
	aio->a_cancel_fn  = NULL;
	aio->a_cancel_arg = NULL;
	if (msg) {
		aio->a_msg = msg;
	}

	aio->a_expire     = NNI_TIME_NEVER;
	aio->a_sleep      = false;
	aio->a_use_expire = false;
	nni_mtx_unlock(&eq->eq_mtx);

	if (sync) {
		nni_task_exec(&aio->a_task);
	} else {
		nni_task_dispatch(&aio->a_task);
	}
}

void
nni_aio_finish(nni_aio *aio, int result, size_t count)
{
	nni_aio_finish_impl(aio, result, count, NULL, false);
}

void
nni_aio_finish_sync(nni_aio *aio, int result, size_t count)
{
	nni_aio_finish_impl(aio, result, count, NULL, true);
}

void
nni_aio_finish_error(nni_aio *aio, int result)
{
	nni_aio_finish_impl(aio, result, 0, NULL, false);
}

void
nni_aio_finish_msg(nni_aio *aio, nni_msg *msg)
{
	NNI_ASSERT(msg != NULL);
	nni_aio_finish_impl(aio, 0, nni_msg_len(msg), msg, false);
}

void
nni_aio_list_init(nni_list *list)
{
	NNI_LIST_INIT(list, nni_aio, a_prov_node);
}

void
nni_aio_list_append(nni_list *list, nni_aio *aio)
{
	nni_aio_list_remove(aio);
	nni_list_append(list, aio);
}

void
nni_aio_list_remove(nni_aio *aio)
{
	nni_list_node_remove(&aio->a_prov_node);
}

int
nni_aio_list_active(nni_aio *aio)
{
	return (nni_list_node_active(&aio->a_prov_node));
}

void
nni_aio_completions_init(nni_aio_completions *clp)
{
	*clp = NULL;
}

void
nni_aio_completions_add(
    nni_aio_completions *clp, nni_aio *aio, int result, size_t count)
{
	NNI_ASSERT(!nni_aio_list_active(aio));
	aio->a_reap_node.rn_next = *clp;
	aio->a_result            = result;
	aio->a_count             = count;
	*clp                     = aio;
}

void
nni_aio_completions_run(nni_aio_completions *clp)
{
	nni_aio *aio;
	nni_aio *cl = *clp;
	*clp        = NULL;

	while ((aio = cl) != NULL) {
		cl                       = (void *) aio->a_reap_node.rn_next;
		aio->a_reap_node.rn_next = NULL;
		nni_aio_finish_sync(aio, aio->a_result, aio->a_count);
	}
}

static void
nni_aio_expire_add(nni_aio *aio)
{
	nni_aio_expire_q *eq = aio->a_expire_q;

	nni_list_append(&eq->eq_list, aio);

	if (eq->eq_next > aio->a_expire) {
		eq->eq_next = aio->a_expire;
		nni_cv_wake(&eq->eq_cv);
	}
}

static void
nni_aio_expire_rm(nni_aio *aio)
{
	nni_list_node_remove(&aio->a_expire_node);

}

static void
nni_aio_expire_loop(void *arg)
{
	nni_aio_expire_q *q   = arg;
	nni_mtx          *mtx = &q->eq_mtx;
	nni_cv           *cv  = &q->eq_cv;
	nni_time          now;
	uint32_t          exp_idx;
	nni_aio          *expires[NNI_EXPIRE_BATCH];

	nni_thr_set_name(NULL, "nng:aio:expire");

	nni_mtx_lock(mtx);

	for (;;) {
		nni_aio *aio;
		int      rv;
		nni_time next;

		next = q->eq_next;
		now  = nni_clock();

		aio = nni_list_first(&q->eq_list);
		if ((aio == NULL) && (q->eq_exit)) {
			nni_mtx_unlock(mtx);
			return;
		}
		if (now < next) {
			nni_cv_until(cv, next);
			continue;
		}
		q->eq_next = NNI_TIME_NEVER;
		exp_idx    = 0;
		while (aio != NULL) {
			if ((aio->a_expire < now) &&
			    (exp_idx < NNI_EXPIRE_BATCH)) {
				nni_aio *nxt;

				expires[exp_idx++] = aio;
				nxt = nni_list_next(&q->eq_list, aio);
				nni_list_remove(&q->eq_list, aio);
				aio->a_expiring = true;
				aio             = nxt;
				continue;
			}
			if (aio->a_expire < q->eq_next) {
				q->eq_next = aio->a_expire;
			}
			aio = nni_list_next(&q->eq_list, aio);
		}

		for (uint32_t i = 0; i < exp_idx; i++) {
			aio = expires[i];
			rv  = aio->a_expire_ok ? 0 : NNG_ETIMEDOUT;

			nni_aio_cancel_fn cancel_fn  = aio->a_cancel_fn;
			void             *cancel_arg = aio->a_cancel_arg;

			aio->a_cancel_fn  = NULL;
			aio->a_cancel_arg = NULL;

			if (cancel_fn != NULL) {
				nni_mtx_unlock(mtx);
				cancel_fn(aio, cancel_arg, rv);
				nni_mtx_lock(mtx);
			}
			aio->a_expiring = false;
		}
		nni_cv_wake(cv);

		if (now < q->eq_next) {
			nni_cv_until(cv, q->eq_next);
		}
	}
}

void *
nni_aio_get_prov_data(nni_aio *aio)
{
	return (aio->a_prov_data);
}

void
nni_aio_set_prov_data(nni_aio *aio, void *data)
{
	aio->a_prov_data = data;
}

void
nni_aio_get_iov(nni_aio *aio, unsigned *nio_p, nni_iov **iov_p)
{
	*nio_p = aio->a_nio;
	*iov_p = aio->a_iov;
}

void
nni_aio_normalize_timeout(nni_aio *aio, nng_duration dur)
{
	if (aio->a_timeout == NNG_DURATION_DEFAULT) {
		aio->a_timeout = dur;
	}
}

void
nni_aio_bump_count(nni_aio *aio, size_t n)
{
	aio->a_count += n;
}

size_t
nni_aio_iov_count(nni_aio *aio)
{
	size_t residual = 0;

	for (unsigned i = 0; i < aio->a_nio; i++) {
		residual += aio->a_iov[i].iov_len;
	}
	return (residual);
}

size_t
nni_aio_iov_clamp_len(size_t len, size_t count)
{
	size_t headroom = (size_t) INT_MAX - count;
	return (len < headroom ? len : headroom);
}

size_t
nni_aio_iov_advance(nni_aio *aio, size_t n)
{
	size_t residual = n;
	while (n) {
		NNI_ASSERT(aio->a_nio != 0);
		if (aio->a_iov[0].iov_len > n) {
			aio->a_iov[0].iov_len -= n;
			NNI_INCPTR(aio->a_iov[0].iov_buf, n);
			return (0);
		}
		residual -= aio->a_iov[0].iov_len;
		n -= aio->a_iov[0].iov_len;
		aio->a_nio--;
		for (unsigned i = 0; i < aio->a_nio; i++) {
			aio->a_iov[i] = aio->a_iov[i + 1];
		}
	}
	return (residual);
}

static void
nni_sleep_cancel(nng_aio *aio, void *arg, int rv)
{
	NNI_ARG_UNUSED(arg);
	nni_aio_expire_q *eq = aio->a_expire_q;

	nni_mtx_lock(&eq->eq_mtx);
	if (!aio->a_sleep) {
		nni_mtx_unlock(&eq->eq_mtx);
		return;
	}

	aio->a_sleep = false;
	nni_aio_expire_rm(aio);
	nni_mtx_unlock(&eq->eq_mtx);

	nni_aio_finish_error(aio, rv);
}

void
nni_sleep_aio(nng_duration ms, nng_aio *aio)
{
	int rv;
	if (nni_aio_begin(aio) != 0) {
		return;
	}
	aio->a_expire_ok = true;
	aio->a_sleep     = true;
	switch (aio->a_timeout) {
	case NNG_DURATION_DEFAULT:
	case NNG_DURATION_INFINITE:
		break;
	default:
		if (ms > aio->a_timeout) {
			aio->a_expire_ok = false;
			ms               = aio->a_timeout;
		}
	}
	aio->a_expire = nni_clock() + ms;

	if ((rv = nni_aio_schedule(aio, nni_sleep_cancel, NULL)) != 0) {
		nni_aio_finish_error(aio, rv);
	}
}

static void
nni_aio_expire_q_free(nni_aio_expire_q *eq)
{
	if (eq == NULL) {
		return;
	}
	if (!eq->eq_exit) {
		nni_mtx_lock(&eq->eq_mtx);
		eq->eq_exit = true;
		nni_cv_wake(&eq->eq_cv);
		nni_mtx_unlock(&eq->eq_mtx);
	}

	nni_thr_fini(&eq->eq_thr);
	nni_cv_fini(&eq->eq_cv);
	nni_mtx_fini(&eq->eq_mtx);
	NNI_FREE_STRUCT(eq);
}

static nni_aio_expire_q *
nni_aio_expire_q_alloc(void)
{
	nni_aio_expire_q *eq;

	if ((eq = NNI_ALLOC_STRUCT(eq)) == NULL) {
		return (NULL);
	}
	nni_mtx_init(&eq->eq_mtx);
	nni_cv_init(&eq->eq_cv, &eq->eq_mtx);
	NNI_LIST_INIT(&eq->eq_list, nni_aio, a_expire_node);
	eq->eq_next = NNI_TIME_NEVER;
	eq->eq_exit = false;

	if (nni_thr_init(&eq->eq_thr, nni_aio_expire_loop, eq) != 0) {
		nni_aio_expire_q_free(eq);
		return (NULL);
	}

	nni_thr_run(&eq->eq_thr);
	return (eq);
}

void
nni_aio_sys_fini(void)
{
	for (int i = 0; i < nni_aio_expire_q_cnt; i++) {
		nni_aio_expire_q_free(nni_aio_expire_q_list[i]);
	}
	nni_free(nni_aio_expire_q_list,
	    sizeof(nni_aio_expire_q *) * nni_aio_expire_q_cnt);
	nni_aio_expire_q_cnt  = 0;
	nni_aio_expire_q_list = NULL;
}

int
nni_aio_sys_init(void)
{
	int num_thr;
	int max_thr;

#ifndef NNG_MAX_EXPIRE_THREADS
#define NNG_MAX_EXPIRE_THREADS 8
#endif

#ifndef NNG_NUM_EXPIRE_THREADS
#define NNG_NUM_EXPIRE_THREADS (nni_plat_ncpu())
#endif

	max_thr = (int) nni_init_get_param(
	    NNG_INIT_MAX_EXPIRE_THREADS, NNG_MAX_EXPIRE_THREADS);

	num_thr = (int) nni_init_get_param(
	    NNG_INIT_NUM_EXPIRE_THREADS, NNG_NUM_EXPIRE_THREADS);

	if ((max_thr > 0) && (num_thr > max_thr)) {
		num_thr = max_thr;
	}
	if (num_thr < 1) {
		num_thr = 1;
	}
	nni_init_set_effective(NNG_INIT_NUM_EXPIRE_THREADS, num_thr);
	nni_aio_expire_q_list =
	    nni_zalloc(sizeof(nni_aio_expire_q *) * num_thr);
	nni_aio_expire_q_cnt = num_thr;
	for (int i = 0; i < num_thr; i++) {
		nni_aio_expire_q *eq;
		if ((eq = nni_aio_expire_q_alloc()) == NULL) {
			nni_aio_sys_fini();
			return (NNG_ENOMEM);
		}
		nni_aio_expire_q_list[i] = eq;
	}

	return (0);
}
