#include <errno.h>
#include <string.h>

#include <zephyr/device.h>
#include <zephyr/kernel.h>
#include <zephyr/usb/class/usbd_hid.h>
#include <zephyr/usb/usbd.h>

#include <app_event_manager.h>
#include <caf/events/power_event.h>

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

#include "hid_report_descriptor.h"
#include "hid_protocol_event.h"
#include "hid_report_event.h"
#include "keyboard_led_state_event.h"
#include "mode_event.h"

#include <zephyr/logging/log.h>
LOG_MODULE_REGISTER(MODULE, LOG_LEVEL_INF);

#define APP_USB_VID 0x1209
#define APP_USB_PID 0x0001

/*
 * 模块目标：
 * 1) 模块内聚控制 USB HID 栈生命周期（初始化/启用/禁用）。
 * 2) 对外发布 HID 协议事件和 NumLock 指示事件。
 * 3) 仅响应 mode_event（USB/BLE/2.4G）和 power_event（休眠/唤醒）。
 *
 * 约束：
 * - 启动时只做 USB 设备初始化，不自动 enable USB 栈；
 * - 只有当 mode 切到 USB 且系统非休眠时才 enable；
 * - BLE 逻辑保持不变，不在本模块中触碰。
 */
struct usb_hid_ctx {
	const struct device *boot_dev;
	const struct device *nkro_dev;
	const struct device *raw_dev;

	bool stack_initialized;
	bool stack_enabled;
	bool stack_error;

	bool usb_mode_selected;
	bool pm_suspended;

	bool boot_iface_ready;
	bool nkro_iface_ready;
	bool raw_iface_ready;
	bool boot_in_flight;
	bool nkro_in_flight;

	enum hid_protocol_type current_protocol;
};

static struct usb_hid_ctx g_usb_hid = {
	.current_protocol = HID_PROTO_REPORT,
};

USBD_DEVICE_DEFINE(new_kbd_usbd,
		   DEVICE_DT_GET(DT_NODELABEL(usbd)),
		   APP_USB_VID, APP_USB_PID);
USBD_DESC_LANG_DEFINE(new_kbd_lang);
USBD_DESC_MANUFACTURER_DEFINE(new_kbd_mfr, "new_kbd");
USBD_DESC_PRODUCT_DEFINE(new_kbd_product, "new_kbd composite HID");
USBD_DESC_CONFIG_DEFINE(new_kbd_fs_cfg_desc, "FS Configuration");
USBD_CONFIGURATION_DEFINE(new_kbd_fs_config, 0, 100, &new_kbd_fs_cfg_desc);

static const uint8_t boot_report_desc[] = HID_KEYBOARD_REPORT_DESC();
static const uint8_t nkro_report_desc[] = HID_DESC_KEYBOARD_NKRO_CONSUMER();
static const uint8_t raw_report_desc[] = HID_DESC_RAW_64();

/* 统一入口仅处理单字节 LED 报告并发布事件。 */
static void process_usb_led_input_report(uint8_t led_report)
{
	struct keyboard_led_state_event *event = new_keyboard_led_state_event();

	event->led_mask = led_report;
	event->num_lock = (led_report & BIT(0)) != 0U;
	APP_EVENT_SUBMIT(event);
}

static bool should_handle_led_input_from_dev(const struct device *dev)
{
	if (g_usb_hid.current_protocol == HID_PROTO_BOOT)
		return (dev == g_usb_hid.boot_dev);

	return (dev == g_usb_hid.nkro_dev);
}

static bool try_extract_led_mask(const struct device *dev,
				 uint16_t len,
				 const uint8_t *buf,
				 uint8_t *led_mask)
{
	if ((buf == NULL) || (len == 0U))
		return false;

	if (dev == g_usb_hid.boot_dev) {
		*led_mask = buf[0];
		return true;
	}

	if (dev != g_usb_hid.nkro_dev)
		return false;

	if (len >= 2U) {
		if (buf[0] != REPORT_ID_KEYBOARD)
			return false;

		*led_mask = buf[1];
		return true;
	}

	*led_mask = buf[0];
	return true;
}

