#include "rgb_fsm.h"
#include "stm32f7xx_hal.h"
#include "config.h"
#include <string.h>
#include <math.h>

/* ================================================================
 * WS2812 NeoPixel LED Strip Configuration
 * ================================================================ */

#define NUM_LEDS        8
#define LED_STRIP_BITS  (NUM_LEDS * 24)  /* 24 bits per LED (RGB) */

/* ================================================================
 * State Machine Internal State
 * ================================================================ */

typedef struct {
    LedState current_state;             /* Current operational state */
    LedState previous_state;            /* Previous state for transition detection */
    uint32_t state_start_time_ms;       /* When current state started */
    uint32_t last_tick_ms;              /* Last tick time */
    uint8_t animation_frame;            /* Current animation frame (0-255) */
    RgbColor led_colors[NUM_LEDS];      /* Current color of each LED */
} RgbFsm;

static RgbFsm s_rgb = {
    .current_state = LED_STATE_BOOT,
    .previous_state = LED_STATE_BOOT,
    .state_start_time_ms = 0,
    .last_tick_ms = 0,
    .animation_frame = 0
};

/* ================================================================
 * Color Definitions for Each State
 * ================================================================ */

static const RgbColor COLOR_BLUE   = {  0,   0, 255 };
static const RgbColor COLOR_GREEN  = {  0, 255,   0 };
static const RgbColor COLOR_CYAN   = {  0, 255, 255 };
static const RgbColor COLOR_RED    = {255,   0,   0 };
static const RgbColor COLOR_ORANGE = {255, 165,   0 };
static const RgbColor COLOR_OFF    = {  0,   0,   0 };

/* ================================================================
 * Hardware Initialization
 * ================================================================ */

void rgb_fsm_init(void)
{
    /* Enable GPIO and timer clocks */
    __HAL_RCC_GPIOB_CLK_ENABLE();
    __HAL_RCC_TIM3_CLK_ENABLE();

    /* Configure PB4 as TIM3_CH1 PWM output */
    GPIO_InitTypeDef gpio_init = {0};
    gpio_init.Pin = LED_STRIP_PIN;
    gpio_init.Mode = GPIO_MODE_AF_PP;
    gpio_init.Pull = GPIO_NOPULL;
    gpio_init.Speed = GPIO_SPEED_HIGH;
    gpio_init.Alternate = LED_STRIP_AF;
    HAL_GPIO_Init(LED_STRIP_PORT, &gpio_init);

    /* Configure TIM3 for 800 kHz PWM
     * Clock: 216MHz / PSC = output frequency
     * For 800 kHz: PSC = 270, ARR = 100
     * Duty cycle = CCR / ARR
     */
    TIM_HandleTypeDef htim3 = {0};
    htim3.Instance = LED_STRIP_TIM;
    htim3.Init.Prescaler = 270 - 1;         /* 216MHz / 270 = 800kHz clock */
    htim3.Init.CounterMode = TIM_COUNTERMODE_UP;
    htim3.Init.Period = 100 - 1;            /* 800kHz / 100 = 8 kHz PWM */
    htim3.Init.ClockDivision = TIM_CLOCKDIVISION_DIV1;
    htim3.Init.RepetitionCounter = 0;
    HAL_TIM_PWM_Init(&htim3);

    /* Configure PWM on CH1 for WS2812 signal */
    TIM_OC_InitTypeDef oc_init = {0};
    oc_init.OCMode = TIM_OCMODE_PWM1;
    oc_init.Pulse = 0;                     /* Start at 0% duty (off) */
    oc_init.OCPolarity = TIM_OCPOLARITY_HIGH;
    oc_init.OCFastMode = TIM_OCFAST_DISABLE;
    HAL_TIM_PWM_ConfigChannel(&htim3, &oc_init, LED_STRIP_CHANNEL);

    /* Start PWM generation */
    HAL_TIM_PWM_Start(LED_STRIP_TIM, LED_STRIP_CHANNEL);

    /* Initialize state */
    s_rgb.current_state = LED_STATE_BOOT;
    s_rgb.previous_state = LED_STATE_BOOT;
    s_rgb.state_start_time_ms = 0;
    s_rgb.animation_frame = 0;
    memset(s_rgb.led_colors, 0, sizeof(s_rgb.led_colors));
}

/* ================================================================
 * LED Update Function
 * ================================================================ */

static void rgb_update_led_strip(void)
{
    /* Calculate total brightness from LED colors */
    uint32_t total_brightness = 0;
    for (int i = 0; i < NUM_LEDS; i++) {
        total_brightness += s_rgb.led_colors[i].r;
        total_brightness += s_rgb.led_colors[i].g;
        total_brightness += s_rgb.led_colors[i].b;
    }
    /* Normalize to 0-100 (PWM duty cycle) */
    uint32_t duty = (total_brightness / (NUM_LEDS * 3)) * 100 / 255;
    if (duty > 100) duty = 100;

    TIM3->CCR1 = (duty * 100) / 100;  /* Set duty cycle */
}

/* ================================================================
 * Animation Functions
 * ================================================================ */

static void animate_boot(uint32_t elapsed_ms)
{
    /* Blue pulse at 0.5 Hz (2 second period) */
    uint32_t period_ms = 2000;
    uint8_t phase = (elapsed_ms % period_ms) * 255 / period_ms;
    uint8_t brightness = (uint8_t)(128 + 127 * sin(2 * 3.14159 * phase / 256));

    RgbColor color = COLOR_BLUE;
    color.b = (color.b * brightness) / 255;

    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = color;
    }
    s_rgb.animation_frame = phase;
}

