#include <stdbool.h>
#include <stdint.h>
#include <string.h>

#include <app_event_manager.h>

#define MODULE hid_flowctrl_module
#include <caf/events/module_state_event.h>
#include <caf/events/power_event.h>

#include <zephyr/logging/log.h>

#include "hid_channel_state_event.h"
#include "hid_report_sent_event.h"
#include "hid_tx_report_event.h"
#include "keyboard_core.h"
#include "keyboard_hid_report_event.h"
#include "mode_switch_event.h"

LOG_MODULE_REGISTER(MODULE, LOG_LEVEL_INF);

#define HID_FLOWCTRL_FIFO_DEPTH 32U
#define HID_FLOWCTRL_REPORT_DATA_MAX KEYBOARD_NKRO_REPORT_SIZE

struct pending_report {
	bool valid;
	enum keyboard_report_type report_type;
	enum keyboard_protocol_mode protocol_mode;
	size_t size;
	uint8_t data[HID_FLOWCTRL_REPORT_DATA_MAX];
};

struct queued_report {
	enum keyboard_report_type report_type;
	enum keyboard_protocol_mode protocol_mode;
	size_t size;
	uint8_t data[HID_FLOWCTRL_REPORT_DATA_MAX];
};

struct hid_channel_state_data {
	uint8_t report_ready_bm;
	enum keyboard_protocol_mode protocol_mode;
};

struct in_flight_report {
	bool active;
	enum hid_send_channel channel;
	enum keyboard_report_type report_type;
	uint16_t sequence;
};

static struct hid_channel_state_data channel_state[HID_SEND_CH_COUNT] = {
	[HID_SEND_CH_USB_KEYS] = {
		.protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT,
	},
	[HID_SEND_CH_USB_CONSUMER] = {
		.protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT,
	},
	[HID_SEND_CH_BLE_SHARED] = {
		.protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT,
	},
};
static struct pending_report pending_keys;
static struct pending_report pending_consumer_latest;
static struct queued_report consumer_fifo[HID_FLOWCTRL_FIFO_DEPTH];
static uint8_t consumer_fifo_head;
static uint8_t consumer_fifo_tail;
static uint8_t consumer_fifo_count;
static struct in_flight_report in_flight[HID_SEND_CH_COUNT];
static enum mode_switch_mode current_mode;
static uint16_t next_sequence;
static bool initialized;
static bool running;

static bool current_mode_to_channel(enum keyboard_report_type report_type,
				    enum hid_send_channel *channel)
{
	if (channel == NULL) {
		return false;
	}

	switch (current_mode) {
	case MODE_SWITCH_USB:
		*channel = (report_type == KEYBOARD_REPORT_TYPE_KEYS) ?
			   HID_SEND_CH_USB_KEYS :
			   HID_SEND_CH_USB_CONSUMER;
		return true;

	case MODE_SWITCH_BLE:
		*channel = HID_SEND_CH_BLE_SHARED;
		return true;

	default:
		return false;
	}
}

static void clear_pending_reports(void)
{
	memset(&pending_keys, 0, sizeof(pending_keys));
	memset(&pending_consumer_latest, 0, sizeof(pending_consumer_latest));
	consumer_fifo_head = 0U;
	consumer_fifo_tail = 0U;
	consumer_fifo_count = 0U;
	memset(&in_flight, 0, sizeof(in_flight));
}

static void consumer_fifo_push(enum keyboard_report_type report_type,
			       enum keyboard_protocol_mode protocol_mode,
			       const uint8_t *data, size_t size)
{
	if (consumer_fifo_count == HID_FLOWCTRL_FIFO_DEPTH) {
		LOG_WRN("Consumer FIFO full, dropping oldest pulse");
		consumer_fifo_head = (consumer_fifo_head + 1U) % HID_FLOWCTRL_FIFO_DEPTH;
		consumer_fifo_count--;
	}

	struct queued_report *entry = &consumer_fifo[consumer_fifo_tail];

	entry->report_type = report_type;
	entry->protocol_mode = protocol_mode;
	entry->size = size;
	memcpy(entry->data, data, size);

	consumer_fifo_tail = (consumer_fifo_tail + 1U) % HID_FLOWCTRL_FIFO_DEPTH;
	consumer_fifo_count++;
}

static bool consumer_fifo_pop(struct queued_report *entry)
{
	if (consumer_fifo_count == 0U) {
		return false;
	}

	*entry = consumer_fifo[consumer_fifo_head];
	consumer_fifo_head = (consumer_fifo_head + 1U) % HID_FLOWCTRL_FIFO_DEPTH;
	consumer_fifo_count--;

	return true;
}