static int hid_stub_get_report(const struct device *dev,
			       uint8_t type, uint8_t id,
			       uint16_t len, uint8_t *buf)
{
	ARG_UNUSED(dev);
	ARG_UNUSED(type);
	ARG_UNUSED(id);
	ARG_UNUSED(len);
	ARG_UNUSED(buf);
	return -ENOTSUP;
}

static int hid_stub_set_report(const struct device *dev,
			       uint8_t type, uint8_t id,
			       uint16_t len, const uint8_t *buf)
{
	ARG_UNUSED(type);
	ARG_UNUSED(id);

	if (!should_handle_led_input_from_dev(dev)) {
		return 0;
	}

	uint8_t led_mask;

	if (!try_extract_led_mask(dev, len, buf, &led_mask)) {
		return 0;
	}

	LOG_INF("hid_stub_set_report led_mask=0x%02x", led_mask);
	process_usb_led_input_report(led_mask);

	return 0;
}

static void hid_stub_set_idle(const struct device *dev, uint8_t id, uint32_t duration)
{
	ARG_UNUSED(dev);
	ARG_UNUSED(id);
	ARG_UNUSED(duration);
}

static uint32_t hid_stub_get_idle(const struct device *dev, uint8_t id)
{
	ARG_UNUSED(dev);
	ARG_UNUSED(id);
	return 0;
}

static void hid_stub_set_protocol(const struct device *dev, uint8_t proto)
{
	ARG_UNUSED(dev);

	enum hid_protocol_type new_protocol =
		(proto == HID_PROTOCOL_BOOT) ? HID_PROTO_BOOT : HID_PROTO_REPORT;

	if (g_usb_hid.current_protocol == new_protocol) {
		return;
	}

	g_usb_hid.current_protocol = new_protocol;

	/*
	 * 按需求：USB HID 在连接后收到 set_protocol 时上报 hid_protocol_event。
	 * 这里额外检查接口 ready，避免在未枚举完成阶段上报无意义协议切换。
	 */
	if (g_usb_hid.boot_iface_ready || g_usb_hid.nkro_iface_ready) {
		struct hid_protocol_event *event = new_hid_protocol_event();

		event->transport = HID_TRANSPORT_USB;
		event->protocol = new_protocol;
		APP_EVENT_SUBMIT(event);
	}
}

static void hid_stub_input_done(const struct device *dev, const uint8_t *report)
{
	ARG_UNUSED(report);

	/*
	 * 发送完成回调：
	 * - 仅在这里清除“在途发送”标志，确保“上一包未完成则丢弃新包”的策略可闭环；
	 * - 若收到未知 dev 的回调，仅记录告警，避免静默状态错乱。
	 */
	if (dev == g_usb_hid.boot_dev) {
		g_usb_hid.boot_in_flight = false;
		return;
	}

	if (dev == g_usb_hid.nkro_dev) {
		g_usb_hid.nkro_in_flight = false;
		return;
	}

	LOG_WRN("input_done from unknown HID dev: %p", (void *)dev);
}

static void hid_stub_output_report(const struct device *dev, uint16_t len, const uint8_t *buf)
{
	if (!should_handle_led_input_from_dev(dev)) {
		return;
	}

	uint8_t led_mask;

	if (!try_extract_led_mask(dev, len, buf, &led_mask)) {
		return;
	}

	LOG_INF("hid_stub_output_report led_mask=0x%02x", led_mask);
	process_usb_led_input_report(led_mask);
}

static void hid_iface_ready_cb(const struct device *dev, bool ready)
{
	if (dev == g_usb_hid.boot_dev) {
		g_usb_hid.boot_iface_ready = ready;
		if (!ready) {
			g_usb_hid.boot_in_flight = false;
		}
	} else if (dev == g_usb_hid.nkro_dev) {
		g_usb_hid.nkro_iface_ready = ready;
		if (!ready) {
			g_usb_hid.nkro_in_flight = false;
		}
	} else if (dev == g_usb_hid.raw_dev) {
		g_usb_hid.raw_iface_ready = ready;
	}

	if (ready) {
		/* 连接可用后同步一次当前协议，让 keyboard_module 与传输侧编码一致。 */
		struct hid_protocol_event *event = new_hid_protocol_event();

		event->transport = HID_TRANSPORT_USB;
		event->protocol = g_usb_hid.current_protocol;
		APP_EVENT_SUBMIT(event);
	}

}