static void animate_idle(uint32_t elapsed_ms)
{
    /* Green breathe at 0.5 Hz (2 second period) */
    uint32_t period_ms = 2000;
    uint8_t phase = (elapsed_ms % period_ms) * 255 / period_ms;
    uint8_t brightness = (uint8_t)(128 + 127 * sin(2 * 3.14159 * phase / 256));

    RgbColor color = COLOR_GREEN;
    color.g = (color.g * brightness) / 255;

    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = color;
    }
    s_rgb.animation_frame = phase;
}

static void animate_armed(uint32_t elapsed_ms)
{
    /* Solid green (no animation) */
    (void)elapsed_ms;
    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = COLOR_GREEN;
    }
    s_rgb.animation_frame = 255;
}

static void animate_nav(uint32_t elapsed_ms)
{
    /* Cyan spin (rotating pattern at 1 Hz) */
    uint32_t period_ms = 1000;
    uint8_t phase = (elapsed_ms % period_ms) * 8 / period_ms;

    for (int i = 0; i < NUM_LEDS; i++) {
        if (i == phase) {
            s_rgb.led_colors[i] = COLOR_CYAN;
        } else if (i == (phase + 7) % 8) {
            s_rgb.led_colors[i] = (RgbColor){0, 128, 128};
        } else {
            s_rgb.led_colors[i] = COLOR_OFF;
        }
    }
    s_rgb.animation_frame = (uint8_t)phase * 32;
}

static void animate_error(uint32_t elapsed_ms)
{
    /* Red flash at 2 Hz (500ms period) */
    uint32_t period_ms = 500;
    uint8_t brightness = ((elapsed_ms % period_ms) < 250) ? 255 : 0;

    RgbColor color = COLOR_RED;
    color.r = (color.r * brightness) / 255;

    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = color;
    }
    s_rgb.animation_frame = brightness;
}

static void animate_low_batt(uint32_t elapsed_ms)
{
    /* Orange blink at 1 Hz (1000ms period) */
    uint32_t period_ms = 1000;
    uint8_t brightness = ((elapsed_ms % period_ms) < 500) ? 255 : 0;

    RgbColor color = COLOR_ORANGE;
    color.r = (color.r * brightness) / 255;
    color.g = (color.g * brightness) / 255;

    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = color;
    }
    s_rgb.animation_frame = brightness;
}

static void animate_charging(uint32_t elapsed_ms)
{
    /* Green fill (progressive LEDs lighting up at 1 Hz) */
    uint32_t period_ms = 1000;
    uint8_t phase = (elapsed_ms % period_ms) * 8 / period_ms;

    for (int i = 0; i < NUM_LEDS; i++) {
        if (i < phase) {
            s_rgb.led_colors[i] = COLOR_GREEN;
        } else {
            s_rgb.led_colors[i] = COLOR_OFF;
        }
    }
    s_rgb.animation_frame = phase * 32;
}

static void animate_estop(uint32_t elapsed_ms)
{
    /* Red solid (full intensity, no animation) */
    (void)elapsed_ms;
    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = COLOR_RED;
    }
    s_rgb.animation_frame = 255;
}

/* ================================================================
 * Public API
 * ================================================================ */

bool rgb_fsm_set_state(LedState state)
{
    if (state >= LED_STATE_COUNT) {
        return false;
    }

    if (state == s_rgb.current_state) {
        return false;
    }

    s_rgb.previous_state = s_rgb.current_state;
    s_rgb.current_state = state;
    s_rgb.state_start_time_ms = 0;
    s_rgb.animation_frame = 0;

    return true;
}

LedState rgb_fsm_get_state(void)
{
    return s_rgb.current_state;
}

void rgb_fsm_tick(uint32_t now_ms)
{
    if (s_rgb.state_start_time_ms == 0) {
        s_rgb.state_start_time_ms = now_ms;
        s_rgb.last_tick_ms = now_ms;
        return;
    }

    uint32_t elapsed = now_ms - s_rgb.state_start_time_ms;

    switch (s_rgb.current_state) {
        case LED_STATE_BOOT:
            animate_boot(elapsed);
            break;
        case LED_STATE_IDLE:
            animate_idle(elapsed);
            break;
        case LED_STATE_ARMED:
            animate_armed(elapsed);
            break;
        case LED_STATE_NAV:
            animate_nav(elapsed);
            break;
        case LED_STATE_ERROR:
            animate_error(elapsed);
            break;
        case LED_STATE_LOW_BATT:
            animate_low_batt(elapsed);
            break;
        case LED_STATE_CHARGING:
            animate_charging(elapsed);
            break;
        case LED_STATE_ESTOP:
            animate_estop(elapsed);
            break;
        default:
            rgb_fsm_all_off();
            break;
    }

    rgb_update_led_strip();
    s_rgb.last_tick_ms = now_ms;
}

bool rgb_fsm_set_color(uint8_t led_index, RgbColor color)
{
    if (led_index >= NUM_LEDS) {
        return false;
    }

    s_rgb.led_colors[led_index] = color;
    rgb_update_led_strip();
    return true;
}

void rgb_fsm_all_off(void)
{
    for (int i = 0; i < NUM_LEDS; i++) {
        s_rgb.led_colors[i] = COLOR_OFF;
    }
    rgb_update_led_strip();
}

uint8_t rgb_fsm_get_animation_frame(void)
{
    return s_rgb.animation_frame;
}