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

#include <app_event_manager.h>

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

#include <caf/key_id.h>
#include <zephyr/logging/log.h>
#include <zephyr/sys/byteorder.h>
#include <zephyr/sys/util.h>

#include "keyboard_core.h"
#include "keyboard_hid_report_event.h"
#include "mode_switch_event.h"
#include "set_protocol_event.h"

LOG_MODULE_REGISTER(MODULE, LOG_LEVEL_INF);

#define KEYBOARD_USAGE_FIRST_MODIFIER 0xE0U
#define KEYBOARD_USAGE_LAST_MODIFIER 0xE7U
#define KEYBOARD_USAGE_ERROR_ROLLOVER 0x01U
#define KEYBOARD_BOOT_RESERVED_BYTE 0x00U

enum key_usage_type {
	KEY_USAGE_TYPE_KEYBOARD,
	KEY_USAGE_TYPE_CONSUMER,
};

struct keymap_entry {
	uint16_t key_id;
	uint8_t usage_type;
	uint16_t usage_id;
};

struct keyboard_state {
	uint8_t modifiers;
	uint8_t keys_bitmap[KEYBOARD_NKRO_BITMAP_BYTES];
	uint32_t consumer_bits;
};

struct keyboard_reports_cache {
	uint8_t boot_report[KEYBOARD_BOOT_REPORT_SIZE];
	uint8_t nkro_report[KEYBOARD_NKRO_REPORT_SIZE];
	uint8_t consumer_report[KEYBOARD_CONSUMER_REPORT_SIZE];
	bool boot_valid;
	bool nkro_valid;
	bool consumer_valid;
};

static const struct keymap_entry keymap[] = {
	{ KEY_ID(0, 1), KEY_USAGE_TYPE_KEYBOARD, 0x0053 }, /* num lock */
	{ KEY_ID(0, 2), KEY_USAGE_TYPE_KEYBOARD, 0x005F }, /* keypad 7 */
	{ KEY_ID(0, 3), KEY_USAGE_TYPE_KEYBOARD, 0x005C }, /* keypad 4 */
	{ KEY_ID(0, 4), KEY_USAGE_TYPE_KEYBOARD, 0x0059 }, /* keypad 1 */
	{ KEY_ID(0, 5), KEY_USAGE_TYPE_KEYBOARD, 0x0062 }, /* keypad 0 */
	{ KEY_ID(1, 1), KEY_USAGE_TYPE_KEYBOARD, 0x0054 }, /* keypad / */
	{ KEY_ID(1, 2), KEY_USAGE_TYPE_KEYBOARD, 0x0060 }, /* keypad 8 */
	{ KEY_ID(1, 3), KEY_USAGE_TYPE_KEYBOARD, 0x005D }, /* keypad 5 */
	{ KEY_ID(1, 4), KEY_USAGE_TYPE_KEYBOARD, 0x005A }, /* keypad 2 */
	{ KEY_ID(1, 5), KEY_USAGE_TYPE_KEYBOARD, 0x0063 }, /* keypad . */
	{ KEY_ID(2, 1), KEY_USAGE_TYPE_KEYBOARD, 0x0055 }, /* keypad * */
	{ KEY_ID(2, 2), KEY_USAGE_TYPE_KEYBOARD, 0x0061 }, /* keypad 9 */
	{ KEY_ID(2, 3), KEY_USAGE_TYPE_KEYBOARD, 0x005E }, /* keypad 6 */
	{ KEY_ID(2, 4), KEY_USAGE_TYPE_KEYBOARD, 0x005B }, /* keypad 3 */
	{ KEY_ID(3, 0), KEY_USAGE_TYPE_CONSUMER, KEYBOARD_CONSUMER_CTRL_MUTE },
	{ KEY_ID(3, 1), KEY_USAGE_TYPE_KEYBOARD, 0x0056 }, /* keypad - */
	{ KEY_ID(3, 3), KEY_USAGE_TYPE_KEYBOARD, 0x0057 }, /* keypad + */
	{ KEY_ID(3, 5), KEY_USAGE_TYPE_KEYBOARD, 0x0058 }, /* keypad enter */
};