static const struct hid_device_ops boot_hid_ops = {
	.iface_ready = hid_iface_ready_cb,
	.get_report = hid_stub_get_report,
	.set_report = hid_stub_set_report,
	.set_idle = hid_stub_set_idle,
	.get_idle = hid_stub_get_idle,
	.set_protocol = hid_stub_set_protocol,
	.input_report_done = hid_stub_input_done,
	.output_report = hid_stub_output_report,
};

static const struct hid_device_ops report_hid_ops = {
	.iface_ready = hid_iface_ready_cb,
	.get_report = hid_stub_get_report,
	.set_report = hid_stub_set_report,
	.set_idle = hid_stub_set_idle,
	.get_idle = hid_stub_get_idle,
	.set_protocol = hid_stub_set_protocol,
	.input_report_done = hid_stub_input_done,
	.output_report = hid_stub_output_report,
};

static const struct hid_device_ops raw_hid_ops = {
	.iface_ready = hid_iface_ready_cb,
	.get_report = hid_stub_get_report,
	.set_report = hid_stub_set_report,
	.set_idle = hid_stub_set_idle,
	.get_idle = hid_stub_get_idle,
	.input_report_done = hid_stub_input_done,
	.output_report = hid_stub_output_report,
};

static void usbd_msg_cb(struct usbd_context *const usbd_ctx,
			const struct usbd_msg *const msg)
{
	switch (msg->type) {
	case USBD_MSG_VBUS_READY:
		if (g_usb_hid.pm_suspended) {
			LOG_INF("VBUS ready: submit wake_up_event");
			APP_EVENT_SUBMIT(new_wake_up_event());
		}

		/*
		 * 只有在 USB 模式下才允许拉起 USB 栈。
		 * 这样即使插着线，只要用户切到 BLE/2.4G，也不会强制进入 USB HID。
		 */
		if (usbd_can_detect_vbus(usbd_ctx) && g_usb_hid.stack_enabled) {
			(void)usbd_enable(usbd_ctx);
		}
		break;

	case USBD_MSG_VBUS_REMOVED:
		break;

	case USBD_MSG_SUSPEND:
	case USBD_MSG_RESUME:
	case USBD_MSG_CONFIGURATION:
		break;

	case USBD_MSG_UDC_ERROR:
	case USBD_MSG_STACK_ERROR:
		LOG_ERR("USBD stack error message: %d", msg->type);
		g_usb_hid.stack_error = true;
		break;

	default:
		break;
	}
}

static bool usb_hid_devices_ready(void)
{
	if (!device_is_ready(g_usb_hid.boot_dev)) {
		LOG_ERR("HID boot device is not ready");
		return false;
	}

	if (!device_is_ready(g_usb_hid.nkro_dev)) {
		LOG_ERR("HID nkro device is not ready");
		return false;
	}

	if (!device_is_ready(g_usb_hid.raw_dev)) {
		LOG_ERR("HID raw device is not ready");
		return false;
	}

	if (!device_is_ready(DEVICE_DT_GET(DT_NODELABEL(usbd)))) {
		LOG_ERR("USBD device is not ready");
		return false;
	}

	return true;
}

static int usb_hid_register_hid_devices(void)
{
	int err = hid_device_register(g_usb_hid.boot_dev,
				      boot_report_desc, sizeof(boot_report_desc),
				      &boot_hid_ops);
	if (err) {
		LOG_ERR("hid_device_register(boot) failed: %d", err);
		return err;
	}

	err = hid_device_register(g_usb_hid.nkro_dev,
				  nkro_report_desc, sizeof(nkro_report_desc),
				  &report_hid_ops);
	if (err) {
		LOG_ERR("hid_device_register(nkro) failed: %d", err);
		return err;
	}

	err = hid_device_register(g_usb_hid.raw_dev,
				  raw_report_desc, sizeof(raw_report_desc),
				  &raw_hid_ops);
	if (err) {
		LOG_ERR("hid_device_register(raw) failed: %d", err);
		return err;
	}

	return 0;
}

