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

#include <app_event_manager.h>

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

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

#include <pb_decode.h>
#include <pb_encode.h>

#include <proto/device_comm.pb.h>

#include "theme_rgb_update_event.h"
#include "time_sync_event.h"
#include "cdc_proto_tx_event.h"
#include "function_bitmap_update_event.h"
#include "hid_led_event.h"
#include "key_function_event.h"
#include "protocol_module.h"
#include "usb_device_state_event.h"

LOG_MODULE_REGISTER(MODULE, LOG_LEVEL_INF);

#define PROTOCOL_VERSION 1U
#define PROTOCOL_VENDOR_ID 0x1915U
#define PROTOCOL_PRODUCT_ID 0x52F0U
#define PROTOCOL_FIRMWARE_MAJOR 0U
#define PROTOCOL_FIRMWARE_MINOR 0U
#define PROTOCOL_CAPABILITY_FLAGS (BIT(0) | BIT(1) | BIT(2) | BIT(3))

static bool initialized;
static bool running;
static bool keyboard_core_ready;
static bool usb_active;
static bool hello_done;

static bool type_matches_body(uint8_t type, const CdcPacketBody *body)
{
	switch (type) {
	case CDC_PROTO_TYPE_HELLO_REQ:
		return body->which_body == CdcPacketBody_hello_req_tag;
	case CDC_PROTO_TYPE_HELLO_RSP:
		return body->which_body == CdcPacketBody_hello_rsp_tag;
	case CDC_PROTO_TYPE_BITMAP:
		return body->which_body == CdcPacketBody_bitmap_tag;
	case CDC_PROTO_TYPE_FUNCTION_KEY_EVENT:
		return body->which_body == CdcPacketBody_function_key_event_tag;
	case CDC_PROTO_TYPE_LED_STATE:
		return body->which_body == CdcPacketBody_led_state_tag;
	case CDC_PROTO_TYPE_TIME_SYNC:
		return body->which_body == CdcPacketBody_time_sync_tag;
	case CDC_PROTO_TYPE_THEME_RGB:
		return body->which_body == CdcPacketBody_theme_rgb_tag;
	case CDC_PROTO_TYPE_ACK:
		return body->which_body == CdcPacketBody_ack_tag;
	case CDC_PROTO_TYPE_ERROR:
		return body->which_body == CdcPacketBody_error_tag;
	default:
		return false;
	}
}

static int decode_body(const uint8_t *payload, size_t payload_len,
		       CdcPacketBody *body)
{
	pb_istream_t stream;

	if ((payload == NULL) || (body == NULL)) {
		return -EINVAL;
	}

	*body = (CdcPacketBody)CdcPacketBody_init_zero;
	stream = pb_istream_from_buffer(payload, payload_len);

	if (!pb_decode(&stream, CdcPacketBody_fields, body)) {
		LOG_WRN("pb_decode failed: %s", PB_GET_ERROR(&stream));
		return -EBADMSG;
	}

	return 0;
}

static int encode_body(const CdcPacketBody *body, uint8_t *payload,
		       size_t payload_buf_size, size_t *payload_len)
{
	pb_ostream_t stream;

	if ((body == NULL) || (payload == NULL) || (payload_len == NULL)) {
		return -EINVAL;
	}

	stream = pb_ostream_from_buffer(payload, payload_buf_size);
	if (!pb_encode(&stream, CdcPacketBody_fields, body)) {
		LOG_WRN("pb_encode failed: %s", PB_GET_ERROR(&stream));
		return -EIO;
	}

	*payload_len = stream.bytes_written;
	return 0;
}

