#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_boot_event.h"
#include "hid_protocol_event.h"
#include "hid_tx_done_event.h"
#include "hid_tx_event.h"
#include "hid_vendor_mask_event.h"
#include "keyboard_led_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_iface {
	const struct device *dev;
	bool iface_ready;
	bool in_flight;
};

enum usb_hid_stack_state {
	USB_HID_STACK_STATE_OFF,
	USB_HID_STACK_STATE_READY,
	USB_HID_STACK_STATE_ACTIVE,
	USB_HID_STACK_STATE_ERROR,
};

struct usb_hid_policy {
	bool usb_mode_selected;
	bool pm_suspended;
};

struct usb_hid_ctx {
	struct usb_hid_iface boot;
	struct usb_hid_iface nkro;
	struct usb_hid_iface raw;
	enum usb_hid_stack_state stack_state;
	struct usb_hid_policy policy;
	enum hid_protocol_type current_protocol;
};

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

static void submit_usb_tx_done(enum hid_tx_kind kind, bool success)
{
	hid_tx_done_event_submit(kind, success);
}

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();

static bool usb_hid_stack_is_active(void)
{
	return g_usb_hid.stack_state == USB_HID_STACK_STATE_ACTIVE;
}

static bool usb_hid_stack_is_error(void)
{
	return g_usb_hid.stack_state == USB_HID_STACK_STATE_ERROR;
}

static bool usb_hid_should_be_active(void)
{
	return g_usb_hid.policy.usb_mode_selected && !g_usb_hid.policy.pm_suspended;
}

static struct usb_hid_iface *usb_hid_iface_from_dev(const struct device *dev)
{
	if (dev == g_usb_hid.boot.dev) {
		return &g_usb_hid.boot;
	}

	if (dev == g_usb_hid.nkro.dev) {
		return &g_usb_hid.nkro;
	}

	if (dev == g_usb_hid.raw.dev) {
		return &g_usb_hid.raw;
	}

	return NULL;
}

static void usb_hid_clear_runtime_iface_state(void)
{
	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;
	g_usb_hid.raw.in_flight = false;
}

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 bool try_extract_vendor_mask(const struct device *dev,
				    uint16_t len,
				    const uint8_t *buf,
				    const uint8_t **mask_data,
				    size_t *mask_len)
{
	if ((buf == NULL) || (len < 1U)) {
		return false;
	}

	if (dev != g_usb_hid.nkro.dev) {
		return false;
	}

	if (buf[0] != REPORT_ID_VENDOR) {
		return false;
	}

	if ((len - 1U) != HID_VENDOR_PAYLOAD_SIZE) {
		return false;
	}

	*mask_data = &buf[1];
	*mask_len = len - 1U;
	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)) {
		const uint8_t *mask_data;
		size_t mask_len;

		if (try_extract_vendor_mask(dev, len, buf, &mask_data, &mask_len)) {
			LOG_INF("hid_stub_set_report vendor mask len=%u", mask_len);
			hid_vendor_mask_event_submit(mask_data, mask_len);
		}

		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);
	keyboard_led_event_submit(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) {
		hid_protocol_event_submit(new_protocol);
	}
}

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

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

	if (iface) {
		iface->in_flight = false;
		submit_usb_tx_done((dev == g_usb_hid.boot.dev) ?
				   HID_TX_KIND_BOOT : HID_TX_KIND_REPORT,
				   true);
		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)) {
		const uint8_t *mask_data;
		size_t mask_len;

		if (try_extract_vendor_mask(dev, len, buf, &mask_data, &mask_len)) {
			LOG_INF("hid_stub_output_report vendor mask len=%u", mask_len);
			hid_vendor_mask_event_submit(mask_data, mask_len);
		}

		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);
	keyboard_led_event_submit(led_mask);
}

