#include "usb.h"
#include "usb_conf.h"
#include "command_handler.h"

#include <zephyr/logging/log.h>
#include <zephyr/device.h>
#include <zephyr/drivers/uart.h>
#include <zephyr/kernel.h>
#include <zephyr/sys/ring_buffer.h>

// PICO-SDK
#include "pico/bootrom.h"

LOG_MODULE_REGISTER(usb, LOG_LEVEL_INF);

// DEVICE
const struct device *const uart_dev = DEVICE_DT_GET_ONE(zephyr_cdc_acm_uart);
static struct usbd_context *usb_context;

// RX THREAD
static struct k_thread usb_rx_thread_data;
static k_tid_t usb_rx_thread_id = NULL;
#define USB_RX_THREAD_STACK_SIZE 2048
K_THREAD_STACK_DEFINE(usb_rx_thread_stack, USB_RX_THREAD_STACK_SIZE);

// RX BUFFER
#define RING_BUF_SIZE 255
static uint8_t ring_buffer[RING_BUF_SIZE];
static struct ring_buf ringbuf;
struct k_sem rx_semaphore;

// TX THREAD
static struct k_thread usb_tx_thread_data;
static k_tid_t usb_tx_thread_id = NULL;
#define USB_TX_THREAD_STACK_SIZE 512
K_THREAD_STACK_DEFINE(usb_tx_thread_stack, USB_TX_THREAD_STACK_SIZE);

// TX BUFFER (add ack and nack at the and as static)
#define TX_BUFFER_SIZE 10
struct command_message_t usb_tx_buffer[TX_BUFFER_SIZE + 2];
struct command_message_t *usb_tx_buf_ptr;
char usb_tx_ptr_msgq_buffer[(TX_BUFFER_SIZE + 2) * sizeof(struct command_message_t *)];
struct k_msgq usb_tx_ptr_msgq;

// ACK / NACK messages
#define RETURN_ACK true


static void interrupt_handler(const struct device *dev, void *user_data) {
	ARG_UNUSED(user_data);

	while (true) {
		uart_irq_update(dev);

		if (uart_irq_is_pending(dev) <= 0) {
			break;
		}

		if (uart_irq_rx_ready(dev)) {
			int recv_len, rb_len;
			uint8_t buffer[64];
			size_t len = MIN(ring_buf_space_get(&ringbuf), sizeof(buffer));

			if (len == 0) {
				// ring buffer full, drops package(s)
			    uart_irq_rx_disable(dev);
				k_sem_give(&rx_semaphore);
				break;
			}

			recv_len = uart_fifo_read(dev, buffer, len);
			if (recv_len < 0) {
				LOG_ERR("Failed to read UART FIFO");
				recv_len = 0;
			};

			rb_len = ring_buf_put(&ringbuf, buffer, recv_len);
			if (rb_len < recv_len) {
				LOG_ERR("Drop %u bytes", recv_len - rb_len);
			}

			k_sem_give(&rx_semaphore);
		}
	}
}

static void usb_rx_thread(void *p1, void *p2, void *p3) {
	ARG_UNUSED(p1);
	ARG_UNUSED(p2);
	ARG_UNUSED(p3);
	struct command_message_t msg;
	command_message_init(&msg);
	
	LOG_INF("USB command processing thread started");
	
	while (1) {
		k_sem_take(&rx_semaphore, K_FOREVER);

		int len;

		// While ring buffer has data
		do {
			uint8_t buf_prefix;
			len = ring_buf_get(&ringbuf, &buf_prefix, 1);

			if (len && (buf_prefix == COMMAND_PREFIX)) {
				uint8_t buf_header[4];
				len = ring_buf_get(&ringbuf, buf_header, 4);

				if ((len == 4) && (buf_header[1] == COMMAND_ID) && (buf_header[0] <= COMMAND_DATA_SIZE)) {
					msg.length = buf_header[0];
					msg.command = buf_header[2];
					msg.crc = buf_header[3];

					if (msg.length) {
						len = ring_buf_get(&ringbuf, msg.data, msg.length);
					}

					uint8_t calculated_crc = command_calculate_crc(&msg);
					if (calculated_crc != msg.crc) {
						if (RETURN_ACK) {
							// Send NACK
							usb_send_command(&usb_tx_buffer[TX_BUFFER_SIZE + 1]);
						}
						continue;
					}

					int ret = command_handler(&msg);
					if (ret == 0) {
						if (RETURN_ACK) {
							// Send ACK
							usb_send_command(&usb_tx_buffer[TX_BUFFER_SIZE]);
						}
					}
					else {
						if (RETURN_ACK) {
							// Send NACK
							usb_send_command(&usb_tx_buffer[TX_BUFFER_SIZE + 1]);
						}
					}

				}
				else {
					// Command_id did not match, ignore
					continue;
				}
			}
			else {
				// Prefix did not match, ignore
				continue;
			}
		} while (len > 0);

		uart_irq_rx_enable(uart_dev);
	}
	
	LOG_INF("USB command processing thread exiting");
}