static const uint16_t consumer_usage_map[KEYBOARD_CONSUMER_CTRL_COUNT] = {
	[KEYBOARD_CONSUMER_CTRL_MUTE] = 0x00E2,
	[KEYBOARD_CONSUMER_CTRL_VOLUME_UP] = 0x00E9,
	[KEYBOARD_CONSUMER_CTRL_VOLUME_DOWN] = 0x00EA,
	[KEYBOARD_CONSUMER_CTRL_PLAY_PAUSE] = 0x00CD,
	[KEYBOARD_CONSUMER_CTRL_NEXT_TRACK] = 0x00B5,
	[KEYBOARD_CONSUMER_CTRL_PREV_TRACK] = 0x00B6,
};

static struct keyboard_state keyboard_state;
static struct keyboard_reports_cache reports_cache;
static enum keyboard_protocol_mode protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT;
static enum mode_switch_mode current_mode;
static bool initialized;
static bool running;
static bool mode_valid;

static const struct keymap_entry *keymap_get(uint16_t key_id)
{
	size_t left = 0;
	size_t right = ARRAY_SIZE(keymap);

	while (left < right) {
		size_t mid = left + ((right - left) / 2U);

		if (keymap[mid].key_id == key_id) {
			return &keymap[mid];
		}

		if (keymap[mid].key_id < key_id) {
			left = mid + 1U;
		} else {
			right = mid;
		}
	}

	return NULL;
}

static bool usage_is_modifier(uint16_t usage_id)
{
	return IN_RANGE(usage_id, KEYBOARD_USAGE_FIRST_MODIFIER,
			KEYBOARD_USAGE_LAST_MODIFIER);
}

static bool keyboard_key_update(uint16_t usage_id, bool pressed)
{
	if (usage_is_modifier(usage_id)) {
		uint8_t new_modifiers = keyboard_state.modifiers;

		WRITE_BIT(new_modifiers, usage_id - KEYBOARD_USAGE_FIRST_MODIFIER, pressed);
		if (new_modifiers == keyboard_state.modifiers) {
			return false;
		}

		keyboard_state.modifiers = new_modifiers;
		return true;
	}

	if (usage_id > KEYBOARD_NKRO_USAGE_MAX) {
		LOG_WRN("Unsupported keyboard usage 0x%04x", usage_id);
		return false;
	}

	uint8_t byte_idx = usage_id / 8U;
	uint8_t bit_idx = usage_id % 8U;
	bool was_pressed = (keyboard_state.keys_bitmap[byte_idx] & BIT(bit_idx)) != 0U;

	if (was_pressed == pressed) {
		return false;
	}

	WRITE_BIT(keyboard_state.keys_bitmap[byte_idx], bit_idx, pressed);
	return true;
}

static bool consumer_key_update(uint16_t consumer_id, bool pressed)
{
	if (consumer_id >= KEYBOARD_CONSUMER_CTRL_COUNT) {
		LOG_WRN("Unsupported consumer id %u", consumer_id);
		return false;
	}

	bool was_pressed = (keyboard_state.consumer_bits & BIT(consumer_id)) != 0U;

	if (was_pressed == pressed) {
		return false;
	}

	WRITE_BIT(keyboard_state.consumer_bits, consumer_id, pressed);
	return true;
}

static void keyboard_state_clear(void)
{
	memset(&keyboard_state, 0, sizeof(keyboard_state));
}

static void reports_cache_invalidate(void)
{
	reports_cache.boot_valid = false;
	reports_cache.nkro_valid = false;
	reports_cache.consumer_valid = false;
}

