#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_report_sent_event.h"
#include "hid_transport_state_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_transport_state_data {
	bool ready;
	bool keys_ready;
	bool consumer_ready;
	enum keyboard_protocol_mode protocol_mode;
};

struct in_flight_report {
	bool active;
	enum hid_transport transport;
	enum keyboard_report_type report_type;
	uint16_t sequence;
};

static struct hid_transport_state_data transport_state[HID_TRANSPORT_COUNT] = {
	[HID_TRANSPORT_USB] = {
		.protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT,
	},
	[HID_TRANSPORT_BLE] = {
		.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;
static enum mode_switch_mode current_mode;
static uint16_t next_sequence;
static bool initialized;
static bool running;

static bool mode_to_transport(enum mode_switch_mode mode, enum hid_transport *transport)
{
	switch (mode) {
	case MODE_SWITCH_USB:
		*transport = HID_TRANSPORT_USB;
		return true;

	case MODE_SWITCH_BLE:
		*transport = HID_TRANSPORT_BLE;
		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 transport_can_send_report(enum keyboard_report_type report_type)
{
	enum hid_transport transport;
	struct hid_transport_state_data *state;

	if (!mode_to_transport(current_mode, &transport) || in_flight.active) {
		return false;
	}

	state = &transport_state[transport];

	if (!state->ready) {
		return false;
	}

	if (report_type == KEYBOARD_REPORT_TYPE_KEYS) {
		return state->keys_ready;
	}

	return state->consumer_ready;
}

static void submit_hid_tx_report_event(enum hid_transport transport,
				       enum keyboard_report_type report_type,
				       enum keyboard_protocol_mode protocol_mode,
				       const uint8_t *data, size_t size)
{
	struct hid_tx_report_event *event = new_hid_tx_report_event(size);

	event->transport = transport;
	event->report_type = report_type;
	event->protocol_mode = protocol_mode;
	event->sequence = next_sequence++;
	memcpy(event->dyndata.data, data, size);

	in_flight.active = true;
	in_flight.transport = transport;
	in_flight.report_type = report_type;
	in_flight.sequence = event->sequence;

	APP_EVENT_SUBMIT(event);
}

static void try_send_next(void)
{
	struct queued_report queued;
	enum hid_transport transport;
	struct hid_transport_state_data *state;

	if (!running || in_flight.active || !mode_to_transport(current_mode, &transport)) {
		return;
	}

	state = &transport_state[transport];

	if (!state->ready) {
		return;
	}

	if (pending_keys.valid && transport_can_send_report(KEYBOARD_REPORT_TYPE_KEYS)) {
		if (pending_keys.protocol_mode != state->protocol_mode) {
			LOG_WRN("Drop stale keys report after protocol change");
			pending_keys.valid = false;
		} else {
			submit_hid_tx_report_event(transport,
						   pending_keys.report_type,
						   pending_keys.protocol_mode,
						   pending_keys.data,
						   pending_keys.size);
			pending_keys.valid = false;
			return;
		}
	}

	if ((consumer_fifo_count > 0U) &&
	    transport_can_send_report(KEYBOARD_REPORT_TYPE_CONSUMER) &&
	    consumer_fifo_pop(&queued)) {
		if (queued.protocol_mode != state->protocol_mode) {
			LOG_WRN("Drop stale consumer report after protocol change");
		} else {
			submit_hid_tx_report_event(transport,
						   queued.report_type,
						   queued.protocol_mode,
						   queued.data,
						   queued.size);
			return;
		}
	}

	if (pending_consumer_latest.valid &&
	    transport_can_send_report(KEYBOARD_REPORT_TYPE_CONSUMER)) {
		if (pending_consumer_latest.protocol_mode != state->protocol_mode) {
			LOG_WRN("Drop stale latest consumer report after protocol change");
			pending_consumer_latest.valid = false;
		} else {
			submit_hid_tx_report_event(transport,
						   pending_consumer_latest.report_type,
						   pending_consumer_latest.protocol_mode,
						   pending_consumer_latest.data,
						   pending_consumer_latest.size);
			pending_consumer_latest.valid = false;
		}
	}
}

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_transport_state_event(const struct hid_transport_state_event *event)
{
	enum hid_transport active_transport;
	struct hid_transport_state_data *state;

	if (event->transport >= HID_TRANSPORT_COUNT) {
		return false;
	}

	state = &transport_state[event->transport];

	state->ready = event->ready;
	state->keys_ready = event->keys_ready;
	state->consumer_ready = event->consumer_ready;

	if (state->protocol_mode != event->protocol_mode) {
		state->protocol_mode = event->protocol_mode;

		if (mode_to_transport(current_mode, &active_transport) &&
		    (active_transport == event->transport)) {
			pending_keys.valid = false;
			pending_consumer_latest.valid = false;
			consumer_fifo_head = 0U;
			consumer_fifo_tail = 0U;
			consumer_fifo_count = 0U;
		}
	}

	if (!state->ready &&
	    mode_to_transport(current_mode, &active_transport) &&
	    (active_transport == event->transport)) {
		consumer_fifo_head = 0U;
		consumer_fifo_tail = 0U;
		consumer_fifo_count = 0U;
		in_flight.active = false;
	}

	try_send_next();
	return false;
}

static bool handle_hid_report_sent_event(const struct hid_report_sent_event *event)
{
	if (!in_flight.active || (event->transport != in_flight.transport)) {
		return false;
	}

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

	in_flight.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;
	transport_state[HID_TRANSPORT_USB].ready = false;
	transport_state[HID_TRANSPORT_USB].keys_ready = false;
	transport_state[HID_TRANSPORT_USB].consumer_ready = false;
	transport_state[HID_TRANSPORT_USB].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	transport_state[HID_TRANSPORT_BLE].ready = false;
	transport_state[HID_TRANSPORT_BLE].keys_ready = false;
	transport_state[HID_TRANSPORT_BLE].consumer_ready = false;
	transport_state[HID_TRANSPORT_BLE].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_transport_state_event(aeh)) {
		return handle_hid_transport_state_event(cast_hid_transport_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_transport_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);