/* -----------------------------------------------------------------------
 * main.c
 * Self-balancing car main loop
 * Target: STM32F103C8T6 (Blue Pill)
 * Toolchain: ARM GCC
 * ----------------------------------------------------------------------- */

#include "main.h"
#include "mpu6050.h"
#include "motor.h"
#include "encoder.h"
#include "pid.h"
#include "balancer.h"

/* Peripheral handles */
I2C_HandleTypeDef hi2c1;
TIM_HandleTypeDef htim2;    /* PWM for motors */
TIM_HandleTypeDef htim3;    /* Encoder quadrature */
UART_HandleTypeDef huart1;  /* HC-05 Bluetooth */

/* System state */
static Balancer bal;

static uint32_t tick_us;

/* -----------------------------------------------------------------------
 * System clock: 72MHz (HSE 8MHz * 9 PLL)
 * SysTick: 1ms for HAL_Delay
 * ----------------------------------------------------------------------- */
static void SystemClock_Config(void)
{
    RCC_OscInitTypeDef RCC_OscInitStruct = {0};
    RCC_ClkInitTypeDef RCC_ClkInitStruct = {0};

    RCC_OscInitStruct.OscillatorType = RCC_OSCILLATORTYPE_HSE;
    RCC_OscInitStruct.HSEState = RCC_HSE_ON;
    RCC_OscInitStruct.HSEPredivValue = RCC_HSE_PREDIV_DIV1;
    RCC_OscInitStruct.PLL.PLLState = RCC_PLL_ON;
    RCC_OscInitStruct.PLL.PLLSource = RCC_PLLSOURCE_HSE;
    RCC_OscInitStruct.PLL.PLLMUL = RCC_PLL_MUL9;
    HAL_RCC_OscConfig(&RCC_OscInitStruct);

    RCC_ClkInitStruct.ClockType = RCC_CLOCKTYPE_HCLK | RCC_CLOCKTYPE_SYSCLK
                                | RCC_CLOCKTYPE_PCLK1 | RCC_CLOCKTYPE_PCLK2;
    RCC_ClkInitStruct.SYSCLKSource = RCC_SYSCLKSOURCE_PLLCLK;
    RCC_ClkInitStruct.AHBCLKDivider = RCC_SYSCLK_DIV1;
    RCC_ClkInitStruct.APB1CLKDivider = RCC_HCLK_DIV2;
    RCC_ClkInitStruct.APB2CLKDivider = RCC_HCLK_DIV1;
    HAL_RCC_ClockConfig(&RCC_ClkInitStruct, FLASH_LATENCY_2);
}

/* -----------------------------------------------------------------------
 * GPIO init
 * ----------------------------------------------------------------------- */
static void MX_GPIO_Init(void)
{
    GPIO_InitTypeDef GPIO_InitStruct = {0};

    __HAL_RCC_GPIOA_CLK_ENABLE();
    __HAL_RCC_GPIOB_CLK_ENABLE();
    __HAL_RCC_GPIOC_CLK_ENABLE();

    /* LED */
    HAL_GPIO_WritePin(GPIOC, GPIO_PIN_13, GPIO_PIN_SET);
    GPIO_InitStruct.Pin = GPIO_PIN_13;
    GPIO_InitStruct.Mode = GPIO_MODE_OUTPUT_PP;
    GPIO_InitStruct.Pull = GPIO_NOPULL;
    GPIO_InitStruct.Speed = GPIO_SPEED_FREQ_LOW;
    HAL_GPIO_Init(GPIOC, &GPIO_InitStruct);

    /* Motor IN1, IN2 pins */
    GPIO_InitStruct.Pin = GPIO_PIN_1 | GPIO_PIN_2 | GPIO_PIN_4 | GPIO_PIN_5;
    GPIO_InitStruct.Mode = GPIO_MODE_OUTPUT_PP;
    GPIO_InitStruct.Pull = GPIO_NOPULL;
    GPIO_InitStruct.Speed = GPIO_SPEED_FREQ_MEDIUM;
    HAL_GPIO_Init(GPIOA, &GPIO_InitStruct);
}

/* -----------------------------------------------------------------------
 * I2C1 init (MPU6050, OLED)
 * PB6=SCL, PB7=SDA, 400kHz
 * ----------------------------------------------------------------------- */