static int encode_hello_rsp(uint8_t *rsp_payload, size_t rsp_payload_buf_size,
			    size_t *rsp_payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

	body.which_body = CdcPacketBody_hello_rsp_tag;
	body.body.hello_rsp.protocol_version = PROTOCOL_VERSION;
	body.body.hello_rsp.vendor_id = PROTOCOL_VENDOR_ID;
	body.body.hello_rsp.product_id = PROTOCOL_PRODUCT_ID;
	body.body.hello_rsp.firmware_major = PROTOCOL_FIRMWARE_MAJOR;
	body.body.hello_rsp.firmware_minor = PROTOCOL_FIRMWARE_MINOR;
	body.body.hello_rsp.capability_flags = PROTOCOL_CAPABILITY_FLAGS;

	return encode_body(&body, rsp_payload, rsp_payload_buf_size, rsp_payload_len);
}

static int encode_ack(uint8_t acked_type, uint8_t *rsp_payload,
		      size_t rsp_payload_buf_size, size_t *rsp_payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

	body.which_body = CdcPacketBody_ack_tag;
	body.body.ack.acked_type = acked_type;

	return encode_body(&body, rsp_payload, rsp_payload_buf_size, rsp_payload_len);
}

static int encode_error(uint8_t error_type, ErrorCode error_code,
			uint8_t *rsp_payload, size_t rsp_payload_buf_size,
			size_t *rsp_payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

	body.which_body = CdcPacketBody_error_tag;
	body.body.error.error_type = error_type;
	body.body.error.error_code = error_code;

	return encode_body(&body, rsp_payload, rsp_payload_buf_size, rsp_payload_len);
}

static int encode_function_key_event(uint16_t usage, uint8_t action,
				     uint8_t *payload,
				     size_t payload_buf_size,
				     size_t *payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

	body.which_body = CdcPacketBody_function_key_event_tag;
	body.body.function_key_event.usage = usage;
	body.body.function_key_event.action = (KeyAction)action;

	return encode_body(&body, payload, payload_buf_size, payload_len);
}

static int encode_led_state(uint32_t led_mask, uint8_t *payload,
			    size_t payload_buf_size, size_t *payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

	body.which_body = CdcPacketBody_led_state_tag;
	body.body.led_state.led_mask = led_mask;

	return encode_body(&body, payload, payload_buf_size, payload_len);
}

static int submit_cdc_proto_tx_event(uint8_t type, const uint8_t *payload,
				     size_t payload_len)
{
	struct cdc_proto_tx_event *event;

	if ((payload == NULL) && (payload_len > 0U)) {
		return -EINVAL;
	}

	event = new_cdc_proto_tx_event(payload_len);
	event->type = type;
	if (payload_len > 0U) {
		memcpy(event->dyndata.data, payload, payload_len);
	}

	APP_EVENT_SUBMIT(event);
	return 0;
}

static int submit_function_bitmap_update_event(const Bitmap *bitmap)
{
	struct function_bitmap_update_event *event;

	if ((bitmap == NULL) ||
	    (bitmap->usage_bitmap.size != KEYBOARD_PROTOCOL_BITMAP_BYTES)) {
		return -EINVAL;
	}

	event = new_function_bitmap_update_event();
	memcpy(event->bitmap, bitmap->usage_bitmap.bytes,
	       KEYBOARD_PROTOCOL_BITMAP_BYTES);
	APP_EVENT_SUBMIT(event);

	return 0;
}

static int submit_theme_rgb_update_event(const ThemeRgb *theme_rgb)
{
	struct theme_rgb_update_event *event;

	if ((theme_rgb == NULL) ||
	    (theme_rgb->red > 255U) ||
	    (theme_rgb->green > 255U) ||
	    (theme_rgb->blue > 255U)) {
		return -EINVAL;
	}

	event = new_theme_rgb_update_event();
	event->theme.r = (uint8_t)theme_rgb->red;
	event->theme.g = (uint8_t)theme_rgb->green;
	event->theme.b = (uint8_t)theme_rgb->blue;
	APP_EVENT_SUBMIT(event);

	return 0;
}

static int submit_time_sync_event(const TimeSync *time_sync)
{
	struct time_sync_event *event;

	if (time_sync == NULL) {
		return -EINVAL;
	}

	event = new_time_sync_event();
	event->version = time_sync->version;
	event->flags = time_sync->flags;
	event->timezone_min = time_sync->timezone_min;
	event->utc_ms = time_sync->utc_ms;
	event->accuracy_ms = time_sync->accuracy_ms;
	APP_EVENT_SUBMIT(event);

	return 0;
}

