#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 "motor_driver.h"
#include "mode_manager.h"
#include "config.h"
#include "status.h"
#include "safety.h"
#include "crsf.h"
#include "i2c1.h"
#include "bmp280.h"
#include "mag.h"
#include "jetson_cmd.h"
#include "battery.h"
#include <math.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 */
extern volatile uint8_t  cdc_recal_request;  /* set by G command */
extern volatile uint32_t cdc_rx_count;       /* total CDC packets received */
extern volatile uint8_t  cdc_estop_request;
extern volatile uint8_t  cdc_estop_clear_request;

/*
 * Apply a PID tuning command string from the USB terminal.
 * Format: P<kp>  I<ki>  D<kd>  T<setpoint_deg>  M<max_speed>  ?
 * Returns a short ack string written into `reply` (max reply_sz bytes).
 */
static void apply_pid_cmd(balance_t *b, const char *cmd, char *reply, int reply_sz) {
    float val = 0.0f;
    /* Simple ASCII float parse — avoids sscanf in hot path */
    if (cmd[0] == '?') {
        snprintf(reply, reply_sz,
            "{\"kp\":%d,\"ki\":%d,\"kd\":%d,\"sp\":%d,\"ms\":%d}\n",
            (int)(b->kp * 1000), (int)(b->ki * 1000), (int)(b->kd * 1000),
            (int)(b->setpoint * 10), (int)b->max_speed);
        return;
    }
    /* Parse float after the command letter */
    if (sscanf(cmd + 1, "%f", &val) != 1) {
        snprintf(reply, reply_sz, "{\"err\":\"bad_fmt\"}\n");
        return;
    }
    switch (cmd[0]) {
        case 'P': if (val >= 0.0f && val <= 500.0f) b->kp = val;        break;
        case 'I': if (val >= 0.0f && val <= 50.0f)  b->ki = val;        break;
        case 'K': /* alias for D */ /* fall through */
        case 'D': if (val >= 0.0f && val <= 50.0f)  b->kd = val;        break;
        case 'T': if (val >= -20.0f && val <= 20.0f) b->setpoint = val; break;
        case 'M': if (val >= 0.0f && val <= 1000.0f) b->max_speed = (int16_t)val; break;
    }
    snprintf(reply, reply_sz,
        "{\"kp\":%d,\"ki\":%d,\"kd\":%d,\"sp\":%d,\"ms\":%d}\n",
        (int)(b->kp * 1000), (int)(b->ki * 1000), (int)(b->kd * 1000),
        (int)(b->setpoint * 10), (int)b->max_speed);
}

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();
    /* DCache stays ON — MPU Region 0 in usbd_conf.c marks USB buffers non-cacheable. */
    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;

    /*
     * Gyro bias calibration — hold board still for ~1s.
     * Must run before safety_init() (IWDG not yet started).
     * Blocks arming and streaming until complete.
     */
    if (imu_ret == 0) mpu6000_calibrate();

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

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

    /* Init motor driver */
    motor_driver_t motors;
    motor_driver_init(&motors);

    /* Init CRSF/ELRS receiver on UART4 (PA0/PA1) with DMA */
    crsf_init();

    /* Init mode manager (RC/autonomous blend; CH6 mode switch) */
    mode_manager_t mode;
    mode_manager_init(&mode);

    /* Init battery ADC (PC1/ADC3 — Vbat divider 11:1) for CRSF telemetry */
    battery_init();

    /* Probe I2C1 for optional sensors — skip gracefully if not found */
    int baro_chip = 0;  /* chip_id: 0x58=BMP280, 0x60=BME280, 0=absent */
    mag_type_t mag_type = MAG_NONE;
    if (i2c1_init() == 0) {
        int chip = bmp280_init();
        baro_chip = (chip > 0) ? chip : 0;
        mag_type = mag_init();
    }

    /*
     * IWDG starts AFTER all peripheral inits — avoids reset during mpu6000_init()
     * which takes ~510ms (well above the 50ms WATCHDOG_TIMEOUT_MS).
     * Once started, safety_refresh() MUST be called every WATCHDOG_TIMEOUT_MS
     * or MCU resets. IWDG cannot be stopped once started (hardware enforced).
     */
    safety_init();

    char buf[320];
    int len;
    uint32_t tx_count = 0;   /* telemetry frames sent to host */

    IMUData imu;
    uint32_t send_tick = 0;
    uint32_t balance_tick = 0;
    uint32_t esc_tick = 0;
    uint32_t crsf_telem_tick = 0;  /* CRSF uplink telemetry TX timer */
    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();

        /* Feed hardware watchdog — must happen every WATCHDOG_TIMEOUT_MS */
        safety_refresh();

        /* Mode manager: update RC liveness, CH6 mode selection, blend ramp */
        mode_manager_update(&mode, now);

        if (cdc_estop_request) {
            EstopSource src = (cdc_estop_request == 1) ? ESTOP_REMOTE : ESTOP_CELLULAR_TIMEOUT;
            cdc_estop_request = 0;
            safety_remote_estop(src);
            safety_arm_cancel();
            balance_disarm(&bal);
            motor_driver_estop(&motors);
        }
        if (cdc_estop_clear_request) {
            cdc_estop_clear_request = 0;
            if (safety_remote_estop_active() && bal.state == BALANCE_DISARMED)
                safety_remote_estop_clear();
        }

        /* RC CH5 kill switch: disarm immediately if RC is alive and CH5 off.
         * Applies regardless of active mode (CH5 always has kill authority). */
        if (mode.rc_alive && !crsf_state.armed && bal.state == BALANCE_ARMED) {
            safety_arm_cancel();
            balance_disarm(&bal);
        }
        /* RC failsafe: disarm if signal lost AFTER RC was previously alive.
         * Prevents auto-disarm in USB-only mode (crsf_state.last_rx_ms == 0). */
        if (bal.state == BALANCE_ARMED &&
                crsf_state.last_rx_ms != 0 &&
                (now - crsf_state.last_rx_ms) > CRSF_FAILSAFE_MS) {
            safety_arm_cancel();
            balance_disarm(&bal);
        }

        /* Tilt fault buzzer alert (one-shot on fault edge) */
        safety_alert_tilt_fault(bal.state == BALANCE_TILT_FAULT);

        /* RC arm/disarm via CH5 switch (CRSF) — edge detect, same hold interlock */
        {
            static bool s_rc_armed_prev = false;
            bool rc_armed_now = safety_rc_alive(now) && crsf_state.armed;
            if (rc_armed_now && !s_rc_armed_prev) {
                /* Rising edge: start arm hold (motors enable after ARMING_HOLD_MS) */
                if (!safety_remote_estop_active() &&
                        mpu6000_is_calibrated() &&
                        bal.state == BALANCE_DISARMED && fabsf(bal.pitch_deg) < 10.0f) {
                    safety_arm_start(now);
                }
            }
            if (!rc_armed_now && s_rc_armed_prev) {
                /* Falling edge: cancel pending arm or disarm if already armed */
                safety_arm_cancel();
                if (bal.state == BALANCE_ARMED) balance_disarm(&bal);
            }
            s_rc_armed_prev = rc_armed_now;
        }

        /* Handle arm/disarm requests from USB with arm-hold interlock */
        if (cdc_arm_request) {
            cdc_arm_request = 0;
            if (!safety_remote_estop_active() &&
                    mpu6000_is_calibrated() &&
                    bal.state == BALANCE_DISARMED && fabsf(bal.pitch_deg) < 10.0f) {
                safety_arm_start(now);
            }
        }
        /* Complete arming once hold timer expires */
        if (safety_arm_ready(now) && bal.state == BALANCE_DISARMED) {
            safety_arm_cancel();
            balance_arm(&bal);
        }
        if (cdc_disarm_request) {
            cdc_disarm_request = 0;
            safety_arm_cancel();
            balance_disarm(&bal);
        }

        /* Gyro recalibration — disarm first, then re-sample bias (~1s blocked) */
        if (cdc_recal_request) {
            cdc_recal_request = 0;
            safety_arm_cancel();
            balance_disarm(&bal);
            if (imu_ret == 0) mpu6000_calibrate();
        }

        /* Handle PID tuning commands from USB (P/I/D/T/M/?) */
        if (cdc_cmd_ready) {
            cdc_cmd_ready = 0;
            char reply[96];
            apply_pid_cmd(&bal, (const char *)cdc_cmd_buf, reply, sizeof(reply));
            CDC_Transmit((uint8_t *)reply, strlen(reply));
        }

        /* Handle Jetson C<speed>,<steer> command (parsed here, not in ISR) */
        if (jetson_cmd_ready) {
            jetson_cmd_ready = 0;
            jetson_cmd_process();
        }

        /* Balance PID at 1kHz — apply Jetson speed offset when active */
        if (imu_ret == 0 && now - balance_tick >= 1) {
            balance_tick = now;
            float base_sp = bal.setpoint;
            if (jetson_cmd_is_active(now)) bal.setpoint += jetson_cmd_sp_offset();
            balance_update(&bal, &imu, dt);
            bal.setpoint = base_sp;
        }

        /* Latch estop on tilt fault or disarm */
        if (bal.state == BALANCE_TILT_FAULT) {
            motor_driver_estop(&motors);
        } else if (bal.state == BALANCE_DISARMED && motors.estop &&
                   !safety_remote_estop_active()) {
            motor_driver_estop_clear(&motors);
        }

        /* Feed autonomous steer from Jetson into mode manager.
         * mode_manager_get_steer() blends it with RC steer per active mode. */
        if (jetson_cmd_is_active(now))
            mode_manager_set_auto_cmd(&mode, jetson_cmd_steer(), 0);

        /* Send to hoverboard ESC at 50Hz (every 20ms) */
        if (now - esc_tick >= 20) {
            esc_tick = now;
            if (bal.state == BALANCE_ARMED) {
                /* Blended steer (RC ↔ auto per mode) + RC speed bias */
                int16_t steer      = mode_manager_get_steer(&mode);
                int16_t spd_bias   = mode_manager_get_speed_bias(&mode);
                int32_t speed      = (int32_t)bal.motor_cmd + spd_bias;
                if (speed >  1000) speed =  1000;
                if (speed < -1000) speed = -1000;
                motor_driver_update(&motors, (int16_t)speed, steer, now);
            } else {
                /* Always send zero while disarmed to prevent ESC timeout */
                motor_driver_update(&motors, 0, 0, now);
            }
        }

        /* CRSF telemetry uplink — battery voltage + arm state at 1 Hz.
         * Sends battery sensor frame (0x08) and flight mode frame (0x21)
         * back to ELRS TX module so the pilot's handset OSD shows Vbat + state. */
        if (now - crsf_telem_tick >= (1000u / CRSF_TELEMETRY_HZ)) {
            crsf_telem_tick = now;
            uint32_t vbat_mv  = battery_read_mv();
            uint8_t  soc_pct  = battery_estimate_pct(vbat_mv);
            crsf_send_battery(vbat_mv, 0u, soc_pct);
            crsf_send_flight_mode(bal.state == BALANCE_ARMED);
        }

        /* USB telemetry at 50Hz (only when streaming enabled and calibration done) */
        if (cdc_streaming && mpu6000_is_calibrated() && now - send_tick >= 20) {
            send_tick = now;
            if (imu_ret == 0) {
                float err = bal.setpoint - bal.pitch_deg;
                /* Build JSON incrementally; append optional sensor fields */
                char *p = buf;
                int rem = (int)sizeof(buf);
                int n;
                n = snprintf(p, rem,
                    "{\"p\":%d,\"r\":%d,\"e\":%d,\"ig\":%d,\"m\":%d,\"s\":%d,\"y\":%d",
                    (int)(bal.pitch_deg * 10),   /* pitch degrees x10 */
                    (int)(imu.roll * 10),        /* roll degrees x10 */
                    (int)(err * 10),             /* PID error x10 */
                    (int)(bal.integral * 10),    /* integral x10 (windup monitor) */
                    (int)bal.motor_cmd,          /* ESC command -1000..+1000 */
                    (int)bal.state,
                    (int)(imu.yaw * 10));        /* yaw degrees x10 (gyro-integrated) */
                p += n; rem -= n;
                n = snprintf(p, rem, ",\"md\":%d",
                    (int)mode_manager_active(&mode));  /* 0=MANUAL,1=ASSISTED,2=AUTO */
                p += n; rem -= n;
                if (mag_type != MAG_NONE) {
                    int16_t hd = mag_read_heading();
                    if (hd >= 0)
                        n = snprintf(p, rem, ",\"hd\":%d", hd); /* heading deg×10 */
                    else
                        n = snprintf(p, rem, ",\"hd\":-1");      /* not ready */
                    p += n; rem -= n;
                }
                if (baro_chip) {
                    int32_t pres_pa; int16_t temp_x10;
                    bmp280_read(&pres_pa, &temp_x10);
                    int32_t alt_cm = bmp280_pressure_to_alt_cm(pres_pa);
                    /* alt cm, temp °C×10, pressure hPa×10 (Pa÷10 = hPa×10) */
                    n = snprintf(p, rem, ",\"alt\":%ld,\"t\":%d,\"pa\":%ld",
                        (long)alt_cm, (int)temp_x10, (long)(pres_pa / 10));
                    p += n; rem -= n;
                    if (baro_chip == 0x60) {  /* BME280: add humidity %RH×10 */
                        int16_t hum_x10 = bmp280_read_humidity();
                        if (hum_x10 >= 0) {
                            n = snprintf(p, rem, ",\"h\":%d", (int)hum_x10);
                            p += n; rem -= n;
                        }
                    }
                }
                if (safety_rc_alive(now)) {
                    /* RSSI, link quality, and CH1–CH4 mapped to µs (1000–2000) */
                    n = snprintf(p, rem, ",\"rssi\":%d,\"lq\":%d"
                        ",\"ch1\":%d,\"ch2\":%d,\"ch3\":%d,\"ch4\":%d",
                        (int)crsf_state.rssi_dbm, (int)crsf_state.link_quality,
                        (int)crsf_to_range(crsf_state.channels[0], 1000, 2000),
                        (int)crsf_to_range(crsf_state.channels[1], 1000, 2000),
                        (int)crsf_to_range(crsf_state.channels[2], 1000, 2000),
                        (int)crsf_to_range(crsf_state.channels[3], 1000, 2000));
                    p += n; rem -= n;
                }
                /* Jetson active flag, USB TX/RX packet counters */
                int es; EstopSource _rs = safety_get_estop();
                if (_rs >= ESTOP_REMOTE) es = (int)_rs;
                else if (bal.state == BALANCE_TILT_FAULT) es = ESTOP_TILT;
                else es = ESTOP_CLEAR;
                n = snprintf(p, rem, ",\"ja\":%d,\"txc\":%u,\"rxc\":%u,\"es\":%d}\n",
                    jetson_cmd_is_active(now) ? 1 : 0,
                    (unsigned)tx_count,
                    (unsigned)cdc_rx_count, es);
                len = (int)(p + n - buf);
            } else {
                len = snprintf(buf, sizeof(buf), "{\"err\":%d}\n", imu_ret);
            }
            CDC_Transmit((uint8_t *)buf, len);
            tx_count++;
        }

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