static void MX_I2C1_Init(void)
{
    hi2c1.Instance = I2C1;
    hi2c1.Init.ClockSpeed = 400000;
    hi2c1.Init.DutyCycle = I2C_DUTYCYCLE_2;
    hi2c1.Init.OwnAddress1 = 0;
    hi2c1.Init.AddressingMode = I2C_ADDRESSINGMODE_7BIT;
    hi2c1.Init.DualAddressMode = I2C_DUALADDRESS_DISABLE;
    hi2c1.Init.GeneralCallMode = I2C_GENERALCALL_DISABLE;
    hi2c1.Init.NoStretchMode = I2C_NOSTRETCH_DISABLE;
    HAL_I2C_Init(&hi2c1);
}

/* -----------------------------------------------------------------------
 * TIM2: PWM for motors (CH1=PA0=L, CH2=PA3=R)
 * 72MHz / 72 = 1MHz → 1kHz PWM (ARR=1000-1)
 * ----------------------------------------------------------------------- */
static void MX_TIM2_Init(void)
{
    TIM_OC_InitTypeDef sConfigOC = {0};

    htim2.Instance = TIM2;
    htim2.Init.Prescaler = 72 - 1;
    htim2.Init.CounterMode = TIM_COUNTERMODE_UP;
    htim2.Init.Period = 1000 - 1;
    htim2.Init.ClockDivision = TIM_CLOCKDIVISION_DIV1;
    htim2.Init.AutoReloadPreload = TIM_AUTORELOAD_PRELOAD_ENABLE;
    HAL_TIM_PWM_Init(&htim2);

    sConfigOC.OCMode = TIM_OCMODE_PWM1;
    sConfigOC.Pulse = 0;
    sConfigOC.OCPolarity = TIM_OCPOLARITY_HIGH;
    sConfigOC.OCFastMode = TIM_OCFAST_DISABLE;
    HAL_TIM_PWM_ConfigChannel(&htim2, &sConfigOC, TIM_CHANNEL_1);
    HAL_TIM_PWM_ConfigChannel(&htim2, &sConfigOC, TIM_CHANNEL_2);
}

/* -----------------------------------------------------------------------
 * TIM3: Quadrature encoder mode (CH1=PA6, CH2=PA7 = left)
 *                           (CH3=PB0, CH4=PB1 = right)
 * ----------------------------------------------------------------------- */
static void MX_TIM3_Init(void)
{
    TIM_Encoder_InitTypeDef sEncoder = {0};
    TIM_MasterConfigTypeDef sMasterConfig = {0};

    htim3.Instance = TIM3;
    htim3.Init.Prescaler = 0;
    htim3.Init.CounterMode = TIM_COUNTERMODE_UP;
    htim3.Init.Period = 0xFFFF;
    htim3.Init.ClockDivision = TIM_CLOCKDIVISION_DIV1;
    htim3.Init.AutoReloadPreload = TIM_AUTORELOAD_PRELOAD_DISABLE;

    sEncoder.EncoderMode = TIM_ENCODERMODE_TI12;
    sEncoder.IC1Polarity = TIM_ICPOLARITY_RISING;
    sEncoder.IC1Selection = TIM_ICSELECTION_DIRECTTI;
    sEncoder.IC1Prescaler = TIM_ICPSC_DIV1;
    sEncoder.IC1Filter = 0;
    sEncoder.IC2Polarity = TIM_ICPOLARITY_RISING;
    sEncoder.IC2Selection = TIM_ICSELECTION_DIRECTTI;
    sEncoder.IC2Prescaler = TIM_ICPSC_DIV1;
    sEncoder.IC2Filter = 0;
    HAL_TIM_Encoder_Init(&htim3, &sEncoder);
}

/* -----------------------------------------------------------------------
 * USART1: HC-05 Bluetooth (PA9=TX, PA10=RX, 115200 8N1)
 * ----------------------------------------------------------------------- */
static void MX_USART1_UART_Init(void)
{
    huart1.Instance = USART1;
    huart1.Init.BaudRate = 115200;
    huart1.Init.WordLength = UART_WORDLENGTH_8B;
    huart1.Init.StopBits = UART_STOPBITS_1;
    huart1.Init.Parity = UART_PARITY_NONE;
    huart1.Init.Mode = UART_MODE_TX_RX;
    huart1.Init.HwFlowCtl = UART_HWCONTROL_NONE;
    huart1.Init.OverSampling = UART_OVERSAMPLING_16;
    HAL_UART_Init(&huart1);
}

