#include "power_mgmt.h"
#include "config.h"
#include "stm32f7xx_hal.h"
#include <string.h>

/* ---- Internal state ---- */
static PowerState  s_state             = PM_ACTIVE;
static uint32_t    s_last_active       = 0;
static uint32_t    s_fade_start        = 0;
static bool        s_sleep_req         = false;
static bool        s_peripherals_gated = false;

/* ---- EXTI wake-source configuration ---- */
/*
 * EXTI1  → PA1 (UART4_RX / CRSF): falling edge (UART start bit)
 * EXTI7  → PB7 (USART1_RX / JLink): falling edge
 * EXTI4  → PC4 (MPU6000 INT): already configured by mpu6000_init();
 *           we just ensure IMR bit is set.
 *
 * GPIO pins remain in their current AF mode; EXTI is pad-level and
 * fires independently of the AF setting.
 */
static void enable_wake_exti(void)
{
    __HAL_RCC_SYSCFG_CLK_ENABLE();

    /* EXTI1: PA1 (UART4_RX) — SYSCFG EXTICR1[7:4] = 0000 (PA) */
    SYSCFG->EXTICR[0] = (SYSCFG->EXTICR[0] & ~(0xFu << 4)) | (0x0u << 4);
    EXTI->FTSR |=  (1u << 1);
    EXTI->RTSR &= ~(1u << 1);
    EXTI->PR    =  (1u << 1);   /* clear pending */
    EXTI->IMR  |=  (1u << 1);
    HAL_NVIC_SetPriority(EXTI1_IRQn, 5, 0);
    HAL_NVIC_EnableIRQ(EXTI1_IRQn);

    /* EXTI7: PC7 (USART6_RX / Jetson UART) — SYSCFG EXTICR2[15:12] = 0010 (PC)
     * Changed from PB7 (JLink) to PC7 (Jetson) — Jetson is primary interface.
     * JLink wake is handled by Jetson timeout instead. */
    SYSCFG->EXTICR[1] = (SYSCFG->EXTICR[1] & ~(0xFu << 12)) | (0x2u << 12);
    EXTI->FTSR |=  (1u << 7);
    EXTI->RTSR &= ~(1u << 7);
    EXTI->PR    =  (1u << 7);
    EXTI->IMR  |=  (1u << 7);
    HAL_NVIC_SetPriority(EXTI9_5_IRQn, 5, 0);
    HAL_NVIC_EnableIRQ(EXTI9_5_IRQn);

    /* EXTI4: PC4 (MPU6000 INT) — handler in mpu6000.c; just ensure IMR set */
    EXTI->IMR |= (1u << 4);
}

static void disable_wake_exti(void)
{
    /* Mask UART RX wake EXTIs now that UART peripherals handle traffic */
    EXTI->IMR &= ~(1u << 1);
    EXTI->IMR &= ~(1u << 7);
    /* Leave EXTI4 (IMU data-ready) always unmasked */
}

/* ---- Peripheral clock gating ---- */
/*
 * Clock-only gate (no force-reset): peripheral register state is preserved.
 * On re-enable, DMA circular transfers resume without reinitialisation.
 */
static void gate_peripherals(void)
{
    if (s_peripherals_gated) return;
    __HAL_RCC_SPI3_CLK_DISABLE();    /* I2S3 / audio amplifier */
    __HAL_RCC_SPI2_CLK_DISABLE();    /* OSD MAX7456 */
    // __HAL_RCC_USART6_CLK_DISABLE(); // kept active for Jetson UART  /* legacy Jetson CDC */
    /* UART5 kept active — hoverboard ESC needs continuous comms */
    /* __HAL_RCC_UART5_CLK_DISABLE(); */
    s_peripherals_gated = true;
}

static void ungate_peripherals(void)
{
    if (!s_peripherals_gated) return;
    __HAL_RCC_SPI3_CLK_ENABLE();
    __HAL_RCC_SPI2_CLK_ENABLE();
    __HAL_RCC_USART6_CLK_ENABLE();
    __HAL_RCC_UART5_CLK_ENABLE();
    s_peripherals_gated = false;
}

/* ---- PLL clock restore after STOP mode ---- */
/*
 * After STOP wakeup SYSCLK = HSI (16 MHz). Re-lock PLL for 216 MHz.
 * PLLM=8, PLLN=216, PLLP=2, PLLQ=9 — STM32F722 @ 216 MHz, HSI source.
 *
 * HAL_RCC_ClockConfig() calls HAL_InitTick() which resets uwTick to 0;
 * save and restore it so existing timeouts remain valid across sleep.
 */
extern volatile uint32_t uwTick;

static void restore_clocks(void)
{
    uint32_t saved_tick = uwTick;

    RCC_OscInitTypeDef osc = {0};
    osc.OscillatorType      = RCC_OSCILLATORTYPE_HSI;
    osc.HSIState            = RCC_HSI_ON;
    osc.HSICalibrationValue = RCC_HSICALIBRATION_DEFAULT;
    osc.PLL.PLLState        = RCC_PLL_ON;
    osc.PLL.PLLSource       = RCC_PLLSOURCE_HSI;
    osc.PLL.PLLM            = 8;
    osc.PLL.PLLN            = 216;
    osc.PLL.PLLP            = RCC_PLLP_DIV2;
    osc.PLL.PLLQ            = 9;
    HAL_RCC_OscConfig(&osc);

    RCC_ClkInitTypeDef clk = {0};
    clk.ClockType      = RCC_CLOCKTYPE_HCLK | RCC_CLOCKTYPE_SYSCLK |
                         RCC_CLOCKTYPE_PCLK1 | RCC_CLOCKTYPE_PCLK2;
    clk.SYSCLKSource   = RCC_SYSCLKSOURCE_PLLCLK;
    clk.AHBCLKDivider  = RCC_SYSCLK_DIV1;
    clk.APB1CLKDivider = RCC_HCLK_DIV4;   /* 54 MHz */
    clk.APB2CLKDivider = RCC_HCLK_DIV2;   /* 108 MHz */
    HAL_RCC_ClockConfig(&clk, FLASH_LATENCY_7);

    uwTick = saved_tick;  /* restore — HAL_InitTick() reset it to 0 */
}

