#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 */


static uint8_t UserRxBuffer[256];
static uint8_t UserTxBuffer[256];

/*
 * 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 && buf[0] == 'S') {
        cdc_streaming = !cdc_streaming;  /* Toggle streaming */
    }
    if (*len >= 1 && buf[0] == 'A') {
        cdc_arm_request = 1;  /* Arm balance — processed in main loop */
    }
    if (*len >= 1 && buf[0] == 'D') {
        cdc_disarm_request = 1;  /* Disarm — processed in main loop */
    }
    if (*len >= 1 && buf[0] == 'R') {
        request_bootloader(); /* Sets magic + resets — never returns */
    }
    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;
    }
    USBD_CDC_SetTxBuffer(&hUsbDevice, buf, len);
    return USBD_CDC_TransmitPacket(&hUsbDevice);
}