static bool channel_can_send_report(enum hid_send_channel channel,
				    enum keyboard_report_type report_type,
				    enum keyboard_protocol_mode protocol_mode)
{
	const struct hid_channel_state_data *state = &channel_state[channel];

	if (in_flight[channel].active) {
		return false;
	}

	if (report_type == KEYBOARD_REPORT_TYPE_KEYS) {
		return (state->report_ready_bm & BIT(KEYBOARD_REPORT_TYPE_KEYS)) &&
		       (state->protocol_mode == protocol_mode);
	}

	if (channel == HID_SEND_CH_BLE_SHARED) {
		return (state->report_ready_bm & BIT(KEYBOARD_REPORT_TYPE_CONSUMER)) &&
		       (state->protocol_mode == KEYBOARD_PROTOCOL_MODE_REPORT);
	}

	return (state->report_ready_bm & BIT(KEYBOARD_REPORT_TYPE_CONSUMER)) != 0U;
}

static void try_send_keys(void)
{
	enum hid_send_channel channel;

	if (!pending_keys.valid) {
		return;
	}

	if (!current_mode_to_channel(KEYBOARD_REPORT_TYPE_KEYS, &channel)) {
		return;
	}

	if (!channel_can_send_report(channel, pending_keys.report_type,
				     pending_keys.protocol_mode)) {
		return;
	}

	in_flight[channel].active = true;
	in_flight[channel].channel = channel;
	in_flight[channel].report_type = pending_keys.report_type;
	in_flight[channel].sequence = next_sequence++;
	(void)submit_hid_tx_report_event(channel, pending_keys.report_type,
					 pending_keys.protocol_mode,
					 in_flight[channel].sequence,
					 pending_keys.data, pending_keys.size);
	pending_keys.valid = false;
}

static void try_send_consumer_fifo(void)
{
	struct queued_report queued;
	enum hid_send_channel channel;

	if (consumer_fifo_count == 0U) {
		return;
	}

	if (!current_mode_to_channel(KEYBOARD_REPORT_TYPE_CONSUMER, &channel)) {
		return;
	}

	if (!consumer_fifo_pop(&queued)) {
		return;
	}

	if (!channel_can_send_report(channel, queued.report_type,
				     queued.protocol_mode)) {
		if (queued.protocol_mode == channel_state[channel].protocol_mode) {
			consumer_fifo_push(queued.report_type, queued.protocol_mode,
					   queued.data, queued.size);
		}
		return;
	}

	in_flight[channel].active = true;
	in_flight[channel].channel = channel;
	in_flight[channel].report_type = queued.report_type;
	in_flight[channel].sequence = next_sequence++;
	(void)submit_hid_tx_report_event(channel, queued.report_type,
					 queued.protocol_mode,
					 in_flight[channel].sequence,
					 queued.data, queued.size);
}

static void try_send_consumer_latest(void)
{
	enum hid_send_channel channel;

	if (!pending_consumer_latest.valid) {
		return;
	}

	if (!current_mode_to_channel(KEYBOARD_REPORT_TYPE_CONSUMER, &channel)) {
		return;
	}

	if (!channel_can_send_report(channel,
				     pending_consumer_latest.report_type,
				     pending_consumer_latest.protocol_mode)) {
		return;
	}

	in_flight[channel].active = true;
	in_flight[channel].channel = channel;
	in_flight[channel].report_type = pending_consumer_latest.report_type;
	in_flight[channel].sequence = next_sequence++;
	(void)submit_hid_tx_report_event(channel,
					 pending_consumer_latest.report_type,
					 pending_consumer_latest.protocol_mode,
					 in_flight[channel].sequence,
					 pending_consumer_latest.data,
					 pending_consumer_latest.size);
	pending_consumer_latest.valid = false;
}

static void try_send_next(void)
{
	if (!running) {
		return;
	}

	try_send_keys();
	try_send_consumer_fifo();
	try_send_consumer_latest();
}

static bool handle_keyboard_hid_report_event(
	const struct keyboard_hid_report_event *event)
{
	if (!running ||
	    ((event->mode != MODE_SWITCH_USB) && (event->mode != MODE_SWITCH_BLE))) {
		return false;
	}

	if (event->queue_policy == HID_QUEUE_POLICY_FIFO) {
		consumer_fifo_push(event->report_type,
				   event->protocol_mode,
				   event->dyndata.data,
				   event->dyndata.size);
	} else if (event->report_type == KEYBOARD_REPORT_TYPE_KEYS) {
		pending_keys.valid = true;
		pending_keys.report_type = event->report_type;
		pending_keys.protocol_mode = event->protocol_mode;
		pending_keys.size = event->dyndata.size;
		memcpy(pending_keys.data, event->dyndata.data, event->dyndata.size);
	} else {
		pending_consumer_latest.valid = true;
		pending_consumer_latest.report_type = event->report_type;
		pending_consumer_latest.protocol_mode = event->protocol_mode;
		pending_consumer_latest.size = event->dyndata.size;
		memcpy(pending_consumer_latest.data, event->dyndata.data,
		       event->dyndata.size);
	}

	try_send_next();
	return false;
}