/* -----------------------------------------------------------------------
 * Microsecond delay using DWT cycle counter
 * ----------------------------------------------------------------------- */
static void DelayUs_Init(void)
{
    CoreDebug->DEMCR |= CoreDebug_DEMCR_TRCENA_Msk;
    DWT->CYCCNT = 0;
    DWT->CTRL |= DWT_CTRL_CYCCNTENA_Msk;
}

static uint32_t Micros(void)
{
    return DWT->CYCCNT / 72;    /* 72MHz → 72 cycles per µs */
}

static void DelayUs(uint32_t us)
{
    uint32_t start = Micros();
    while ((Micros() - start) < us);
}

/* -----------------------------------------------------------------------
 * UART printf helper (simple, non-blocking)
 * ----------------------------------------------------------------------- */
#include <stdio.h>
#include <stdarg.h>

static void UART_Printf(const char *fmt, ...)
{
    char buf[128];
    va_list args;
    va_start(args, fmt);
    int len = vsnprintf(buf, sizeof(buf), fmt, args);
    va_end(args);

    if (len > 0)
        HAL_UART_Transmit(&huart1, (uint8_t*)buf, len, 10);
}

/* -----------------------------------------------------------------------
 * UART command parser for PID tuning
 * Format: "p=30.0 i=0.5 d=1.5\r\n"
 *         "stop\r\n"
 *         "start\r\n"
 *         "s=0.0\r\n"  (speed reference)
 * ----------------------------------------------------------------------- */
static void ParseUARTCommand(const char *cmd)
{
    if (strncmp(cmd, "p=", 2) == 0) {
        float kp, ki, kd;
        sscanf(cmd, "p=%f i=%f d=%f", &kp, &ki, &kd);
        PID_SetGains(&bal.pid_angle, kp, ki, kd);
        UART_Printf("OK: Kp=%.2f Ki=%.2f Kd=%.2f\r\n", kp, ki, kd);
    } else if (strncmp(cmd, "sp=", 3) == 0) {
        float kp, ki;
        sscanf(cmd, "sp=%f si=%f", &kp, &ki);
        PID_SetGains(&bal.pid_speed, kp, ki, 0);
        UART_Printf("OK: Spd_Kp=%.2f Spd_Ki=%.2f\r\n", kp, ki);
    } else if (strncmp(cmd, "start", 5) == 0) {
        Balancer_Calibrate(&bal);
        bal.running = 1;
        UART_Printf("OK: started\r\n");
    } else if (strncmp(cmd, "stop", 4) == 0) {
        Balancer_EmergencyStop(&bal);
        UART_Printf("OK: stopped\r\n");
    } else if (strncmp(cmd, "s=", 2) == 0) {
        float speed;
        sscanf(cmd, "s=%f", &speed);
        bal.speed_ref = speed;
        UART_Printf("OK: speed_ref=%.1f\r\n", speed);
    } else if (strncmp(cmd, "?state", 6) == 0) {
        UART_Printf("pitch=%.1f angle_ref=%.1f mot_l=%d mot_r=%d\r\n",
                    bal.att.pitch, bal.angle_ref,
                    (int)bal.motor_output_l, (int)bal.motor_output_r);
    }
}

static void UART_RxCallback(uint8_t byte)
{
    static char line[64];
    static uint8_t idx = 0;

    if (byte == '\n' || byte == '\r') {
        line[idx] = '\0';
        if (idx > 0)
            ParseUARTCommand(line);
        idx = 0;
    } else if (idx < sizeof(line) - 1) {
        line[idx++] = (char)byte;
    }
}

/* -----------------------------------------------------------------------
 * main()
 * ----------------------------------------------------------------------- */
