#include <stdio.h>
#include <string.h>
#include "pico/stdlib.h"
#include "pico/unique_id.h"
#include "hardware/uart.h"
#include "hardware/irq.h"
#include "hardware/gpio.h"
#include "hardware/clocks.h"
#include "tusb.h"

//--------------------------------------------------------------------+
// UART Configuration
//--------------------------------------------------------------------+
#define UART_ID          uart0
#define UART_IRQ         UART0_IRQ
#define UART_TX_PIN      0
#define UART_RX_PIN      1
#define DEFAULT_BAUD     115200

// Activity / Heartbeat LED Pin
#define LED_PIN_WS       16
#define LED_PIN_PICO     25

// UART RX Ring Buffer for capture
#define BUFFER_SIZE      4096

static volatile struct {
    uint8_t  buf[BUFFER_SIZE];
    volatile uint32_t head;
    volatile uint32_t tail;
} uart_rx_ring;

static inline bool ring_is_empty(void) {
    return uart_rx_ring.head == uart_rx_ring.tail;
}

static inline void ring_push(uint8_t c) {
    uint32_t next = (uart_rx_ring.head + 1) % BUFFER_SIZE;
    if (next != uart_rx_ring.tail) {
        uart_rx_ring.buf[uart_rx_ring.head] = c;
        uart_rx_ring.head = next;
    }
}

static inline uint8_t ring_pop(void) {
    uint8_t c = uart_rx_ring.buf[uart_rx_ring.tail];
    uart_rx_ring.tail = (uart_rx_ring.tail + 1) % BUFFER_SIZE;
    return c;
}

//--------------------------------------------------------------------+
// UART Interrupt Handler
//--------------------------------------------------------------------+
static void on_uart_rx(void) {
    while (uart_is_readable(UART_ID)) {
        uint8_t ch = (uint8_t)uart_getc(UART_ID);
        ring_push(ch);
    }
}

//--------------------------------------------------------------------+
// Hardware UART Setup
//--------------------------------------------------------------------+
static void setup_uart(uint32_t baudrate) {
    uart_init(UART_ID, baudrate);
    gpio_set_function(UART_TX_PIN, GPIO_FUNC_UART);
    gpio_set_function(UART_RX_PIN, GPIO_FUNC_UART);
    gpio_pull_up(UART_RX_PIN);

    uart_set_hw_flow(UART_ID, false, false);
    uart_set_format(UART_ID, 8, 1, UART_PARITY_NONE);
    uart_set_fifo_enabled(UART_ID, true);

    // Set up and enable the interrupt handlers
    irq_set_exclusive_handler(UART_IRQ, on_uart_rx);
    irq_set_enabled(UART_IRQ, true);

    // Enable UART RX interrupt
    uart_set_irq_enables(UART_ID, true, false);
}

//--------------------------------------------------------------------+
// USB Descriptors
//--------------------------------------------------------------------+
#define USB_VID   0x2E8A // Raspberry Pi
#define USB_PID   0x000A // Raspberry Pi Pico CDC
#define USB_BCD   0x0200

tusb_desc_device_t const desc_device = {
    .bLength            = sizeof(tusb_desc_device_t),
    .bDescriptorType    = TUSB_DESC_DEVICE,
    .bcdUSB             = USB_BCD,
    .bDeviceClass       = TUSB_CLASS_MISC,
    .bDeviceSubClass    = MISC_SUBCLASS_COMMON,
    .bDeviceProtocol    = MISC_PROTOCOL_IAD,
    .bMaxPacketSize0    = CFG_TUD_ENDPOINT0_SIZE,
    .idVendor           = USB_VID,
    .idProduct          = USB_PID,
    .bcdDevice          = 0x0100,
    .iManufacturer      = 0x01,
    .iProduct           = 0x02,
    .iSerialNumber      = 0x03,
    .bNumConfigurations = 0x01
};

extern "C" uint8_t const *tud_descriptor_device_cb(void) {
    return (uint8_t const *) &desc_device;
}

enum {
    ITF_NUM_CDC = 0,
    ITF_NUM_CDC_DATA,
    ITF_NUM_TOTAL
};

#define EPNUM_CDC_NOTIF   0x81
#define EPNUM_CDC_OUT     0x02
#define EPNUM_CDC_IN      0x82

#define CONFIG_TOTAL_LEN  (TUD_CONFIG_DESC_LEN + TUD_CDC_DESC_LEN)

uint8_t const desc_configuration[] = {
    TUD_CONFIG_DESCRIPTOR(1, ITF_NUM_TOTAL, 0, CONFIG_TOTAL_LEN, 0x00, 100),
    TUD_CDC_DESCRIPTOR(ITF_NUM_CDC, 4, EPNUM_CDC_NOTIF, 8, EPNUM_CDC_OUT, EPNUM_CDC_IN, 64),
};

extern "C" uint8_t const *tud_descriptor_configuration_cb(uint8_t index) {
    (void) index;
    return desc_configuration;
}

enum {
    STRID_LANGID = 0,
    STRID_MANUFACTURER,
    STRID_PRODUCT,
    STRID_SERIAL,
    STRID_CDC_INTERFACE,
};

static char const *string_desc_arr[] = {
    (const char[]) { 0x09, 0x04 },
    "Waveshare",
    "RP2350 USB-UART Adapter",
    NULL,
    "RP2350 CDC UART",
};