static void build_boot_report(uint8_t report[KEYBOARD_BOOT_REPORT_SIZE])
{
	size_t key_count = 0;

	memset(report, 0, KEYBOARD_BOOT_REPORT_SIZE);
	report[0] = keyboard_state.modifiers;
	report[1] = KEYBOARD_BOOT_RESERVED_BYTE;

	for (uint16_t usage_id = 0; usage_id <= KEYBOARD_NKRO_USAGE_MAX; usage_id++) {
		uint8_t byte_idx = usage_id / 8U;
		uint8_t bit_idx = usage_id % 8U;

		if ((keyboard_state.keys_bitmap[byte_idx] & BIT(bit_idx)) == 0U) {
			continue;
		}

		if (key_count == (KEYBOARD_BOOT_REPORT_SIZE - 2U)) {
			memset(&report[2], KEYBOARD_USAGE_ERROR_ROLLOVER,
			       KEYBOARD_BOOT_REPORT_SIZE - 2U);
			return;
		}

		report[2U + key_count] = (uint8_t)usage_id;
		key_count++;
	}
}

static void build_nkro_report(uint8_t report[KEYBOARD_NKRO_REPORT_SIZE])
{
	report[0] = keyboard_state.modifiers;
	memcpy(&report[1], keyboard_state.keys_bitmap, KEYBOARD_NKRO_BITMAP_BYTES);
}

static uint16_t active_consumer_usage_get(void)
{
	for (uint8_t consumer_id = 0; consumer_id < KEYBOARD_CONSUMER_CTRL_COUNT; consumer_id++) {
		if ((keyboard_state.consumer_bits & BIT(consumer_id)) != 0U) {
			return consumer_usage_map[consumer_id];
		}
	}

	return 0U;
}

static void build_consumer_report(uint8_t report[KEYBOARD_CONSUMER_REPORT_SIZE])
{
	sys_put_le16(active_consumer_usage_get(), report);
}

static void submit_keyboard_report_event(enum keyboard_report_type report_type,
					 const uint8_t *data, size_t size)
{
	struct keyboard_hid_report_event *event =
		new_keyboard_hid_report_event(size);

	event->mode = current_mode;
	event->report_type = report_type;
	event->protocol_mode = protocol_mode;
	memcpy(event->dyndata.data, data, size);

	APP_EVENT_SUBMIT(event);
}

static void emit_keys_report(bool force)
{
	uint8_t report_buf[KEYBOARD_NKRO_REPORT_SIZE];
	uint8_t report_size;
	uint8_t *cache_buf;
	bool *cache_valid;

	if (!mode_valid) {
		return;
	}

	if (protocol_mode == KEYBOARD_PROTOCOL_MODE_BOOT) {
		build_boot_report(report_buf);
		report_size = KEYBOARD_BOOT_REPORT_SIZE;
		cache_buf = reports_cache.boot_report;
		cache_valid = &reports_cache.boot_valid;
	} else {
		build_nkro_report(report_buf);
		report_size = KEYBOARD_NKRO_REPORT_SIZE;
		cache_buf = reports_cache.nkro_report;
		cache_valid = &reports_cache.nkro_valid;
	}

	if (!force && *cache_valid && (memcmp(cache_buf, report_buf, report_size) == 0)) {
		return;
	}

	memcpy(cache_buf, report_buf, report_size);
	*cache_valid = true;

	submit_keyboard_report_event(KEYBOARD_REPORT_TYPE_KEYS, report_buf, report_size);
}

static void emit_consumer_report(bool force)
{
	uint8_t report_buf[KEYBOARD_CONSUMER_REPORT_SIZE];

	if (!mode_valid) {
		return;
	}

	build_consumer_report(report_buf);
	if (!force && reports_cache.consumer_valid &&
	    (memcmp(reports_cache.consumer_report, report_buf,
		    KEYBOARD_CONSUMER_REPORT_SIZE) == 0)) {
		return;
	}

	memcpy(reports_cache.consumer_report, report_buf, KEYBOARD_CONSUMER_REPORT_SIZE);
	reports_cache.consumer_valid = true;

	submit_keyboard_report_event(KEYBOARD_REPORT_TYPE_CONSUMER,
				     report_buf,
				     KEYBOARD_CONSUMER_REPORT_SIZE);
}

