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

/* ================================================================
 * WS2812B NeoPixel protocol via PWM
 * ================================================================
 * 800 kHz PWM → 1.25 µs per cycle
 * Bit encoding:
 *   "0": High 350 ns  (40% duty)  → ~3/8 of 1.25 µs
 *   "1": High 700 ns  (56% duty)  → ~7/10 of 1.25 µs
 * Reset: Low > 50 µs (automatic with DMA ring and reload)
 *
 * Implementation: DMA copies PWM duty values from buffer.
 * Each bit needs one PWM cycle; 192 bits total (24 bits/LED × 8 LEDs).
 */

#define LED_BITS_PER_COLOR  8u
#define LED_BITS_PER_LED    (LED_BITS_PER_COLOR * 3u)  /* RGB */
#define LED_TOTAL_BITS      (LED_BITS_PER_LED * LED_STRIP_NUM_LEDS)
#define LED_PWM_PERIOD      (216000000 / LED_STRIP_FREQ_HZ)  /* 216 MHz / 800 kHz */

/* PWM duty values for bit encoding (out of LED_PWM_PERIOD) */
#define LED_BIT_0_DUTY      (LED_PWM_PERIOD * 40 / 100)   /* ~350 ns high */
#define LED_BIT_1_DUTY      (LED_PWM_PERIOD * 56 / 100)   /* ~700 ns high */

/* ================================================================
 * LED buffer and animation state
 * ================================================================
 */

typedef struct {
    RGBColor leds[LED_STRIP_NUM_LEDS];
    uint32_t pwm_buf[LED_TOTAL_BITS];  /* DMA buffer: PWM duty values */
} LEDBuffer;

/* LED state machine */
typedef struct {
    LEDState current_state;
    LEDState next_state;
    uint32_t state_start_ms;
    uint8_t animation_phase;  /* 0-255 for continuous animations */
} LEDAnimState;

static LEDBuffer      s_led_buf = {0};
static LEDAnimState   s_anim = {0};
static TIM_HandleTypeDef s_tim_handle = {0};

/* ================================================================
 * Helper functions
 * ================================================================
 */

static void rgb_to_pwm_buffer(const RGBColor *colors, uint8_t num_leds)
{
    /* Encode LED colors into PWM duty values for WS2812B transmission.
     * GRB byte order (WS2812B standard), MSB first. */
    uint32_t buf_idx = 0;

    for (uint8_t led = 0; led < num_leds; led++) {
        uint8_t g = colors[led].g;
        uint8_t r = colors[led].r;
        uint8_t b = colors[led].b;

        /* GRB byte order */
        uint8_t bytes[3] = {g, r, b};

        for (int byte_idx = 0; byte_idx < 3; byte_idx++) {
            uint8_t byte = bytes[byte_idx];

            /* MSB first — encode 8 bits */
            for (int bit = 7; bit >= 0; bit--) {
                uint8_t bit_val = (byte >> bit) & 1;
                s_led_buf.pwm_buf[buf_idx++] = bit_val ? LED_BIT_1_DUTY : LED_BIT_0_DUTY;
            }
        }
    }
}

static uint8_t sin_u8(uint8_t phase)
{
    /* Approximate sine wave (0-255) from phase (0-255) for breathing effect. */
    static const uint8_t sine_lut[256] = {
        128, 131, 134, 137, 140, 143, 146, 149, 152, 155, 158, 161, 164, 167, 170, 173,
        176, 179, 182, 185, 188, 191, 193, 196, 199, 201, 204, 206, 209, 211, 214, 216,
        218, 221, 223, 225, 227, 229, 231, 233, 235, 236, 238, 240, 241, 243, 244, 245,
        247, 248, 249, 250, 251, 252, 252, 253, 254, 254, 255, 255, 255, 255, 255, 254,
        254, 253, 252, 252, 251, 250, 249, 248, 247, 245, 244, 243, 241, 240, 238, 236,
        235, 233, 231, 229, 227, 225, 223, 221, 218, 216, 214, 211, 209, 206, 204, 201,
        199, 196, 193, 191, 188, 185, 182, 179, 176, 173, 170, 167, 164, 161, 158, 155,
        152, 149, 146, 143, 140, 137, 134, 131, 128, 125, 122, 119, 116, 113, 110, 107,
        104, 101, 98, 95, 92, 89, 86, 83, 80, 77, 74, 71, 68, 65, 62, 59,
        56, 53, 50, 47, 44, 41, 39, 36, 33, 31, 28, 26, 23, 21, 18, 16,
        14, 11, 9, 7, 5, 3, 1, 0, 0, 0, 0, 0, 1, 2, 3, 4,
        5, 7, 8, 10, 11, 13, 15, 17, 19, 21, 23, 26, 28, 31, 33, 36,
        39, 42, 45, 48, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 82, 85,
        88, 92, 95, 99, 102, 105, 109, 113, 116, 120, 124, 127, 131
    };
    return sine_lut[phase];
}

/* ================================================================
 * Animation implementations
 * ================================================================
 */

static void animate_boot(uint32_t elapsed_ms)
{
    /* Blue chase: rotate a single LED around the ring. */
    uint8_t led_idx = (elapsed_ms / 100) % LED_STRIP_NUM_LEDS;  /* 100 ms per LED */

    memset(s_led_buf.leds, 0, sizeof(s_led_buf.leds));
    s_led_buf.leds[led_idx].b = 255;  /* Bright blue */

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

static void animate_armed(void)
{
    /* Solid green: all LEDs constant brightness. */
    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].g = 200;  /* Bright green */
        s_led_buf.leds[i].r = 0;
        s_led_buf.leds[i].b = 0;
    }

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

