#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 "function_bitmap_state_event.h"
#include "function_bitmap_update_event.h"
#include "hid_led_event.h"
#include "proto_rx_event.h"
#include "proto_transport_state_event.h"
#include "proto_tx_event.h"
#include "protocol_module.h"
#include "theme_rgb_update_event.h"
#include "time_sync_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) | BIT(4))
#define PROTOCOL_MAX_MSG_LEN 128U

static bool initialized;
static bool running;

enum proto_session_state {
	PROTO_SESSION_DOWN = 0,
	PROTO_SESSION_WAIT_HELLO,
	PROTO_SESSION_ACTIVE,
};

static enum proto_session_state session_state[PROTO_TRANSPORT_COUNT];

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_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 encode_function_bitmap_state(const uint8_t *bitmap, uint8_t *payload,
					size_t payload_buf_size,
					size_t *payload_len)
{
	CdcPacketBody body = CdcPacketBody_init_zero;

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

	body.which_body = CdcPacketBody_function_key_event_tag;
	body.body.function_key_event.usage_bitmap.size = KEYBOARD_PROTOCOL_BITMAP_BYTES;
	memcpy(body.body.function_key_event.usage_bitmap.bytes, bitmap,
	       KEYBOARD_PROTOCOL_BITMAP_BYTES);

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

static int module_init(void)
{
	for (size_t i = 0; i < ARRAY_SIZE(session_state); i++) {
		session_state[i] = PROTO_SESSION_DOWN;
	}

	return 0;
}

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

	running = true;
	return 0;
}

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

	for (size_t i = 0; i < ARRAY_SIZE(session_state); i++) {
		session_state[i] = PROTO_SESSION_DOWN;
	}

	running = false;
}