/* ---- EXTI IRQ handlers (wake-only: clear pending bit and return) ---- */
/*
 * These handlers fire once on wakeup. After restore_clocks() the respective
 * UART peripherals resume normal DMA/IDLE-interrupt operation.
 *
 * NOTE: If EXTI9_5_IRQHandler is already defined elsewhere in the project,
 * merge that handler with this one.
 */
void EXTI1_IRQHandler(void)
{
    if (EXTI->PR & (1u << 1)) EXTI->PR = (1u << 1);
}

void EXTI9_5_IRQHandler(void)
{
    /* Clear any pending EXTI5-9 lines (PB7 = EXTI7 is our primary wake) */
    uint32_t pr = EXTI->PR & 0x3E0u;
    if (pr) EXTI->PR = pr;
}

/* ---- LED brightness (integer arithmetic, no float, called from main loop) ---- */
/*
 * Triangle wave: 0→255→0 over PM_LED_PERIOD_MS.
 * Only active during PM_SLEEP_PENDING; returns 0 otherwise.
 */
uint8_t power_mgmt_led_brightness(void)
{
    if (s_state != PM_SLEEP_PENDING) return 0u;

    uint32_t phase = (HAL_GetTick() - s_fade_start) % PM_LED_PERIOD_MS;
    uint32_t half  = PM_LED_PERIOD_MS / 2u;
    if (phase < half)
        return (uint8_t)(phase * 255u / half);
    else
        return (uint8_t)((PM_LED_PERIOD_MS - phase) * 255u / half);
}

/* ---- Current estimate ---- */
uint16_t power_mgmt_current_ma(void)
{
    if (s_state == PM_SLEEPING)
        return (uint16_t)PM_CURRENT_STOP_MA;
    uint16_t ma = (uint16_t)PM_CURRENT_BASE_MA;
    if (!s_peripherals_gated) {
        ma += (uint16_t)(PM_CURRENT_AUDIO_MA + PM_CURRENT_OSD_MA +
                         PM_CURRENT_DEBUG_MA);
    }
    return ma;
}

/* ---- Idle elapsed ---- */
uint32_t power_mgmt_idle_ms(void)
{
    return HAL_GetTick() - s_last_active;
}

/* ---- Public API ---- */
void power_mgmt_init(void)
{
    s_state             = PM_ACTIVE;
    s_last_active       = HAL_GetTick();
    s_fade_start        = 0;
    s_sleep_req         = false;
    s_peripherals_gated = false;
    enable_wake_exti();
}

void power_mgmt_activity(void)
{
    s_last_active = HAL_GetTick();
    if (s_state != PM_ACTIVE) {
        s_sleep_req = false;
        s_state     = PM_WAKING;  /* resolved to PM_ACTIVE on next tick() */
    }
}

void power_mgmt_request_sleep(void)
{
    s_sleep_req = true;
}

PowerState power_mgmt_state(void)
{
    return s_state;
}

PowerState power_mgmt_tick(uint32_t now_ms)
{
    switch (s_state) {

    case PM_ACTIVE:
        if (s_sleep_req || (now_ms - s_last_active) >= PM_IDLE_TIMEOUT_MS) {
            s_sleep_req  = false;
            s_fade_start = now_ms;
            s_state      = PM_SLEEP_PENDING;
        }
        break;

    case PM_SLEEP_PENDING:
        if ((now_ms - s_fade_start) >= PM_FADE_MS) {
            gate_peripherals();
            enable_wake_exti();
            s_state = PM_SLEEPING;

            /* Feed IWDG: wakeup <10 ms << WATCHDOG_TIMEOUT_MS (50 ms) */
            IWDG->KR = 0xAAAAu;

            /* === STOP MODE ENTRY — execution resumes here on EXTI wake === */
            HAL_PWR_EnterSTOPMode(PWR_LOWPOWERREGULATOR_ON, PWR_STOPENTRY_WFI);
            /* === WAKEUP POINT (< 10 ms latency) === */

            restore_clocks();
            ungate_peripherals();
            disable_wake_exti();
            s_last_active = HAL_GetTick();
            s_state       = PM_ACTIVE;
        }
        break;

    case PM_SLEEPING:
        /* Unreachable: WFI is inline in PM_SLEEP_PENDING above */
        break;

    case PM_WAKING:
        /* Set by power_mgmt_activity() during SLEEP_PENDING/SLEEPING */
        ungate_peripherals();
        s_state = PM_ACTIVE;
        break;
    }

    return s_state;
}

/* Issue #467: battery low-voltage emergency sleep integration */
static uint32_t s_batt_critical_mv = 0u;

void power_mgmt_notify_battery(uint32_t vbat_mv)
{
    s_batt_critical_mv = vbat_mv;
    if (s_state == PM_SLEEPING || s_state == PM_SLEEP_PENDING) return;
    s_sleep_req  = true;
    s_fade_start = HAL_GetTick();
}