static void usb_msg_cb(struct usbd_context *const ctx, const struct usbd_msg *msg) {
	if (msg->type == USBD_MSG_CDC_ACM_LINE_CODING) {
		// Jump to BOOTSEL when baudrate changes to 1200
		uint32_t baudrate;

		if (uart_line_ctrl_get(msg->dev, UART_LINE_CTRL_BAUD_RATE, &baudrate) == 0) {
			LOG_INF("Baudrate %u", baudrate);

			if (baudrate == 1200) {
				LOG_INF("Entering BOOTSEL...");
				reset_usb_boot(0, 0);
			}
		}
	}
}

static void usb_tx_thread(void *p1, void *p2, void *p3) {
	struct command_message_t *data;

    while (1) {
        k_msgq_get(&usb_tx_ptr_msgq, &data, K_FOREVER);

		if (!device_is_ready(uart_dev)) {
			// FIXME: Don't drop packages
			continue;
		}

		// Message size: prefix + length + id + command + crc + data
		size_t msg_size = COMMAND_HEADER_SIZE + data->length;
		uint8_t *msg_bytes = (uint8_t *)data;

		/* uart_poll_out blocks until sent, ensuring data integrity */
		for (size_t i = 0; i < msg_size; i++) {
			uart_poll_out(uart_dev, msg_bytes[i]);
		}
    }
}

int usb_init() {
	ring_buf_init(&ringbuf, sizeof(ring_buffer), ring_buffer);
	k_sem_init(&rx_semaphore, 0, 1);
	usb_tx_buf_ptr = usb_tx_buffer;
	k_msgq_init(&usb_tx_ptr_msgq, usb_tx_ptr_msgq_buffer, sizeof(struct command_message_t *), TX_BUFFER_SIZE);

	command_create_ack(&usb_tx_buffer[TX_BUFFER_SIZE]);
	command_create_nack(&usb_tx_buffer[TX_BUFFER_SIZE + 1]);

	int ret;

	if (!device_is_ready(uart_dev)) {
		LOG_ERR("CDC ACM device not ready");
		return -ENODEV;
	}

	usb_context = usb_device_init(usb_msg_cb);
	if (usb_context == NULL) {
		LOG_ERR("Failed to initialize USB device");
		return -ENODEV;
	}

	if (!usbd_can_detect_vbus(usb_context)) {
		ret = usbd_enable(usb_context);
		if (ret) {
			LOG_ERR("Failed to enable device support");
			return ret;
		}
	}
	k_msleep(100);

	uart_irq_callback_set(uart_dev, interrupt_handler);
	uart_irq_rx_enable(uart_dev);

	// RX THREAD
	usb_rx_thread_id = k_thread_create(
		&usb_rx_thread_data,
		usb_rx_thread_stack,
		K_THREAD_STACK_SIZEOF(usb_rx_thread_stack),
		usb_rx_thread,
		NULL, NULL, NULL,
		5,
		0,
		K_NO_WAIT
	);

	if (usb_rx_thread_id == NULL) {
		LOG_ERR("Failed to create USB RX thread");
		return -ENOMEM;
	}

	// TX THREAD
	usb_tx_thread_id = k_thread_create(
		&usb_tx_thread_data,
		usb_tx_thread_stack,
		K_THREAD_STACK_SIZEOF(usb_tx_thread_stack),
		usb_tx_thread,
		NULL, NULL, NULL,
		5,
		0,
		K_NO_WAIT
	);

	if (usb_tx_thread_id == NULL) {
		LOG_ERR("Failed to create USB TX thread");
		return -ENOMEM;
	}

	return ret;
}

struct command_message_t* usb_get_next_tx_buf() {
	struct command_message_t *buf = usb_tx_buf_ptr;

	// Increment the buffer pointer
	usb_tx_buf_ptr++;

	if (usb_tx_buf_ptr > &usb_tx_buffer[TX_BUFFER_SIZE-1]) {
		usb_tx_buf_ptr = usb_tx_buffer;
	}

	return buf;
}

int usb_send_command(struct command_message_t *msg) {
	
	k_msgq_put(&usb_tx_ptr_msgq, &msg, K_NO_WAIT);

	return 0;
}