int protocol_module_process_message(enum proto_transport transport,
				    const uint8_t *req_payload,
				    size_t req_payload_len,
				    uint8_t *rsp_payload,
				    size_t rsp_payload_buf_size,
				    size_t *rsp_payload_len)
{
	CdcPacketBody body;
	int err;

	if ((transport >= PROTO_TRANSPORT_COUNT) ||
	    (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 err;
	}

	switch (body.which_body) {
	case CdcPacketBody_hello_req_tag:
		if (session_state[transport] == PROTO_SESSION_DOWN) {
			return -EAGAIN;
		}

		LOG_INF("HelloReq transport:%u protocol_version:%u",
			transport, 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);
		}

		session_state[transport] = PROTO_SESSION_ACTIVE;
		return encode_hello_rsp(rsp_payload, rsp_payload_buf_size, rsp_payload_len);

	case CdcPacketBody_bitmap_tag:
		if (session_state[transport] != PROTO_SESSION_ACTIVE) {
			return encode_error(CdcPacketBody_bitmap_tag,
					    ErrorCode_ERROR_CODE_NOT_READY,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

		if (body.body.bitmap.usage_bitmap.size != KEYBOARD_PROTOCOL_BITMAP_BYTES) {
			return encode_error(CdcPacketBody_bitmap_tag,
					    ErrorCode_ERROR_CODE_INVALID_LENGTH,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

		err = submit_function_bitmap_update_event(
			body.body.bitmap.usage_bitmap.bytes);
		if (err) {
			return encode_error(CdcPacketBody_bitmap_tag,
					    ErrorCode_ERROR_CODE_INVALID_PARAM,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

		return encode_ack(CdcPacketBody_bitmap_tag, rsp_payload,
				  rsp_payload_buf_size, rsp_payload_len);

	case CdcPacketBody_time_sync_tag:
		if (session_state[transport] != PROTO_SESSION_ACTIVE) {
			return encode_error(CdcPacketBody_time_sync_tag,
					    ErrorCode_ERROR_CODE_NOT_READY,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

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

		submit_time_sync_event(body.body.time_sync.version,
				       body.body.time_sync.flags,
				       body.body.time_sync.timezone_min,
				       body.body.time_sync.utc_ms,
				       body.body.time_sync.accuracy_ms);
		return encode_ack(CdcPacketBody_time_sync_tag, rsp_payload,
				  rsp_payload_buf_size, rsp_payload_len);

	case CdcPacketBody_theme_rgb_tag:
		if (session_state[transport] != PROTO_SESSION_ACTIVE) {
			return encode_error(CdcPacketBody_theme_rgb_tag,
					    ErrorCode_ERROR_CODE_NOT_READY,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

		if ((body.body.theme_rgb.red > 255U) ||
		    (body.body.theme_rgb.green > 255U) ||
		    (body.body.theme_rgb.blue > 255U)) {
			return encode_error(CdcPacketBody_theme_rgb_tag,
					    ErrorCode_ERROR_CODE_INVALID_PARAM,
					    rsp_payload, rsp_payload_buf_size,
					    rsp_payload_len);
		}

		submit_theme_rgb_update_event((struct theme_rgb) {
			.r = (uint8_t)body.body.theme_rgb.red,
			.g = (uint8_t)body.body.theme_rgb.green,
			.b = (uint8_t)body.body.theme_rgb.blue,
		});
		return encode_ack(CdcPacketBody_theme_rgb_tag, rsp_payload,
				  rsp_payload_buf_size, rsp_payload_len);

	default:
		LOG_WRN("Unsupported protobuf body case %d", body.which_body);
		return -ENOTSUP;
	}
}

static bool handle_proto_rx_event(const struct proto_rx_event *event)
{
	uint8_t rsp_payload[PROTOCOL_MAX_MSG_LEN];
	size_t rsp_payload_len = 0U;
	int err;

	if (!running) {
		return false;
	}

	err = protocol_module_process_message(event->transport,
					      event->dyndata.data,
					      event->dyndata.size,
					      rsp_payload,
					      sizeof(rsp_payload),
					      &rsp_payload_len);
	if (err) {
		if (err != -ENOTSUP) {
			LOG_WRN("Protocol processing failed (%d)", err);
		}

		return false;
	}

	err = submit_proto_tx_event(event->transport, rsp_payload, rsp_payload_len);
	if (err) {
		LOG_WRN("Proto TX submit failed (%d)", err);
	}

	return false;
}

static bool handle_proto_transport_state_event(
	const struct proto_transport_state_event *event)
{
	if (event->transport >= PROTO_TRANSPORT_COUNT) {
		return false;
	}

	switch (event->state) {
	case PROTO_TRANSPORT_LINK_DOWN:
		session_state[event->transport] = PROTO_SESSION_DOWN;
		break;

	case PROTO_TRANSPORT_LINK_READY:
		session_state[event->transport] = PROTO_SESSION_WAIT_HELLO;
		break;

	default:
		return false;
	}

	return false;
}

static bool handle_function_bitmap_state_event(
	const struct function_bitmap_state_event *event)
{
	uint8_t payload[PROTOCOL_MAX_MSG_LEN];
	size_t payload_len;
	int err;

	if (!running) {
		return false;
	}

	for (enum proto_transport transport = 0; transport < PROTO_TRANSPORT_COUNT;
	     transport++) {
		if (session_state[transport] != PROTO_SESSION_ACTIVE) {
			continue;
		}

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

		err = submit_proto_tx_event(transport, 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[PROTOCOL_MAX_MSG_LEN];
	size_t payload_len;
	int err;
	enum proto_transport transport;

	if (!running) {
		return false;
	}

	transport = (event->transport == HID_TRANSPORT_USB) ?
			    PROTO_TRANSPORT_USB_CDC :
			    PROTO_TRANSPORT_BLE_NUS;

	if (session_state[transport] != PROTO_SESSION_ACTIVE) {
		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_proto_tx_event(transport, payload, payload_len);
	if (err) {
		LOG_WRN("LedState submit failed (%d)", err);
	}

	return false;
}

static bool app_event_handler(const struct app_event_header *aeh)
{
	if (is_proto_rx_event(aeh)) {
		return handle_proto_rx_event(cast_proto_rx_event(aeh));
	}

	if (is_proto_transport_state_event(aeh)) {
		return handle_proto_transport_state_event(
			cast_proto_transport_state_event(aeh));
	}

	if (is_function_bitmap_state_event(aeh)) {
		return handle_function_bitmap_state_event(
			cast_function_bitmap_state_event(aeh));
	}

	if (is_hid_led_event(aeh)) {
		return handle_hid_led_event(cast_hid_led_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;
		}

		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, function_bitmap_state_event);
APP_EVENT_SUBSCRIBE(MODULE, hid_led_event);
APP_EVENT_SUBSCRIBE(MODULE, module_state_event);
APP_EVENT_SUBSCRIBE(MODULE, proto_rx_event);
APP_EVENT_SUBSCRIBE(MODULE, proto_transport_state_event);
APP_EVENT_SUBSCRIBE_EARLY(MODULE, power_down_event);
APP_EVENT_SUBSCRIBE(MODULE, wake_up_event);