static int usb_hid_configure_usbd(void)
{
	int err = usbd_add_descriptor(&new_kbd_usbd, &new_kbd_lang);
	if (err) {
		LOG_ERR("usbd_add_descriptor(lang) failed: %d", err);
		return err;
	}

	err = usbd_add_descriptor(&new_kbd_usbd, &new_kbd_mfr);
	if (err) {
		LOG_ERR("usbd_add_descriptor(mfr) failed: %d", err);
		return err;
	}

	err = usbd_add_descriptor(&new_kbd_usbd, &new_kbd_product);
	if (err) {
		LOG_ERR("usbd_add_descriptor(product) failed: %d", err);
		return err;
	}

	err = usbd_add_configuration(&new_kbd_usbd, USBD_SPEED_FS, &new_kbd_fs_config);
	if (err) {
		LOG_ERR("usbd_add_configuration failed: %d", err);
		return err;
	}

	err = usbd_register_all_classes(&new_kbd_usbd, USBD_SPEED_FS, 1, NULL);
	if (err) {
		LOG_ERR("usbd_register_all_classes failed: %d", err);
		return err;
	}

	return 0;
}

static int usb_hid_init_usbd_stack(void)
{
	int err;

	usbd_device_set_code_triple(&new_kbd_usbd, USBD_SPEED_FS, 0, 0, 0);

	err = usbd_msg_register_cb(&new_kbd_usbd, usbd_msg_cb);
	if (err) {
		LOG_ERR("usbd_msg_register_cb failed: %d", err);
		return err;
	}

	err = usbd_init(&new_kbd_usbd);
	if (err && (err != -EALREADY)) {
		LOG_ERR("usbd_init failed: %d", err);
		return err;
	}

	return 0;
}

static int usb_hid_stack_init(void)
{
	if (g_usb_hid.stack_initialized) {
		return 0;
	}

	g_usb_hid.boot_dev = DEVICE_DT_GET(DT_NODELABEL(hid_dev_0));
	g_usb_hid.nkro_dev = DEVICE_DT_GET(DT_NODELABEL(hid_dev_1));
	g_usb_hid.raw_dev = DEVICE_DT_GET(DT_NODELABEL(raw_hid));

	if (!usb_hid_devices_ready()) {
		return -ENODEV;
	}

	int err = usb_hid_register_hid_devices();
	if (err) {
		return err;
	}

	err = usb_hid_configure_usbd();
	if (err) {
		return err;
	}

	err = usb_hid_init_usbd_stack();
	if (err) {
		return err;
	}

	g_usb_hid.stack_initialized = true;
	return 0;
}

static int usb_hid_set_enabled(bool enable)
{
	int err;

	if (!g_usb_hid.stack_initialized) {
		err = usb_hid_stack_init();
		if (err) {
			return err;
		}
	}

	if (g_usb_hid.stack_enabled == enable) {
		return 0;
	}

	g_usb_hid.stack_enabled = enable;

	if (enable) {
		err = usbd_enable(&new_kbd_usbd);
	} else {
		err = usbd_disable(&new_kbd_usbd);
		g_usb_hid.boot_iface_ready = false;
		g_usb_hid.nkro_iface_ready = false;
		g_usb_hid.raw_iface_ready = false;
		g_usb_hid.boot_in_flight = false;
		g_usb_hid.nkro_in_flight = false;
	}

	if (err && (err != -EALREADY)) {
		LOG_ERR("usbd_%s failed: %d", enable ? "enable" : "disable", err);
		g_usb_hid.stack_error = true;
		return err;
	}

	return 0;
}

static void refresh_usb_state_by_policy(void)
{
	/*
	 * 控制策略：
	 * - USB 模式 + 非休眠：启用 USB HID。
	 * - 其他情况：关闭 USB HID（不销毁初始化结果，后续可快速恢复）。
	 */
	bool should_enable = g_usb_hid.usb_mode_selected && !g_usb_hid.pm_suspended;
	int err = usb_hid_set_enabled(should_enable);

	if (err) {
		LOG_ERR("usb_hid_set_enabled(%d) failed: %d", should_enable, err);
	}
}