static int encode_error_response(uint8_t req_type, ErrorCode error_code,
				 uint8_t *rsp_type, uint8_t *rsp_payload,
				 size_t rsp_payload_buf_size,
				 size_t *rsp_payload_len)
{
	int err;

	err = encode_error(req_type, error_code, rsp_payload,
			   rsp_payload_buf_size, rsp_payload_len);
	if (err) {
		return err;
	}

	*rsp_type = CDC_PROTO_TYPE_ERROR;
	return 0;
}

static int encode_ack_response(uint8_t acked_type, uint8_t *rsp_type,
			       uint8_t *rsp_payload,
			       size_t rsp_payload_buf_size,
			       size_t *rsp_payload_len)
{
	int err;

	err = encode_ack(acked_type, rsp_payload, rsp_payload_buf_size,
			 rsp_payload_len);
	if (err) {
		return err;
	}

	*rsp_type = CDC_PROTO_TYPE_ACK;
	return 0;
}

static int module_init(void)
{
	keyboard_core_ready = false;
	usb_active = false;
	hello_done = false;
	return 0;
}

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

	running = true;
	return 0;
}

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

	hello_done = false;
	running = false;
}

int protocol_module_process_cdc_packet(uint8_t req_type,
				       const uint8_t *req_payload,
				       size_t req_payload_len,
				       uint8_t *rsp_type,
				       uint8_t *rsp_payload,
				       size_t rsp_payload_buf_size,
				       size_t *rsp_payload_len)
{
	CdcPacketBody body;
	int err;

	if ((rsp_type == NULL) || (rsp_payload == NULL) || (rsp_payload_len == NULL)) {
		return -EINVAL;
	}

	if (!running) {
		return -EAGAIN;
	}

	err = decode_body(req_payload, req_payload_len, &body);
	if (err) {
		return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_LENGTH,
					     rsp_type, rsp_payload,
					     rsp_payload_buf_size, rsp_payload_len);
	}

	if (!type_matches_body(req_type, &body)) {
		LOG_WRN("CDC type/body mismatch type:0x%02x body_case:%d",
			req_type, body.which_body);
		return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_PARAM,
					     rsp_type, rsp_payload,
					     rsp_payload_buf_size, rsp_payload_len);
	}

	switch (req_type) {
	case CDC_PROTO_TYPE_HELLO_REQ:
		LOG_INF("HelloReq protocol_version:%u",
			body.body.hello_req.protocol_version);

		if (body.body.hello_req.protocol_version != PROTOCOL_VERSION) {
			LOG_WRN("Unexpected protocol version:%u",
				body.body.hello_req.protocol_version);
		}

		hello_done = true;
		err = encode_hello_rsp(rsp_payload, rsp_payload_buf_size, rsp_payload_len);
		if (err) {
			return err;
		}

		*rsp_type = CDC_PROTO_TYPE_HELLO_RSP;
		return 0;

	case CDC_PROTO_TYPE_BITMAP:
		if (!hello_done) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_NOT_READY,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		if (!keyboard_core_ready) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_NOT_READY,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		if (body.body.bitmap.usage_bitmap.size != KEYBOARD_PROTOCOL_BITMAP_BYTES) {
			LOG_WRN("Bitmap len:%u expected:%u",
				(unsigned int)body.body.bitmap.usage_bitmap.size,
				KEYBOARD_PROTOCOL_BITMAP_BYTES);
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_LENGTH,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		err = submit_function_bitmap_update_event(&body.body.bitmap);
		if (err) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_PARAM,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		return encode_ack_response(req_type, rsp_type, rsp_payload,
					   rsp_payload_buf_size, rsp_payload_len);

	case CDC_PROTO_TYPE_TIME_SYNC:
		if (!hello_done) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_NOT_READY,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		if (body.body.time_sync.version != 1U) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_PARAM,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		err = submit_time_sync_event(&body.body.time_sync);
		if (err) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_PARAM,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		return encode_ack_response(req_type, rsp_type, rsp_payload,
					   rsp_payload_buf_size, rsp_payload_len);

	case CDC_PROTO_TYPE_THEME_RGB:
		if (!hello_done) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_NOT_READY,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		err = submit_theme_rgb_update_event(&body.body.theme_rgb);
		if (err) {
			return encode_error_response(req_type, ErrorCode_ERROR_CODE_INVALID_PARAM,
						     rsp_type, rsp_payload,
						     rsp_payload_buf_size,
						     rsp_payload_len);
		}

		return encode_ack_response(req_type, rsp_type, rsp_payload,
					   rsp_payload_buf_size, rsp_payload_len);

	default:
		LOG_WRN("Unsupported CDC protocol type:0x%02x", req_type);
		return encode_error_response(req_type, ErrorCode_ERROR_CODE_UNKNOWN_TYPE,
					     rsp_type, rsp_payload,
					     rsp_payload_buf_size, rsp_payload_len);
	}
}