static bool handle_hid_channel_state_event(
	const struct hid_channel_state_event *event)
{
	if (event->channel >= HID_SEND_CH_COUNT) {
		return false;
	}

	channel_state[event->channel].report_ready_bm = event->report_ready_bm;
	channel_state[event->channel].protocol_mode = event->protocol_mode;

	if (event->report_ready_bm == 0U) {
		in_flight[event->channel].active = false;
	}

	try_send_next();
	return false;
}

static bool handle_hid_report_sent_event(const struct hid_report_sent_event *event)
{
	if (event->channel >= HID_SEND_CH_COUNT) {
		return false;
	}

	if (!in_flight[event->channel].active) {
		return false;
	}

	if (event->sequence != in_flight[event->channel].sequence) {
		LOG_WRN("Unexpected HID sent sequence %u (expected %u)",
			event->sequence, in_flight[event->channel].sequence);
		return false;
	}

	in_flight[event->channel].active = false;

	if (event->error) {
		LOG_WRN("HID report send failed for seq %u", event->sequence);
	}

	try_send_next();
	return false;
}

static bool handle_mode_switch_event(const struct mode_switch_event *event)
{
	bool mode_changed = (current_mode != event->mode);

	current_mode = event->mode;

	if (mode_changed || ((current_mode != MODE_SWITCH_USB) &&
			     (current_mode != MODE_SWITCH_BLE))) {
		clear_pending_reports();
	}

	try_send_next();
	return false;
}

static int module_init(void)
{
	clear_pending_reports();
	current_mode = MODE_SWITCH_USB;
	memset(channel_state, 0, sizeof(channel_state));
	channel_state[HID_SEND_CH_USB_KEYS].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	channel_state[HID_SEND_CH_USB_CONSUMER].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	channel_state[HID_SEND_CH_BLE_SHARED].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	next_sequence = 1U;
	return 0;
}

static int module_start(void)
{
	if (running) {
		return 0;
	}

	running = true;
	try_send_next();
	return 0;
}

static void module_pause(void)
{
	if (!running) {
		return;
	}

	clear_pending_reports();
	running = false;
}

static bool app_event_handler(const struct app_event_header *aeh)
{
	if (is_keyboard_hid_report_event(aeh)) {
		return handle_keyboard_hid_report_event(
			cast_keyboard_hid_report_event(aeh));
	}

	if (is_hid_channel_state_event(aeh)) {
		return handle_hid_channel_state_event(
			cast_hid_channel_state_event(aeh));
	}

	if (is_hid_report_sent_event(aeh)) {
		return handle_hid_report_sent_event(
			cast_hid_report_sent_event(aeh));
	}

	if (is_mode_switch_event(aeh)) {
		return handle_mode_switch_event(cast_mode_switch_event(aeh));
	}

	if (is_module_state_event(aeh)) {
		const struct module_state_event *event =
			cast_module_state_event(aeh);

		if (check_state(event, MODULE_ID(main), MODULE_STATE_READY)) {
			int err;

			if (!initialized) {
				err = module_init();
				if (err) {
					module_set_state(MODULE_STATE_ERROR);
					return false;
				}

				initialized = true;
			}

			err = module_start();
			if (err) {
				module_set_state(MODULE_STATE_ERROR);
			} else {
				module_set_state(MODULE_STATE_READY);
			}
		}

		return false;
	}

	if (is_power_down_event(aeh)) {
		if (initialized) {
			module_pause();
			module_set_state(MODULE_STATE_STANDBY);
		}

		return false;
	}

	if (is_wake_up_event(aeh)) {
		if (initialized) {
			int err = module_start();

			if (err) {
				module_set_state(MODULE_STATE_ERROR);
			} else {
				module_set_state(MODULE_STATE_READY);
			}
		}

		return false;
	}

	__ASSERT_NO_MSG(false);
	return false;
}

APP_EVENT_LISTENER(MODULE, app_event_handler);
APP_EVENT_SUBSCRIBE(MODULE, keyboard_hid_report_event);
APP_EVENT_SUBSCRIBE(MODULE, hid_channel_state_event);
APP_EVENT_SUBSCRIBE(MODULE, hid_report_sent_event);
APP_EVENT_SUBSCRIBE(MODULE, mode_switch_event);
APP_EVENT_SUBSCRIBE(MODULE, module_state_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, power_down_event);
APP_EVENT_SUBSCRIBE(MODULE, wake_up_event);