#include "stm32f7xx_hal.h"
#include "usbd_core.h"
#include "usbd_cdc.h"
#include "usbd_cdc_if.h"
#include "usbd_desc.h"
#include "mpu6000.h"
#include "balance.h"
#include "hoverboard.h"
#include "config.h"
#include "status.h"
#include <string.h>
#include <stdio.h>

extern volatile uint8_t cdc_streaming;      /* set by S command in CDC RX */
extern volatile uint8_t cdc_arm_request;    /* set by A command */
extern volatile uint8_t cdc_disarm_request; /* set by D command */

USBD_HandleTypeDef hUsbDevice;

static void SystemClock_Config(void) {
    RCC_OscInitTypeDef osc = {0};
    RCC_ClkInitTypeDef clk = {0};
    RCC_PeriphCLKInitTypeDef pclk = {0};
    __HAL_RCC_PWR_CLK_ENABLE();
    __HAL_PWR_VOLTAGESCALING_CONFIG(PWR_REGULATOR_VOLTAGE_SCALE1);
    osc.OscillatorType = RCC_OSCILLATORTYPE_HSE;
    osc.HSEState = RCC_HSE_ON;
    osc.PLL.PLLState = RCC_PLL_ON;
    osc.PLL.PLLSource = RCC_PLLSOURCE_HSE;
    osc.PLL.PLLM = 8; osc.PLL.PLLN = 432; osc.PLL.PLLP = 2; osc.PLL.PLLQ = 9;
    if (HAL_RCC_OscConfig(&osc) != HAL_OK) {
        osc.OscillatorType = RCC_OSCILLATORTYPE_HSI;
        osc.HSIState = RCC_HSI_ON; osc.HSICalibrationValue = RCC_HSICALIBRATION_DEFAULT;
        osc.HSEState = RCC_HSE_OFF; osc.PLL.PLLSource = RCC_PLLSOURCE_HSI; osc.PLL.PLLM = 16;
        if (HAL_RCC_OscConfig(&osc) != HAL_OK) while (1);
    }
    HAL_PWREx_EnableOverDrive();
    pclk.PeriphClockSelection = RCC_PERIPHCLK_CLK48 | RCC_PERIPHCLK_I2C1;
    pclk.Clk48ClockSelection = RCC_CLK48SOURCE_PLLSAIP;
    pclk.PLLSAI.PLLSAIN = 384; pclk.PLLSAI.PLLSAIQ = 7; pclk.PLLSAI.PLLSAIP = RCC_PLLSAIP_DIV8;
    pclk.I2c1ClockSelection = RCC_I2C1CLKSOURCE_PCLK1;
    if (HAL_RCCEx_PeriphCLKConfig(&pclk) != HAL_OK) while (1);
    clk.ClockType = RCC_CLOCKTYPE_SYSCLK|RCC_CLOCKTYPE_HCLK|RCC_CLOCKTYPE_PCLK1|RCC_CLOCKTYPE_PCLK2;
    clk.SYSCLKSource = RCC_SYSCLKSOURCE_PLLCLK;
    clk.AHBCLKDivider = RCC_SYSCLK_DIV1;
    clk.APB1CLKDivider = RCC_HCLK_DIV4; clk.APB2CLKDivider = RCC_HCLK_DIV2;
    HAL_RCC_ClockConfig(&clk, FLASH_LATENCY_7);
}

extern PCD_HandleTypeDef hpcd;
void OTG_FS_IRQHandler(void) { HAL_PCD_IRQHandler(&hpcd); }
void SysTick_Handler(void) { HAL_IncTick(); }

int main(void) {
    SCB_EnableICache();
    SCB_DisableDCache();  /* Must be off before USB starts — STM32F7 DCache/USB coherency */
    checkForBootloader(); /* Check RTC magic BEFORE HAL_Init — must be first thing */
    HAL_Init();
    SystemClock_Config();

    /* USB CDC */
    USBD_Init(&hUsbDevice, &SaltyLab_Desc, 0);
    USBD_RegisterClass(&hUsbDevice, &USBD_CDC);
    USBD_CDC_RegisterInterface(&hUsbDevice, &USBD_CDC_fops);
    USBD_Start(&hUsbDevice);

    status_init();
    HAL_Delay(3000);  /* Wait for USB host to enumerate */

    /* Init IMU (MPU6000 via SPI1 — mpu6000.c wraps icm42688 + complementary filter) */
    int imu_ret = mpu6000_init() ? 0 : -1;

    /* Init hoverboard ESC UART */
    hoverboard_init();

    /* Init balance controller */
    balance_t bal;
    balance_init(&bal);

    char buf[256];
    int len;

    IMUData imu;
    uint32_t send_tick = 0;
    uint32_t balance_tick = 0;
    uint32_t esc_tick = 0;
    const float dt = 1.0f / PID_LOOP_HZ;  /* 1ms at 1kHz */

    while (1) {
        if (imu_ret == 0) mpu6000_read(&imu);

        uint32_t now = HAL_GetTick();

        /* Handle arm/disarm requests from USB */
        if (cdc_arm_request) {
            cdc_arm_request = 0;
            balance_arm(&bal);
        }
        if (cdc_disarm_request) {
            cdc_disarm_request = 0;
            balance_disarm(&bal);
        }

        /* Balance PID at 1kHz */
        if (imu_ret == 0 && now - balance_tick >= 1) {
            balance_tick = now;
            balance_update(&bal, &imu, dt);
        }

        /* Send to hoverboard ESC at 50Hz (every 20ms)
         * Both wheels get same speed for balance, steer=0 for now */
        if (now - esc_tick >= 20) {
            esc_tick = now;
            if (bal.state == BALANCE_ARMED) {
                hoverboard_send(bal.motor_cmd, 0);
            } else {
                hoverboard_send(0, 0);  /* Always send to prevent ESC timeout */
            }
        }

        /* USB telemetry at 50Hz (only when streaming enabled via S command) */
        if (cdc_streaming && now - send_tick >= 20) {
            send_tick = now;
            if (imu_ret == 0) {
                len = snprintf(buf, sizeof(buf),
                    "{\"p\":%d,\"r\":%d,\"m\":%d,\"s\":%d}\n",
                    (int)(imu.pitch * 10),       /* fused pitch degrees x10 */
                    (int)(imu.pitch_rate * 10),  /* gyro pitch rate °/s x10 */
                    (int)bal.motor_cmd,
                    (int)bal.state);
            } else {
                len = snprintf(buf, sizeof(buf), "{\"err\":%d}\n", imu_ret);
            }
            CDC_Transmit((uint8_t *)buf, len);
        }

        status_update(now, (imu_ret == 0), (bal.state == BALANCE_ARMED));
        HAL_Delay(1);
    }
}