int main(void)
{
    HAL_Init();
    SystemClock_Config();
    MX_GPIO_Init();
    MX_I2C1_Init();
    MX_TIM2_Init();
    MX_TIM3_Init();
    MX_USART1_UART_Init();
    DelayUs_Init();

    /* Init MPU6050 */
    if (MPU6050_Init(&hi2c1) != 0) {
        UART_Printf("MPU6050 init FAILED\r\n");
        while (1);  /* halt */
    }
    UART_Printf("MPU6050 OK\r\n");

    /* Init balancer */
    Balancer_Init(&bal);

    /* Calibrate gyro */
    MPU6050_CalibrateGyro(&hi2c1,
                          &bal.gyro_bias_x, &bal.gyro_bias_y, &bal.gyro_bias_z,
                          200);

    UART_Printf("Gyro Bias: %.3f %.3f %.3f\r\n",
                bal.gyro_bias_x, bal.gyro_bias_y, bal.gyro_bias_z);

    bal.calibrated = 1;

    /* Wait for car to be placed upright */
    HAL_Delay(1000);

    uint32_t prev_tick = HAL_GetTick();
    uint32_t speed_tick = prev_tick;
    uint32_t debug_tick = prev_tick;
    MPU6050_Data sensor;
    float dt_angle, dt_speed;

    UART_Printf("Ready. Send 'start' to begin.\r\n");

    /* Main loop */
    while (1)
    {
        uint32_t now = HAL_GetTick();
        dt_angle = (now - prev_tick) * 0.001f;
        dt_speed = (now - speed_tick) * 0.001f;

        /* Read sensor at max rate */
        MPU6050_ReadScaled(&hi2c1, &sensor);

        /* Apply gyro bias */
        sensor.gx -= bal.gyro_bias_x;
        sensor.gy -= bal.gyro_bias_y;
        sensor.gz -= bal.gyro_bias_z;

        bal.sensor = sensor;

        /* Complementary filter at ~1kHz, but attitude updates at angle rate */
        Attitude_Complementary(&bal.att, &sensor, dt_angle, COMP_ALPHA);

        /* Safety check */
        if (fabsf(bal.att.pitch) > ANGLE_LIMIT_DEG) {
            Balancer_EmergencyStop(&bal);
            UART_Printf("FAULT: tilt>%.0fdeg\r\n", ANGLE_LIMIT_DEG);
        }

        /* ===== Angle PID (inner loop) ~200Hz ===== */
        if (now - bal.tick_angle >= LOOP_ANGLE_PERIOD_MS) {
            bal.tick_angle = now;

            if (bal.running) {
                float pid_out = PID_Update(&bal.pid_angle,
                                           bal.angle_ref, bal.att.pitch, dt_angle);
                bal.motor_output = pid_out;
            } else {
                bal.motor_output = 0;
            }

            prev_tick = now;
        }

        /* ===== Speed PID (outer loop) ~50Hz ===== */
        if (now - bal.tick_speed >= LOOP_SPEED_PERIOD_MS) {
            bal.tick_speed = now;

            /* Update encoders */
            Encoder_Update(&bal.enc_l, dt_speed);
            Encoder_Update(&bal.enc_r, dt_speed);

            if (bal.running) {
                float avg_speed = (bal.enc_l.speed_filtered +
                                   bal.enc_r.speed_filtered) * 0.5f;
                float angle_target = PID_Update(&bal.pid_speed,
                                                 bal.speed_ref, avg_speed, dt_speed);
                bal.angle_ref = angle_target;
            } else {
                bal.angle_ref = 0;
            }

            speed_tick = now;
        }

        /* Apply motor output */
        {
            float out_l = bal.motor_output;
            float out_r = bal.motor_output;

            /* Differential steering? Not needed for basic balance. */
            bal.motor_output_l = out_l;
            bal.motor_output_r = out_r;

            Motor_SetSpeedScaled(&bal.motor_l, (int16_t)out_l);
            Motor_SetSpeedScaled(&bal.motor_r, (int16_t)out_r);
        }

        /* Debug output at ~2Hz */
        if (now - debug_tick >= 500) {
            debug_tick = now;
            UART_Printf("pitch=%.1f ref=%.1f out=%d spd=%.1f\r\n",
                        bal.att.pitch, bal.angle_ref,
                        (int)bal.motor_output,
                        (bal.enc_l.speed_filtered + bal.enc_r.speed_filtered) * 0.5f);
            HAL_GPIO_TogglePin(GPIOC, GPIO_PIN_13);
        }

        /* Feed watchdog (IWDG expected in production) */
    }
}