#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 "module_lifecycle.h"
#include "transport_policy_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;
};

struct hid_flowctrl_module_ctx {
	struct module_lifecycle_ctx lc;
	struct hid_channel_state_data channel_state[HID_SEND_CH_COUNT];
	struct pending_report pending_keys;
	struct pending_report pending_consumer_latest;
	struct queued_report consumer_fifo[HID_FLOWCTRL_FIFO_DEPTH];
	uint8_t consumer_fifo_head;
	uint8_t consumer_fifo_tail;
	uint8_t consumer_fifo_count;
	struct in_flight_report in_flight[HID_SEND_CH_COUNT];
	enum hid_transport_policy current_transport;
	uint16_t next_sequence;
};

static int do_init(void);
static int do_start(void);
static int do_stop(void);

static const struct module_lifecycle_cfg lifecycle_cfg = {
	.mode = ML_MODE_POWER,
	.stopped_state = MODULE_STATE_STANDBY,
};

static const struct module_lifecycle_ops lifecycle_ops = {
	.do_init = do_init,
	.do_start = do_start,
	.do_stop = do_stop,
};

static struct hid_flowctrl_module_ctx ctx = {
	.lc = {
		.state = LC_UNINIT,
		.cfg = &lifecycle_cfg,
		.ops = &lifecycle_ops,
	},
	.channel_state = {
		[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 bool current_transport_to_channel(enum keyboard_report_type report_type,
					 enum hid_send_channel *channel)
{
	if (channel == NULL) {
		return false;
	}

	switch (ctx.current_transport) {
	case HID_TRANSPORT_POLICY_USB:
		*channel = (report_type == KEYBOARD_REPORT_TYPE_KEYS) ?
			   HID_SEND_CH_USB_KEYS :
			   HID_SEND_CH_USB_CONSUMER;
		return true;

	case HID_TRANSPORT_POLICY_BLE:
		*channel = HID_SEND_CH_BLE_SHARED;
		return true;

	default:
		return false;
	}
}

static void clear_pending_reports(void)
{
	memset(&ctx.pending_keys, 0, sizeof(ctx.pending_keys));
	memset(&ctx.pending_consumer_latest, 0, sizeof(ctx.pending_consumer_latest));
	ctx.consumer_fifo_head = 0U;
	ctx.consumer_fifo_tail = 0U;
	ctx.consumer_fifo_count = 0U;
	memset(&ctx.in_flight, 0, sizeof(ctx.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 (ctx.consumer_fifo_count == HID_FLOWCTRL_FIFO_DEPTH) {
		LOG_WRN("Consumer FIFO full, dropping oldest pulse");
		ctx.consumer_fifo_head =
			(ctx.consumer_fifo_head + 1U) % HID_FLOWCTRL_FIFO_DEPTH;
		ctx.consumer_fifo_count--;
	}

	struct queued_report *entry = &ctx.consumer_fifo[ctx.consumer_fifo_tail];

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

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

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

	*entry = ctx.consumer_fifo[ctx.consumer_fifo_head];
	ctx.consumer_fifo_head =
		(ctx.consumer_fifo_head + 1U) % HID_FLOWCTRL_FIFO_DEPTH;
	ctx.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 = &ctx.channel_state[channel];

	if (ctx.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 (!ctx.pending_keys.valid) {
		return;
	}

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

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

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

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

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

	if (!current_transport_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 == ctx.channel_state[channel].protocol_mode) {
			consumer_fifo_push(queued.report_type, queued.protocol_mode,
					   queued.data, queued.size);
		}
		return;
	}

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

static void try_send_consumer_latest(void)
{
	enum hid_send_channel channel;

	if (!ctx.pending_consumer_latest.valid) {
		return;
	}

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

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

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

static void try_send_next(void)
{
	if (!module_lifecycle_is_running(&ctx.lc)) {
		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 (!module_lifecycle_is_running(&ctx.lc) ||
	    ((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) {
		ctx.pending_keys.valid = true;
		ctx.pending_keys.report_type = event->report_type;
		ctx.pending_keys.protocol_mode = event->protocol_mode;
		ctx.pending_keys.size = event->dyndata.size;
		memcpy(ctx.pending_keys.data, event->dyndata.data,
		       event->dyndata.size);
	} else {
		ctx.pending_consumer_latest.valid = true;
		ctx.pending_consumer_latest.report_type = event->report_type;
		ctx.pending_consumer_latest.protocol_mode = event->protocol_mode;
		ctx.pending_consumer_latest.size = event->dyndata.size;
		memcpy(ctx.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;
	}

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

	if (event->report_ready_bm == 0U) {
		ctx.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 (!ctx.in_flight[event->channel].active) {
		return false;
	}

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

	ctx.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_transport_policy_event(
	const struct transport_policy_event *event)
{
	bool transport_changed =
		(ctx.current_transport != event->hid_transport);

	ctx.current_transport = event->hid_transport;

	if (transport_changed ||
	    (ctx.current_transport == HID_TRANSPORT_POLICY_NONE)) {
		clear_pending_reports();
	}

	try_send_next();
	return false;
}

static int do_init(void)
{
	clear_pending_reports();
	ctx.current_transport = HID_TRANSPORT_POLICY_USB;
	memset(ctx.channel_state, 0, sizeof(ctx.channel_state));
	ctx.channel_state[HID_SEND_CH_USB_KEYS].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	ctx.channel_state[HID_SEND_CH_USB_CONSUMER].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	ctx.channel_state[HID_SEND_CH_BLE_SHARED].protocol_mode =
		KEYBOARD_PROTOCOL_MODE_REPORT;
	ctx.next_sequence = 1U;
	return 0;
}

static int do_start(void)
{
	if (module_lifecycle_is_running(&ctx.lc)) {
		return 0;
	}

	try_send_next();
	return 0;
}

static int do_stop(void)
{
	if (!module_lifecycle_is_running(&ctx.lc)) {
		return 0;
	}

	clear_pending_reports();

	return 0;
}

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_transport_policy_event(aeh)) {
		return handle_transport_policy_event(
			cast_transport_policy_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)) {
			(void)module_set_lifecycle(&ctx.lc, LC_RUNNING);
		}

		return false;
	}

	if (is_power_down_event(aeh)) {
		if (module_lifecycle_is_initialized(&ctx.lc)) {
			(void)module_set_lifecycle(&ctx.lc, LC_STOPPED);
		}

		return false;
	}

	if (is_wake_up_event(aeh)) {
		if (module_lifecycle_is_initialized(&ctx.lc)) {
			(void)module_set_lifecycle(&ctx.lc, LC_RUNNING);
		}

		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, transport_policy_event);
APP_EVENT_SUBSCRIBE(MODULE, module_state_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, power_down_event);
APP_EVENT_SUBSCRIBE(MODULE, wake_up_event);