#include "usbd_cdc_if.h"
#include "stm32f7xx_hal.h"

extern USBD_HandleTypeDef hUsbDevice;
volatile uint8_t cdc_streaming = 1;  /* auto-stream */
static volatile uint8_t cdc_port_open = 0;  /* set when host asserts DTR */
volatile uint8_t cdc_arm_request = 0;    /* set by A command */
volatile uint8_t cdc_disarm_request = 0; /* set by D command */
volatile uint8_t cdc_recal_request = 0;  /* set by G command — gyro recalibration */

/*
 * PID tuning command buffer.
 * CDC_Receive (USB IRQ) copies multi-char commands here.
 * Main loop polls cdc_cmd_ready, parses, and clears.
 * Commands: P<kp>  I<ki>  D<kd>  T<setpoint>  M<max_speed>  ?
 */
volatile uint8_t  cdc_cmd_ready = 0;
volatile char     cdc_cmd_buf[32];

/*
 * Jetson command buffer (bidirectional protocol).
 * 'H'\n        — heartbeat, ISR updates jetson_hb_tick only (no buf copy needed).
 * 'C'<s>,<t>\n — drive command: ISR copies to buf, main loop parses with sscanf.
 * jetson_hb_tick is also refreshed on every C command.
 */
volatile uint8_t  jetson_cmd_ready = 0;
volatile char     jetson_cmd_buf[32];
volatile uint32_t jetson_hb_tick = 0;  /* HAL_GetTick() of last H or C */

/*
 * USB TX/RX buffers grouped into a single 512-byte aligned struct so that
 * one MPU region (configured in usbd_conf.c) can mark them non-cacheable.
 * Size must be a power-of-2 >= total size for MPU RASR SIZE encoding.
 */
static struct {
    uint8_t tx[256];
    uint8_t rx[256];
} __attribute__((aligned(512))) usb_nc_buf;

#define UserTxBuffer usb_nc_buf.tx
#define UserRxBuffer usb_nc_buf.rx

/* Exported so usbd_conf.c USB_NC_MPU_Config() can set the region base */
void * const usb_nc_buf_base = &usb_nc_buf;

/*
 * Betaflight-proven DFU reboot:
 * 1. Write magic to RTC backup register (persists across soft reset)
 * 2. NVIC_SystemReset() — clean hardware reset
 * 3. Early startup checks magic, clears it, jumps to system bootloader
 *
 * The magic check happens in checkForBootloader() called from main.c
 * before any peripheral init.
 */
#define BOOTLOADER_MAGIC 0xDEADBEEF

static void request_bootloader(void) {
    /* Betaflight-proven: write magic, disable IRQs, reset.
     * checkForBootloader() runs on next boot before anything else. */
    __HAL_RCC_PWR_CLK_ENABLE();
    HAL_PWR_EnableBkUpAccess();
    __HAL_RCC_RTC_ENABLE();

    /* Write magic to RTC backup register 0 */
    RTC->BKP0R = BOOTLOADER_MAGIC;

    __disable_irq();
    NVIC_SystemReset();
}

/*
 * Call this VERY early in main(), before HAL_Init().
 * Checks RTC backup register for magic value left by request_bootloader().
 * If found: clear magic, jump to STM32F7 system bootloader at 0x1FF00000.
 */
void checkForBootloader(void) {
    /*
     * Betaflight-proven bootloader jump for STM32F7.
     * Called VERY early, before HAL_Init/caches/clocks.
     * At this point only RCC PWR is needed to read RTC backup regs.
     */

    /* Enable backup domain access to read RTC backup register */
    __HAL_RCC_PWR_CLK_ENABLE();
    HAL_PWR_EnableBkUpAccess();
    __HAL_RCC_RTC_ENABLE();

    uint32_t magic = RTC->BKP0R;

    if (magic != BOOTLOADER_MAGIC) {
        return; /* Normal boot */
    }

    /* Write POST marker (Betaflight does this so SystemInit can
     * do a second reset if needed — we just clear it) */
    RTC->BKP0R = 0;

    /* Jump to STM32F7 system bootloader at 0x1FF00000.
     * Exactly as Betaflight does it — no cache/VTOR/MEMRMP games needed
     * because we run before any of that is configured. */

    __HAL_RCC_SYSCFG_CLK_ENABLE();

    __set_MSP(*(uint32_t *)0x1FF00000);
    ((void (*)(void))(*(uint32_t *)0x1FF00004))();

    while (1);
}

