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

#include <libcr_ipc/select.h>

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->waiter_typ == _CR_CHAN_SENDER)
			return _CR_SELECT_CLASS_NONBLOCK;
		else
			return _CR_SELECT_CLASS_BLOCKING;
	case _CR_SELECT_OP_SEND:
		if (arg.ch->waiters.front && arg.ch->waiter_typ == _CR_CHAN_RECVER)
			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");
	}
}

void _cr_select_dequeue(void *_waiters, size_t idx) {
	struct _cr_select_waiters *waiters = _waiters;
	for (size_t i = 0; i < waiters->cnt; i++)
		_cr_ipc_dll_remove(&(waiters->args[i].ch->waiters),
		                   &(waiters->nodes[i]));
	waiters->cnt = idx;
}

size_t cr_select_v(size_t arg_cnt, struct cr_select_arg 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])) {
		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 = rand_uint63n(cnt_nonblock);
		for (size_t i = 0, seen = 0; i < arg_cnt; i++) {
			if (_cr_select_getclass(arg_vec[i]) == _CR_SELECT_CLASS_NONBLOCK) {
				if (seen == choice) {
					_cr_chan_xfer(arg_vec[i].op == _CR_SELECT_OP_RECV
					              ? _CR_CHAN_RECVER
					              : _CR_CHAN_SENDER,
					              arg_vec[i].ch,
					              arg_vec[i].val_ptr,
					              arg_vec[i].val_siz);
					return i;
				}
				seen++;
			}
		}
		assert_notreached("should have returned from inside for() loop");
	}

	if (cnt_default) {
		for (size_t i = 0; i < arg_cnt; i++)
			if (_cr_select_getclass(arg_vec[i]) == _CR_SELECT_CLASS_DEFAULT)
				return i;
		assert_notreached("should have returned from inside for() loop");
	}

	struct _cr_select_waiters waiters = {
		.cnt   = arg_cnt,
		.args  = arg_vec,
		.nodes = alloca(sizeof(struct _cr_chan_waiter) * arg_cnt),
	};
	for (size_t i = 0; i < arg_cnt; i++) {
		waiters.nodes[i] = (struct _cr_chan_waiter){
			.cid          = cr_getcid(),
			.val_ptr      = arg_vec[i].val_ptr,
			.dequeue      = _cr_select_dequeue,
			.dequeue_arg1 = &waiters,
			.dequeue_arg2 = i,
		};
		_cr_ipc_dll_push_to_rear(&arg_vec[i].ch->waiters, &waiters.nodes[i]);
	}
	cr_pause_and_yield();
	return waiters.cnt;
}