static void animate_error(uint32_t elapsed_ms)
{
    /* Red blinking: on/off every 250 ms. */
    bool on = ((elapsed_ms / 250) % 2) == 0;

    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].r = on ? 255 : 0;
        s_led_buf.leds[i].g = 0;
        s_led_buf.leds[i].b = 0;
    }

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

static void animate_low_battery(uint32_t elapsed_ms)
{
    /* Yellow pulsing: brightness varies smoothly. */
    uint8_t phase = (elapsed_ms / 20) & 0xFF;  /* Cycle every 5120 ms */
    uint8_t brightness = sin_u8(phase);

    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].r = (brightness * 255) >> 8;
        s_led_buf.leds[i].g = (brightness * 255) >> 8;
        s_led_buf.leds[i].b = 0;
    }

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

static void animate_charging(uint32_t elapsed_ms)
{
    /* Green breathing: smooth brightness modulation. */
    uint8_t phase = (elapsed_ms / 20) & 0xFF;  /* Cycle every 5120 ms */
    uint8_t brightness = sin_u8(phase);

    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].g = (brightness * 255) >> 8;
        s_led_buf.leds[i].r = 0;
        s_led_buf.leds[i].b = 0;
    }

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

static void animate_estop(uint32_t elapsed_ms)
{
    /* Red strobe: on/off every 125 ms (8 Hz). */
    bool on = ((elapsed_ms / 125) % 2) == 0;

    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].r = on ? 255 : 0;
        s_led_buf.leds[i].g = 0;
        s_led_buf.leds[i].b = 0;
    }

    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

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

void led_init(void)
{
    /* Initialize state machine */
    s_anim.current_state = LED_STATE_BOOT;
    s_anim.next_state = LED_STATE_BOOT;
    s_anim.state_start_ms = 0;
    s_anim.animation_phase = 0;

    /* Configure GPIO PB4 as TIM3_CH1 output (AF2) */
    __HAL_RCC_GPIOB_CLK_ENABLE();
    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_FREQ_HIGH;
    gpio_init.Alternate = LED_STRIP_AF;
    HAL_GPIO_Init(LED_STRIP_PORT, &gpio_init);

    /* Configure TIM3: PWM mode, 800 kHz frequency
     * STM32F722 has 216 MHz on APB1; TIM3 is on APB1 (prescaler 4×).
     * APB1 clock: 216 MHz / 4 = 54 MHz
     * For 800 kHz PWM: 54 MHz / 800 kHz = 67.5 → use 67 or 68
     * With ARR = 67: 54 MHz / 68 = 794 kHz ≈ 800 kHz
     */
    __HAL_RCC_TIM3_CLK_ENABLE();

    s_tim_handle.Instance           = LED_STRIP_TIM;
    s_tim_handle.Init.Prescaler     = 0;  /* No prescaler; APB1 = 54 MHz directly */
    s_tim_handle.Init.CounterMode   = TIM_COUNTERMODE_UP;
    s_tim_handle.Init.Period        = LED_PWM_PERIOD - 1;
    s_tim_handle.Init.ClockDivision = TIM_CLOCKDIVISION_DIV1;
    s_tim_handle.Init.RepetitionCounter = 0;

    HAL_TIM_PWM_Init(&s_tim_handle);

    /* Configure TIM3_CH1 for PWM */
    TIM_OC_InitTypeDef oc_init = {0};
    oc_init.OCMode       = TIM_OCMODE_PWM1;
    oc_init.Pulse        = 0;  /* Start at 0% duty */
    oc_init.OCPolarity   = TIM_OCPOLARITY_HIGH;
    oc_init.OCFastMode   = TIM_OCFAST_DISABLE;

    HAL_TIM_PWM_ConfigChannel(&s_tim_handle, &oc_init, LED_STRIP_CHANNEL);
    HAL_TIM_PWM_Start(&s_tim_handle, LED_STRIP_CHANNEL);

    /* Initialize LED buffer with boot state */
    animate_boot(0);
}

void led_set_state(LEDState state)
{
    if (state >= LED_STATE_COUNT) {
        return;
    }
    s_anim.next_state = state;
}

LEDState led_get_state(void)
{
    return s_anim.current_state;
}

void led_set_color(uint8_t r, uint8_t g, uint8_t b)
{
    for (uint8_t i = 0; i < LED_STRIP_NUM_LEDS; i++) {
        s_led_buf.leds[i].r = r;
        s_led_buf.leds[i].g = g;
        s_led_buf.leds[i].b = b;
    }
    rgb_to_pwm_buffer(s_led_buf.leds, LED_STRIP_NUM_LEDS);
}

void led_tick(uint32_t now_ms)
{
    /* State transition */
    if (s_anim.next_state != s_anim.current_state) {
        s_anim.current_state = s_anim.next_state;
        s_anim.state_start_ms = now_ms;
    }

    uint32_t elapsed = now_ms - s_anim.state_start_ms;

    /* Run state-specific animation */
    switch (s_anim.current_state) {
        case LED_STATE_BOOT:
            animate_boot(elapsed);
            break;
        case LED_STATE_ARMED:
            animate_armed();
            break;
        case LED_STATE_ERROR:
            animate_error(elapsed);
            break;
        case LED_STATE_LOW_BATT:
            animate_low_battery(elapsed);
            break;
        case LED_STATE_CHARGING:
            animate_charging(elapsed);
            break;
        case LED_STATE_ESTOP:
            animate_estop(elapsed);
            break;
        default:
            break;
    }
}

bool led_is_animating(void)
{
    /* Static states: ARMED (always) and ERROR (after first blink) */
    /* All others animate continuously */
    return s_anim.current_state != LED_STATE_ARMED;
}
