#include "trikke_usb_transport.h"

#include "driver/usb_serial_jtag.h"
#include "driver/usb_serial_jtag_vfs.h"
#include "freertos/FreeRTOS.h"

#define TRIKKE_USB_TX_BUFFER_SIZE 512
#define TRIKKE_USB_RX_BUFFER_SIZE 128
#define TRIKKE_USB_OPERATION_TIMEOUT_MS 50

static trikke_transport_status_t usb_begin_packet(
    void *context,
    const uint8_t *packet,
    size_t packet_size)
{
    const trikke_usb_transport_t *usb = context;
    if (usb == NULL || !usb->initialized || packet == NULL || packet_size == 0 ||
        packet_size > TRIKKE_USB_TX_BUFFER_SIZE) {
        return TRIKKE_TRANSPORT_FATAL;
    }

    const int written = usb_serial_jtag_write_bytes(
        packet, packet_size,
        pdMS_TO_TICKS(TRIKKE_USB_OPERATION_TIMEOUT_MS));
    if (written == (int)packet_size) {
        return TRIKKE_TRANSPORT_PENDING;
    }
    if (written == 0) {
        return TRIKKE_TRANSPORT_RETRY;
    }
    return TRIKKE_TRANSPORT_FATAL;
}

static trikke_transport_status_t usb_poll_packet(void *context)
{
    const trikke_usb_transport_t *usb = context;
    if (usb == NULL || !usb->initialized) {
        return TRIKKE_TRANSPORT_FATAL;
    }

    const esp_err_t err = usb_serial_jtag_wait_tx_done(
        pdMS_TO_TICKS(TRIKKE_USB_OPERATION_TIMEOUT_MS));
    if (err == ESP_OK) {
        return TRIKKE_TRANSPORT_COMPLETE;
    }
    if (err == ESP_ERR_TIMEOUT) {
        return TRIKKE_TRANSPORT_PENDING;
    }
    return TRIKKE_TRANSPORT_FATAL;
}

esp_err_t trikke_usb_transport_init(
    trikke_usb_transport_t *usb,
    trikke_transport_t *transport)
{
    if (usb == NULL || transport == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    if (usb->initialized) {
        return ESP_ERR_INVALID_STATE;
    }

    usb_serial_jtag_driver_config_t config = {
        .tx_buffer_size = TRIKKE_USB_TX_BUFFER_SIZE,
        .rx_buffer_size = TRIKKE_USB_RX_BUFFER_SIZE,
    };
    const esp_err_t err = usb_serial_jtag_driver_install(&config);
    if (err != ESP_OK) {
        return err;
    }

    usb->initialized = true;
    transport->context = usb;
    transport->begin = usb_begin_packet;
    transport->poll = usb_poll_packet;

    // Route any unexpected VFS output through the installed driver as well, so
    // it cannot race the driver's ISR by touching the hardware FIFO directly.
    usb_serial_jtag_vfs_use_driver();
    return ESP_OK;
}

void trikke_usb_transport_deinit(trikke_usb_transport_t *usb)
{
    if (usb == NULL || !usb->initialized) {
        return;
    }

    usb_serial_jtag_vfs_use_nonblocking();
    usb_serial_jtag_driver_uninstall();
    usb->initialized = false;
}