static void hid_iface_ready_cb(const struct device *dev, bool ready)
{
	struct usb_hid_iface *iface = usb_hid_iface_from_dev(dev);

	if (!iface) {
		return;
	}

	iface->iface_ready = ready;
	if (!ready) {
		iface->in_flight = false;
	}
}

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.policy.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) && usb_hid_stack_is_active()) {
			(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_state = USB_HID_STACK_STATE_ERROR;
		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_state != USB_HID_STACK_STATE_OFF) {
		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_state = USB_HID_STACK_STATE_READY;
	return 0;
}

static int usb_hid_set_enabled(bool enable)
{
	int err;

	if (usb_hid_stack_is_error()) {
		return -EIO;
	}

	if (g_usb_hid.stack_state == USB_HID_STACK_STATE_OFF) {
		err = usb_hid_stack_init();
		if (err) {
			return err;
		}
	}

	if (enable && usb_hid_stack_is_active()) {
		return 0;
	}

	if (!enable && (g_usb_hid.stack_state == USB_HID_STACK_STATE_READY)) {
		return 0;
	}

	if (enable) {
		err = usbd_enable(&new_kbd_usbd);
	} else {
		err = usbd_disable(&new_kbd_usbd);
		usb_hid_clear_runtime_iface_state();
	}

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

	g_usb_hid.stack_state = enable ? USB_HID_STACK_STATE_ACTIVE : USB_HID_STACK_STATE_READY;
	return 0;
}

static void refresh_usb_state_by_policy(void)
{
	/*
	 * 控制策略：
	 * - USB 模式 + 非休眠：启用 USB HID。
	 * - 其他情况：关闭 USB HID（不销毁初始化结果，后续可快速恢复）。
	 */
	bool should_enable = usb_hid_should_be_active();
	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_state = USB_HID_STACK_STATE_ERROR;
		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.policy.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.policy.pm_suspended) {
		/* 避免重复上报 STANDBY 导致 power_manager 在 SUSPENDING 期间反复迭代。 */
		return false;
	}

	g_usb_hid.policy.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.policy.pm_suspended) {
		return false;
	}

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

static bool handle_hid_tx_event(const struct hid_tx_event *event)
{
	if (!g_usb_hid.policy.usb_mode_selected || !usb_hid_stack_is_active()) {
		return false;
	}

	if (event->kind == HID_TX_KIND_BOOT) {
		const uint8_t *payload = hid_tx_event_get_data(event);
		size_t payload_len = hid_tx_event_get_size(event);
		int err;

		if (g_usb_hid.current_protocol != HID_PROTO_BOOT) {
			return false;
		}

		if (!g_usb_hid.boot.iface_ready || !g_usb_hid.boot.dev) {
			submit_usb_tx_done(HID_TX_KIND_BOOT, false);
			return false;
		}
		if (g_usb_hid.boot.in_flight) {
			LOG_WRN("Drop boot tx: previous report not sent");
			submit_usb_tx_done(HID_TX_KIND_BOOT, false);
			return false;
		}

		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);
			submit_usb_tx_done(HID_TX_KIND_BOOT, false);
		} else {
			g_usb_hid.boot.in_flight = true;
		}

		return false;
	}

	/*
	 * USB 侧仅在 active 条件满足时发送：
	 * - 当前 mode 为 USB；
	 * - USB HID 栈已启用且对应接口 ready。
	 */
	if (g_usb_hid.current_protocol != HID_PROTO_REPORT) {
		return false;
	}

	const uint8_t *data = hid_tx_event_get_data(event);
	size_t data_len = hid_tx_event_get_size(event);
	uint8_t report_id;

	if (data_len < 1U) {
		submit_usb_tx_done(HID_TX_KIND_REPORT, false);
		return false;
	}

	report_id = data[0];

	if (!g_usb_hid.nkro.iface_ready || !g_usb_hid.nkro.dev) {
		submit_usb_tx_done(HID_TX_KIND_REPORT, false);
		return false;
	}

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

	/* Report 协议下 dyndata 是 [report_id|payload]，可直接透传。 */
	int err = hid_device_submit_report(g_usb_hid.nkro.dev, data_len, data);
	if (err) {
		LOG_WRN("USB report send failed id=0x%02x err=%d", report_id, err);
		submit_usb_tx_done(HID_TX_KIND_REPORT, false);
	} 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_tx_event(aeh)) {
		return handle_hid_tx_event(cast_hid_tx_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_tx_event);