static bool handle_key_function_event(const struct key_function_event *event)
{
	uint8_t payload[64];
	size_t payload_len;
	int err;

	if (!running || !usb_active || !hello_done) {
		return false;
	}

	err = encode_function_key_event(event->usage, event->action, payload,
					sizeof(payload), &payload_len);
	if (err) {
		LOG_WRN("FunctionKeyEvent encode failed (%d)", err);
		return false;
	}

	err = submit_cdc_proto_tx_event(CDC_PROTO_TYPE_FUNCTION_KEY_EVENT,
					payload, payload_len);
	if (err) {
		LOG_WRN("FunctionKeyEvent submit failed (%d)", err);
	}

	return false;
}

static bool handle_hid_led_event(const struct hid_led_event *event)
{
	uint8_t payload[64];
	size_t payload_len;
	int err;

	if (!running || !usb_active || !hello_done) {
		return false;
	}

	err = encode_led_state(event->led_bm, payload, sizeof(payload), &payload_len);
	if (err) {
		LOG_WRN("LedState encode failed (%d)", err);
		return false;
	}

	err = submit_cdc_proto_tx_event(CDC_PROTO_TYPE_LED_STATE, payload, payload_len);
	if (err) {
		LOG_WRN("LedState submit failed (%d)", err);
	}

	return false;
}

static bool handle_usb_device_state_event(const struct usb_device_state_event *event)
{
	usb_active = (event->state == USB_DEVICE_STATE_ACTIVE);
	if (!usb_active) {
		hello_done = false;
	}

	return false;
}

static bool app_event_handler(const struct app_event_header *aeh)
{
	if (is_key_function_event(aeh)) {
		return handle_key_function_event(cast_key_function_event(aeh));
	}

	if (is_hid_led_event(aeh)) {
		return handle_hid_led_event(cast_hid_led_event(aeh));
	}

	if (is_usb_device_state_event(aeh)) {
		return handle_usb_device_state_event(cast_usb_device_state_event(aeh));
	}

	if (is_module_state_event(aeh)) {
		const struct module_state_event *event = cast_module_state_event(aeh);
		int err;

		if (check_state(event, MODULE_ID(main), MODULE_STATE_READY)) {
			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 (check_state(event, MODULE_ID(keyboard_core_module), MODULE_STATE_READY)) {
			keyboard_core_ready = true;
			return false;
		}

		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;
	}

	return false;
}

APP_EVENT_LISTENER(MODULE, app_event_handler);
APP_EVENT_SUBSCRIBE(MODULE, hid_led_event);
APP_EVENT_SUBSCRIBE(MODULE, key_function_event);
APP_EVENT_SUBSCRIBE(MODULE, module_state_event);
APP_EVENT_SUBSCRIBE(MODULE, usb_device_state_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, power_down_event);
APP_EVENT_SUBSCRIBE(MODULE, wake_up_event);