static int8_t CDC_Init(void) {
    USBD_CDC_SetTxBuffer(&hUsbDevice, UserTxBuffer, 0);
    USBD_CDC_SetRxBuffer(&hUsbDevice, UserRxBuffer);
    USBD_CDC_ReceivePacket(&hUsbDevice);

    /* Reset TxState so CDC_Transmit works after host (re)connects.
     * Without this, if transmits happen before host opens port,
     * TxState stays BUSY forever since host never ACKs. */
    USBD_CDC_HandleTypeDef *hcdc = (USBD_CDC_HandleTypeDef *)hUsbDevice.pClassData;
    if (hcdc) hcdc->TxState = 0;

    return USBD_OK;
}

static int8_t CDC_DeInit(void) { return USBD_OK; }

static int8_t CDC_Control(uint8_t cmd, uint8_t *pbuf, uint16_t length) {
    (void)pbuf; (void)length;
    if (cmd == 0x22) {  /* CDC_SET_CONTROL_LINE_STATE — host opened port */
        cdc_port_open = 1;
        cdc_streaming = 1;
    }
    return USBD_OK;
}

static int8_t CDC_Receive(uint8_t *buf, uint32_t *len) {
    if (*len < 1) goto done;

    switch (buf[0]) {
        case 'S': cdc_streaming = !cdc_streaming; break;
        case 'A': cdc_arm_request = 1;            break;
        case 'D': cdc_disarm_request = 1;         break;
        case 'G': cdc_recal_request = 1;          break; /* gyro recalibration */
        case 'R': request_bootloader();            break; /* never returns */

        /*
         * PID tuning: P<kp>  I<ki>  D<kd>  T<setpoint>  M<max_speed>  ?
         * Copy to cmd buffer; main loop parses float (avoids sscanf in IRQ).
         */
        case 'P': case 'I': case 'K': case 'T': case 'M': case '?': {
            uint32_t copy_len = *len < 31 ? *len : 31;
            for (uint32_t i = 0; i < copy_len; i++) cdc_cmd_buf[i] = (char)buf[i];
            cdc_cmd_buf[copy_len] = '\0';
            cdc_cmd_ready = 1;
            break;
        }

        /* Jetson heartbeat — just refresh the tick, no buffer copy needed */
        case 'H':
            jetson_hb_tick = HAL_GetTick();
            break;

        /* Jetson drive command: C<speed>,<steer>\n
         * Copy to buffer; main loop parses ints (keeps sscanf out of ISR). */
        case 'C': {
            uint32_t copy_len = *len < 31 ? *len : 31;
            for (uint32_t i = 0; i < copy_len; i++) jetson_cmd_buf[i] = (char)buf[i];
            jetson_cmd_buf[copy_len] = '\0';
            jetson_hb_tick = HAL_GetTick(); /* C command also refreshes heartbeat */
            jetson_cmd_ready = 1;
            break;
        }

        default: break;
    }

done:
    USBD_CDC_SetRxBuffer(&hUsbDevice, UserRxBuffer);
    USBD_CDC_ReceivePacket(&hUsbDevice);
    return USBD_OK;
}

USBD_CDC_ItfTypeDef USBD_CDC_fops = { CDC_Init, CDC_DeInit, CDC_Control, CDC_Receive };

uint8_t CDC_Transmit(uint8_t *buf, uint16_t len) {
    USBD_CDC_HandleTypeDef *hcdc = (USBD_CDC_HandleTypeDef *)hUsbDevice.pClassData;
    if (hcdc == NULL) return USBD_FAIL;
    if (hcdc->TxState != 0) {
        /* If stuck busy (no host ACK), force reset after a while */
        static uint32_t busy_count = 0;
        if (++busy_count > 100) { hcdc->TxState = 0; busy_count = 0; }
        return USBD_BUSY;
    }
    /* Always copy into the static UserTxBuffer so the USB hardware reads
     * from a known fixed SRAM address — never from the caller's stack.
     * The USB TXFE IRQ fires asynchronously; a stack buffer could be
     * overwritten by the time the FIFO is loaded. */
    if (len > sizeof(UserTxBuffer)) len = sizeof(UserTxBuffer);
    memcpy(UserTxBuffer, buf, len);
    USBD_CDC_SetTxBuffer(&hUsbDevice, UserTxBuffer, len);
    return USBD_CDC_TransmitPacket(&hUsbDevice);
}