static uint16_t _desc_str[64];

extern "C" uint16_t const *tud_descriptor_string_cb(uint8_t index, uint16_t langid) {
    (void) langid;
    size_t chr_count = 0;

    switch (index) {
        case STRID_LANGID:
            memcpy(&_desc_str[1], string_desc_arr[0], 2);
            chr_count = 1;
            break;

        case STRID_SERIAL: {
            pico_unique_board_id_t id;
            pico_get_unique_board_id(&id);
            char serial_str[2 * PICO_UNIQUE_BOARD_ID_SIZE_BYTES + 1];
            pico_get_unique_board_id_string(serial_str, sizeof(serial_str));
            chr_count = strlen(serial_str);
            if (chr_count > 32) chr_count = 32;
            for (size_t i = 0; i < chr_count; i++) {
                _desc_str[1 + i] = serial_str[i];
            }
            break;
        }

        default: {
            if (!(index < sizeof(string_desc_arr) / sizeof(string_desc_arr[0]))) {
                return NULL;
            }
            const char *str = string_desc_arr[index];
            if (!str) return NULL;

            chr_count = strlen(str);
            size_t const max_count = sizeof(_desc_str) / sizeof(_desc_str[0]) - 1;
            if (chr_count > max_count) chr_count = max_count;

            for (size_t i = 0; i < chr_count; i++) {
                _desc_str[1 + i] = str[i];
            }
            break;
        }
    }

    _desc_str[0] = (uint16_t) ((TUSB_DESC_STRING << 8) | (2 * chr_count + 2));
    return _desc_str;
}

//--------------------------------------------------------------------+
// TinyUSB CDC Callbacks
//--------------------------------------------------------------------+

extern "C" void tud_cdc_line_state_cb(uint8_t itf, bool dtr, bool rts) {
    (void) itf;
    (void) dtr;
    (void) rts;
}

extern "C" void tud_cdc_line_coding_cb(uint8_t itf, cdc_line_coding_t const* p_line_coding) {
    (void) itf;

    if (p_line_coding->bit_rate > 0) {
        uart_set_baudrate(UART_ID, p_line_coding->bit_rate);
    }

    uint data_bits = (p_line_coding->data_bits >= 5 && p_line_coding->data_bits <= 8)
                     ? p_line_coding->data_bits : 8;

    uint stop_bits = (p_line_coding->stop_bits == 2) ? 2 : 1;

    uart_parity_t parity = UART_PARITY_NONE;
    if (p_line_coding->parity == 1) {
        parity = UART_PARITY_ODD;
    } else if (p_line_coding->parity == 2) {
        parity = UART_PARITY_EVEN;
    }

    uart_set_format(UART_ID, data_bits, stop_bits, parity);
}

//--------------------------------------------------------------------+
// Main Bridge Task
//--------------------------------------------------------------------+
static void cdc_uart_bridge_task(void) {
    // 1. Drain any available hardware UART RX bytes directly into ring buffer
    while (uart_is_readable(UART_ID)) {
        ring_push((uint8_t)uart_getc(UART_ID));
    }

    // 2. Forward UART RX -> USB CDC TX
    if (!ring_is_empty()) {
        uint32_t available = tud_cdc_write_available();
        if (available > 0) {
            uint8_t tx_buf[64];
            uint32_t count = 0;
            while (count < sizeof(tx_buf) && count < available && !ring_is_empty()) {
                tx_buf[count++] = ring_pop();
            }
            if (count > 0) {
                tud_cdc_write(tx_buf, count);
                tud_cdc_write_flush();
            }
        }
    }

    // 3. Forward USB CDC RX -> UART TX
    if (tud_cdc_available()) {
        uint8_t usb_rx_buf[64];
        uint32_t count = tud_cdc_read(usb_rx_buf, sizeof(usb_rx_buf));
        if (count > 0) {
            uart_write_blocking(UART_ID, usb_rx_buf, count);
        }
    }
}

//--------------------------------------------------------------------+
// Main Entry
//--------------------------------------------------------------------+
int main(void) {
    stdio_init_all();

    // Initialize Activity LED
    gpio_init(LED_PIN_WS);
    gpio_set_dir(LED_PIN_WS, GPIO_OUT);
    gpio_put(LED_PIN_WS, 0);

    #ifdef LED_PIN_PICO
    gpio_init(LED_PIN_PICO);
    gpio_set_dir(LED_PIN_PICO, GPIO_OUT);
    gpio_put(LED_PIN_PICO, 0);
    #endif

    // Initialize UART0 on GP0 (TX) and GP1 (RX)
    setup_uart(DEFAULT_BAUD);

    // Initialize TinyUSB Device stack
    tusb_init();

    uint32_t last_blink = 0;
    bool led_state = false;

    while (true) {
        // Process TinyUSB device events
        tud_task();

        // Process bidirectional data transfer
        cdc_uart_bridge_task();

        // Heartbeat blink (toggle every 250ms)
        uint32_t now = to_ms_since_boot(get_absolute_time());
        if (now - last_blink >= 250) {
            last_blink = now;
            led_state = !led_state;
            gpio_put(LED_PIN_WS, led_state);
            #ifdef LED_PIN_PICO
            gpio_put(LED_PIN_PICO, led_state);
            #endif
        }
    }

    return 0;
}