/* libcr_ipc/chan.c - Simple channels for libcr
 *
 * Copyright (C) 2024-2025  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <string.h>  /* for memcpy() */

#include <libcr/coroutine.h> /* for cid_t, cr_* */
#include <libmisc/assert.h>
#include <libmisc/rand.h>

#define IMPLEMENTATION_FOR_LIBCR_IPC_CHAN_H YES
#include <libcr_ipc/chan.h>

struct _cr_select_waiter {
	cid_t                            cid;
	struct _cr_select_arg_list_node *arg_vec;
	size_t                           arg_cnt;

	size_t                           ret;
};

enum cr_select_class {
	CR_SELECT_CLASS_DEFAULT,
	CR_SELECT_CLASS_BLOCKING,
	CR_SELECT_CLASS_NONBLOCK,
};

static inline enum cr_select_class cr_select_getclass(struct _cr_select_arg *arg) {
	switch (arg->op) {
	case _CR_SELECT_OP_RECV:
		if (arg->ch->waiters.front && arg->ch->waiters.front->val.op == _CR_SELECT_OP_SEND)
			return CR_SELECT_CLASS_NONBLOCK;
		else
			return CR_SELECT_CLASS_BLOCKING;
	case _CR_SELECT_OP_SEND:
		if (arg->ch->waiters.front && arg->ch->waiters.front->val.op == _CR_SELECT_OP_RECV)
			return CR_SELECT_CLASS_NONBLOCK;
		else
			return CR_SELECT_CLASS_BLOCKING;
	case _CR_SELECT_OP_DEFAULT:
		return CR_SELECT_CLASS_DEFAULT;
	default:
		assert_notreached("invalid arg->op");
	}
}

size_t cr_select_v(size_t arg_cnt, struct _cr_select_arg_list_node arg_vec[]) {
	size_t  cnt_blocking = 0;
	size_t  cnt_nonblock = 0;
	size_t  cnt_default  = 0;

	assert(arg_cnt);
	assert(arg_vec);
	cr_assert_in_coroutine();

	for (size_t i = 0; i < arg_cnt; i++) {
		switch (cr_select_getclass(&arg_vec[i].val)) {
		case CR_SELECT_CLASS_BLOCKING:
			cnt_blocking++;
			break;
		case CR_SELECT_CLASS_NONBLOCK:
			cnt_nonblock++;
			break;
		case CR_SELECT_CLASS_DEFAULT:
			cnt_default++;
			break;
		}
	}

	if (cnt_nonblock) {
		size_t choice_among_nonblock = rand_uint63n(cnt_nonblock);
		size_t choice_among_all = arg_cnt;
		for (size_t i = 0, seen = 0; i < choice_among_all; i++) {
			if (cr_select_getclass(&arg_vec[i].val) == CR_SELECT_CLASS_NONBLOCK) {
				if (seen == choice_among_nonblock)
					choice_among_all = i;
				seen++;
			}
		}
		assert(choice_among_all < arg_cnt);

		struct _cr_select_arg *this = &arg_vec[choice_among_all].val;
		assert(this->ch->waiters.front);
		struct _cr_select_arg *other = &this->ch->waiters.front->val;
		assert(this->val_siz == other->val_siz);
		assert(this->ch == other->ch);
		switch (this->op) {
		case _CR_SELECT_OP_SEND:
			assert(other->op == _CR_SELECT_OP_RECV);
			memcpy(other->val_ptr, this->val_ptr, this->val_siz);
			break;
		case _CR_SELECT_OP_RECV:
			assert(other->op == _CR_SELECT_OP_SEND);
			memcpy(this->val_ptr, other->val_ptr, this->val_siz);
			break;
		case _CR_SELECT_OP_DEFAULT:
			assert_notreached("_CR_SELECT_OP_DEFAULT is not CR_SELECT_CLASS_NONBLOCK");
		}
		struct _cr_select_waiter *waiter = other->waiter;
		for (size_t i = 0; i < waiter->arg_cnt; i++) {
			waiter->arg_vec[i].val.ch->nwaiters--;
			dlist_remove(&waiter->arg_vec[i].val.ch->waiters, &waiter->arg_vec[i]);
			if (&waiter->arg_vec[i].val == other)
				waiter->ret = i;
		}
		cr_unpause(waiter->cid);
		cr_yield();
		return choice_among_all;
	}

	if (cnt_default) {
		for (size_t i = 0; i < arg_cnt; i++)
			if (cr_select_getclass(&arg_vec[i].val) == CR_SELECT_CLASS_DEFAULT)
				return i;
		assert_notreached("should have returned from inside for() loop");
	}

	assert(cnt_blocking && cnt_blocking == arg_cnt);

	struct _cr_select_waiter waiter = {
		.cid     = cr_getcid(),
		.arg_vec = arg_vec,
		.arg_cnt = arg_cnt,
	};
	for (size_t i = 0; i < arg_cnt; i++) {
		arg_vec[i].val.waiter = &waiter;
		arg_vec[i].val.ch->nwaiters++;
		dlist_push_to_rear(&arg_vec[i].val.ch->waiters, &arg_vec[i]);
	}
	cr_pause_and_yield();
	return waiter.ret;
}