static bool handle_module_state_event(const struct module_state_event *event)
{
	if (!check_state(event, MODULE_ID(main), MODULE_STATE_READY)) {
		return false;
	}

	int err = usb_hid_stack_init();
	if (err) {
		LOG_ERR("USB HID stack init failed: %d", err);
		g_usb_hid.stack_error = true;
		module_set_state(MODULE_STATE_ERROR);
		return false;
	}

	module_set_state(MODULE_STATE_READY);
	return false;
}

static bool handle_mode_event(const struct mode_event *event)
{
	g_usb_hid.usb_mode_selected = (event->mode_type == MODE_TYPE_USB);
	refresh_usb_state_by_policy();
	return false;
}

static bool handle_power_down_event(void)
{
	if (g_usb_hid.pm_suspended) {
		/* 避免重复上报 STANDBY 导致 power_manager 在 SUSPENDING 期间反复迭代。 */
		return false;
	}

	g_usb_hid.pm_suspended = true;
	refresh_usb_state_by_policy();
	module_set_state(MODULE_STATE_STANDBY);
	return false;
}

static bool handle_wake_up_event(void)
{
	if (!g_usb_hid.pm_suspended) {
		return false;
	}

	g_usb_hid.pm_suspended = false;
	refresh_usb_state_by_policy();
	module_set_state(MODULE_STATE_READY);
	return false;
}

static bool handle_hid_report_event(const struct hid_report_event *event)
{
	/*
	 * USB 侧仅在 active 条件满足时发送：
	 * - 当前 mode 为 USB；
	 * - USB HID 栈已启用且对应接口 ready。
	 */
	if (!g_usb_hid.usb_mode_selected || !g_usb_hid.stack_enabled) {
		return false;
	}

	uint8_t report_id;

	if (event->protocol != g_usb_hid.current_protocol) {
		return false;
	}

	if (event->protocol == HID_PROTO_BOOT) {
		const uint8_t *payload = event->dyndata.data;
		size_t payload_len = event->dyndata.size;

		if (!g_usb_hid.boot_iface_ready || !g_usb_hid.boot_dev) {
			return false;
		}
		if (g_usb_hid.boot_in_flight) {
			LOG_WRN("Drop boot report: previous report not sent");
			return false;
		}

		int err = hid_device_submit_report(g_usb_hid.boot_dev,
						   payload_len,
						   payload);
		if (err) {
			LOG_WRN("USB boot report send failed err=%d", err);
		} else {
			g_usb_hid.boot_in_flight = true;
		}
		return false;
	}

	if (event->dyndata.size < 1U) {
		return false;
	}

	report_id = event->dyndata.data[0];

	if (!g_usb_hid.nkro_iface_ready || !g_usb_hid.nkro_dev) {
		return false;
	}

	if ((report_id != REPORT_ID_KEYBOARD) && (report_id != REPORT_ID_CONSUMER)) {
		return false;
	}
	if (g_usb_hid.nkro_in_flight) {
		LOG_WRN("Drop report id=0x%02x: previous report not sent", report_id);
		return false;
	}

	/* Report 协议下 dyndata 是 [report_id|payload]，可直接透传。 */
	int err = hid_device_submit_report(g_usb_hid.nkro_dev, event->dyndata.size, event->dyndata.data);
	if (err) {
		LOG_WRN("USB report send failed id=0x%02x err=%d", report_id, err);
	} else {
		g_usb_hid.nkro_in_flight = true;
	}

	return false;
}

static bool app_event_handler(const struct app_event_header *aeh)
{
	if (is_module_state_event(aeh)) {
		return handle_module_state_event(cast_module_state_event(aeh));
	}

	if (is_mode_event(aeh)) {
		return handle_mode_event(cast_mode_event(aeh));
	}

	if (is_power_down_event(aeh)) {
		return handle_power_down_event();
	}

	if (is_wake_up_event(aeh)) {
		return handle_wake_up_event();
	}

	if (is_hid_report_event(aeh)) {
		return handle_hid_report_event(cast_hid_report_event(aeh));
	}

	__ASSERT_NO_MSG(false);
	return false;
}

APP_EVENT_LISTENER(MODULE, app_event_handler);
APP_EVENT_SUBSCRIBE(MODULE, module_state_event);
APP_EVENT_SUBSCRIBE(MODULE, mode_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, power_down_event);
APP_EVENT_SUBSCRIBE(MODULE, wake_up_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, hid_report_event);