#include "stm32f7xx_hal.h"
#include "usbd_core.h"
#include "usbd_cdc.h"
#include "usbd_cdc_if.h"
#include "usbd_desc.h"
#include "icm42688.h"
#include "bmp280.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();
    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(1000);

    /* Init IMU */
    int imu_ret = icm42688_init();

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

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

    /* Send init status over USB */
    uint8_t tr[16]; int tlen;
    icm42688_get_trace(tr, &tlen);
    char buf[256];
    int len = snprintf(buf, sizeof(buf),
        "{\"init\":1,\"imu\":%d,\"who\":%d}\n",
        imu_ret, tr[0]);
    CDC_Transmit((uint8_t *)buf, len);

    icm42688_data_t 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) icm42688_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),
                    "{\"ax\":%d,\"ay\":%d,\"az\":%d,\"gx\":%d,\"gy\":%d,\"gz\":%d,"
                    "\"p\":%d,\"m\":%d,\"s\":%d}\n",
                    imu.ax, imu.ay, imu.az, imu.gx, imu.gy, imu.gz,
                    (int)(bal.pitch_deg * 10),   /* pitch x10 for precision without %f */
                    (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);
    }
}