static void emit_all_reports(bool force)
{
	emit_keys_report(force);
	emit_consumer_report(force);
}

static void emit_release_reports(enum mode_switch_mode mode)
{
	struct keyboard_hid_report_event *event;
	uint8_t keys_report[KEYBOARD_NKRO_REPORT_SIZE] = { 0 };
	uint8_t consumer_report[KEYBOARD_CONSUMER_REPORT_SIZE] = { 0 };
	size_t keys_report_size =
		(protocol_mode == KEYBOARD_PROTOCOL_MODE_BOOT) ?
			KEYBOARD_BOOT_REPORT_SIZE : KEYBOARD_NKRO_REPORT_SIZE;

	event = new_keyboard_hid_report_event(keys_report_size);
	event->mode = mode;
	event->report_type = KEYBOARD_REPORT_TYPE_KEYS;
	event->protocol_mode = protocol_mode;
	memcpy(event->dyndata.data, keys_report, keys_report_size);
	APP_EVENT_SUBMIT(event);

	event = new_keyboard_hid_report_event(KEYBOARD_CONSUMER_REPORT_SIZE);
	event->mode = mode;
	event->report_type = KEYBOARD_REPORT_TYPE_CONSUMER;
	event->protocol_mode = protocol_mode;
	memcpy(event->dyndata.data, consumer_report, KEYBOARD_CONSUMER_REPORT_SIZE);
	APP_EVENT_SUBMIT(event);
}

static int module_init(void)
{
	keyboard_state_clear();
	reports_cache_invalidate();
	mode_valid = false;
	protocol_mode = KEYBOARD_PROTOCOL_MODE_REPORT;

	return 0;
}

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

	running = true;
	return 0;
}

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

	if (mode_valid) {
		emit_release_reports(current_mode);
	}

	keyboard_state_clear();
	reports_cache_invalidate();
	mode_valid = false;
	running = false;
}

static bool handle_button_event(const struct button_event *event)
{
	const struct keymap_entry *entry;
	bool changed;

	if (!running) {
		return false;
	}

	entry = keymap_get(event->key_id);
	if (!entry) {
		LOG_WRN("Unmapped key id 0x%04x", event->key_id);
		return false;
	}

	if (entry->usage_type == KEY_USAGE_TYPE_KEYBOARD) {
		changed = keyboard_key_update(entry->usage_id, event->pressed);
		if (changed) {
			emit_keys_report(false);
		}
	} else {
		changed = consumer_key_update(entry->usage_id, event->pressed);
		if (changed) {
			emit_consumer_report(false);
		}
	}

	return false;
}

static bool handle_mode_switch_event(const struct mode_switch_event *event)
{
	bool mode_changed;

	if (!running) {
		current_mode = event->mode;
		return false;
	}

	mode_changed = mode_valid && (current_mode != event->mode);
	if (mode_changed) {
		emit_release_reports(current_mode);
		keyboard_state_clear();
		reports_cache_invalidate();
	}

	current_mode = event->mode;
	mode_valid = true;
	emit_all_reports(true);

	return false;
}

static bool app_event_handler(const struct app_event_header *aeh)
{
	if (is_button_event(aeh)) {
		return handle_button_event(cast_button_event(aeh));
	}

	if (is_set_protocol_event(aeh)) {
		const struct set_protocol_event *event = cast_set_protocol_event(aeh);

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

			if (running && mode_valid && (current_mode == MODE_SWITCH_USB)) {
				emit_keys_report(true);
			}
		}

		return false;
	}

	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, button_event);
APP_EVENT_SUBSCRIBE(MODULE, set_protocol_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);