//! WS2812 status LED driver for CMDR Joystick
//!
//! Provides a small helper around `ws2812_pio` to display firmware state on a
//! single WS2812 RGB LED using RP2040 PIO. The driver offers simple modes
//! (solid/flash) and a `SystemState` facade used by main.

use rp2040_hal::{
    gpio::AnyPin,
    pio::{PIOExt, StateMachineIndex, UninitStateMachine, PIO},
};
use smart_leds::{SmartLedsWrite, RGB8};
use ws2812_pio::Ws2812Direct;

const COLOR_OFF: RGB8 = RGB8 { r: 0, g: 0, b: 0 };
const COLOR_GREEN: RGB8 = RGB8 { r: 10, g: 7, b: 0 };
const COLOR_BLUE: RGB8 = RGB8 { r: 10, g: 4, b: 10 };
const COLOR_ORANGE: RGB8 = RGB8 { r: 5, g: 10, b: 0 };
const COLOR_RED: RGB8 = RGB8 { r: 20, g: 0, b: 0 };
const COLOR_PURPLE: RGB8 = RGB8 { r: 0, g: 10, b: 10 };

/// Status LED modes with clear semantics.
#[allow(dead_code)]
#[derive(PartialEq, Eq, Copy, Clone, Debug)]
pub enum StatusMode {
    Off = 0,
    Normal = 1,
    NormalFlash = 2,
    Activity = 3,
    ActivityFlash = 4,
    Idle = 5,
    Other = 6,
    OtherFlash = 7,
    Warning = 8,
    Error = 9,
    Bootloader = 10,
    Power = 11,
    Suspended = 12,
}

/// Aggregate system state for LED status indication.
#[derive(Clone, Copy)]
pub struct SystemState {
    pub usb_active: bool,
    pub usb_initialized: bool,
    pub usb_suspended: bool,
    pub idle_mode: bool,
    pub calibration_active: bool,
    pub throttle_hold_enable: bool,
    pub vt_enable: bool,
}

#[derive(Copy, Clone)]
enum LedEffect {
    Solid,
    Blink { period_ms: u32 },
    Heartbeat { period_ms: u32 },
}

#[derive(Copy, Clone)]
struct ModeDescriptor {
    color: RGB8,
    effect: LedEffect,
}

const HEARTBEAT_POWER_MS: u32 = 800;
const HEARTBEAT_IDLE_MS: u32 = 3200;
const HEARTBEAT_PAUSE_MS: u32 = 3000;

impl LedEffect {
    fn update_interval_ms(self) -> u32 {
        // Resolve the LED descriptor for the requested status mode.
        match self {
            LedEffect::Solid => 0,
            LedEffect::Blink { period_ms } => period_ms / 2,
            LedEffect::Heartbeat { .. } => 40,
        }
    }

    fn color_for(self, base: RGB8, elapsed_ms: u32) -> RGB8 {
        // Compute the base status mode given system state flags.
        match self {
            LedEffect::Solid => base,
            LedEffect::Blink { period_ms } => {
                if period_ms == 0 {
                    return base;
                }
                let half = (period_ms / 2).max(1);
                if elapsed_ms % period_ms < half {
                    base
                } else {
                    RGB8 { r: 0, g: 0, b: 0 }
                }
            }
            LedEffect::Heartbeat { period_ms } => {
                let period = period_ms.max(1);
                let cycle = period.saturating_add(HEARTBEAT_PAUSE_MS);
                let phase = elapsed_ms % cycle;
                if phase >= period {
                    return COLOR_OFF;
                }

                let half = (period / 2).max(1);
                let ramp = if phase < half {
                    ((phase * 255) / half) as u8
                } else {
                    (((period - phase) * 255) / half) as u8
                };
                scale_color(base, ramp)
            }
        }
    }
}

const fn descriptor_for(mode: StatusMode, base_mode: StatusMode) -> ModeDescriptor {
    match mode {
        StatusMode::Off => ModeDescriptor {
            color: COLOR_OFF,
            effect: LedEffect::Solid,
        },
        StatusMode::Normal => ModeDescriptor {
            color: COLOR_GREEN,
            effect: LedEffect::Solid,
        },
        StatusMode::NormalFlash => ModeDescriptor {
            color: COLOR_GREEN,
            effect: LedEffect::Blink { period_ms: 1000 },
        },
        StatusMode::Activity => ModeDescriptor {
            color: COLOR_BLUE,
            effect: LedEffect::Solid,
        },
        StatusMode::ActivityFlash => ModeDescriptor {
            color: COLOR_BLUE,
            effect: LedEffect::Blink { period_ms: 600 },
        },
        StatusMode::Idle => match base_mode {
            StatusMode::Activity | StatusMode::ActivityFlash => ModeDescriptor {
                color: COLOR_BLUE,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
            StatusMode::Other | StatusMode::OtherFlash => ModeDescriptor {
                color: COLOR_ORANGE,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
            StatusMode::NormalFlash | StatusMode::Normal => ModeDescriptor {
                color: COLOR_GREEN,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
            StatusMode::Warning => ModeDescriptor {
                color: COLOR_RED,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
            StatusMode::Error => ModeDescriptor {
                color: COLOR_RED,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
            _ => ModeDescriptor {
                color: COLOR_GREEN,
                effect: LedEffect::Heartbeat {
                    period_ms: HEARTBEAT_IDLE_MS,
                },
            },
        },
        StatusMode::Other => ModeDescriptor {
            color: COLOR_ORANGE,
            effect: LedEffect::Solid,
        },
        StatusMode::OtherFlash => ModeDescriptor {
            color: COLOR_ORANGE,
            effect: LedEffect::Blink { period_ms: 600 },
        },
        StatusMode::Warning => ModeDescriptor {
            color: COLOR_RED,
            effect: LedEffect::Blink { period_ms: 500 },
        },
        StatusMode::Error => ModeDescriptor {
            color: COLOR_RED,
            effect: LedEffect::Solid,
        },
        StatusMode::Bootloader => ModeDescriptor {
            color: COLOR_PURPLE,
            effect: LedEffect::Solid,
        },
        StatusMode::Power => ModeDescriptor {
            color: COLOR_GREEN,
            effect: LedEffect::Heartbeat {
                period_ms: HEARTBEAT_POWER_MS,
            },
        },
        StatusMode::Suspended => ModeDescriptor {
            color: COLOR_OFF,
            effect: LedEffect::Solid,
        },
    }
}

fn scale_color(base: RGB8, brightness: u8) -> RGB8 {
    if brightness == 0 {
        return COLOR_OFF;
    }

    let components = [base.r, base.g, base.b];
    let mut scaled = [0u8; 3];
    let mut remainders = [0u16; 3];
    let mut total_floor: u16 = 0;
    let mut total_base: u16 = 0;
    let brightness_u16 = brightness as u16;

    for (index, &component) in components.iter().enumerate() {
        total_base += component as u16;
        if component == 0 {
            continue;
        }

        let value = component as u32 * brightness as u32;
        let div = (value / 255) as u16;
        let rem = (value % 255) as u16;

        scaled[index] = div as u8;
        remainders[index] = rem;
        total_floor += div;
    }

    if total_base == 0 {
        return COLOR_OFF;
    }

    let mut target_total = ((total_base as u32 * brightness_u16 as u32) + 127) / 255;
    if target_total > total_base as u32 {
        target_total = total_base as u32;
    }

    let mut extra = target_total.saturating_sub(total_floor as u32) as u16;

    while extra > 0 {
        let mut best_index: Option<usize> = None;
        let mut best_remainder = 0u16;

        for idx in 0..components.len() {
            if components[idx] == 0 {
                continue;
            }
            if remainders[idx] == 0 {
                continue;
            }
            if scaled[idx] as u16 >= components[idx] as u16 {
                continue;
            }

            if remainders[idx] > best_remainder {
                best_remainder = remainders[idx];
                best_index = Some(idx);
            }
        }

        let Some(idx) = best_index else {
            break;
        };

        scaled[idx] += 1;
        remainders[idx] = 0;
        extra -= 1;
    }

    if extra > 0 {
        for idx in 0..components.len() {
            if extra == 0 {
                break;
            }
            if components[idx] == 0 {
                continue;
            }
            if scaled[idx] as u16 >= components[idx] as u16 {
                continue;
            }

            scaled[idx] += 1;
            extra -= 1;
        }
    }

    RGB8 {
        r: scaled[0],
        g: scaled[1],
        b: scaled[2],
    }
}

fn determine_base_mode(system_state: SystemState) -> StatusMode {
    // Compute the base status mode based on USB/calibration/idle state.
    if system_state.usb_suspended {
        StatusMode::Suspended
    } else if system_state.calibration_active {
        StatusMode::ActivityFlash
    } else if !system_state.usb_initialized {
        StatusMode::Power
    } else if !system_state.usb_active {
        StatusMode::NormalFlash
    } else if system_state.vt_enable {
        StatusMode::Activity
    } else if system_state.throttle_hold_enable {
        StatusMode::Other
    } else {
        StatusMode::Normal
    }
}

/// Status LED driver with self-contained state management.
pub struct StatusLed<P, SM, I>
where
    I: AnyPin<Function = P::PinFunction>,
    P: PIOExt,
    SM: StateMachineIndex,
{
    ws2812_direct: Ws2812Direct<P, SM, I>,
    current_mode: StatusMode,
    mode_started_at: Option<u32>,
    last_update_time: Option<u32>,
}

impl<P, SM, I> StatusLed<P, SM, I>
where
    I: AnyPin<Function = P::PinFunction>,
    P: PIOExt,
    SM: StateMachineIndex,
{
    /// Create a new StatusLed instance.
    ///
    /// # Arguments
    /// * `pin` - PIO pin for WS2812 LED
    /// * `pio` - PIO instance
    /// * `sm` - PIO state machine
    /// * `clock_freq` - PIO clock frequency
    pub fn new(
        pin: I,
        pio: &mut PIO<P>,
        sm: UninitStateMachine<(P, SM)>,
        clock_freq: fugit::HertzU32,
    ) -> Self {
        let ws2812_direct = Ws2812Direct::new(pin, pio, sm, clock_freq);
        let mut status = Self {
            ws2812_direct,
            current_mode: StatusMode::Off,
            mode_started_at: None,
            last_update_time: None,
        };
        status.write_color(COLOR_OFF);
        status
    }

    /// Update LED based on system state and current time (ms).
    pub fn update_from_system_state(&mut self, system_state: SystemState, current_time: u32) {
        let base_mode = determine_base_mode(system_state);
        let idle_overlay = system_state.idle_mode
            && system_state.usb_initialized
            && base_mode != StatusMode::Off
            && base_mode != StatusMode::Power;
        let desired_mode = if idle_overlay {
            (StatusMode::Idle, base_mode)
        } else {
            (base_mode, base_mode)
        };

        self.set_mode(desired_mode.0, desired_mode.1, current_time);
    }

    /// Set LED mode directly (explicit override).
    pub fn set_mode(&mut self, mode: StatusMode, base_mode: StatusMode, current_time: u32) {
        // Force update if mode changed
        let force_update = mode != self.current_mode;
        self.current_mode = mode;
        let descriptor = descriptor_for(mode, base_mode);
        if force_update {
            let start_time = match descriptor.effect {
                LedEffect::Heartbeat { period_ms } => {
                    let half = (period_ms / 2).max(1);
                    current_time.saturating_sub(half)
                }
                _ => current_time,
            };
            self.mode_started_at = Some(start_time);
            self.last_update_time = None;
        }

        self.update_display(current_time, force_update, base_mode);
    }

    /// Periodic update for flashing behavior at roughly ~500 ms intervals.
    pub fn update_display(&mut self, current_time: u32, force_update: bool, base_mode: StatusMode) {
        let descriptor = descriptor_for(self.current_mode, base_mode);
        let interval = descriptor.effect.update_interval_ms();

        if !force_update {
            if interval == 0 {
                return;
            }
            if let Some(last) = self.last_update_time {
                if current_time.saturating_sub(last) < interval {
                    return;
                }
            }
        }

        let elapsed = self
            .mode_started_at
            .map(|start| current_time.saturating_sub(start))
            .unwrap_or(0);
        let color = descriptor.effect.color_for(descriptor.color, elapsed);

        self.write_color(color);
        self.last_update_time = Some(current_time);
    }

    /// Write a single color to the LED.
    fn write_color(&mut self, color: RGB8) {
        // Push the color to the WS2812 LED, ignoring transient IO errors.
        let _ = self.ws2812_direct.write([color].iter().copied());
    }
}

impl<P, SM, I> StatusLed<P, SM, I>
where
    I: AnyPin<Function = P::PinFunction>,
    P: PIOExt,
    SM: StateMachineIndex,
{
    /// Legacy interface for compatibility – direct mode update.
    pub fn update(&mut self, mode: StatusMode) {
        // Use a dummy time for immediate updates
        self.set_mode(mode, mode, 0);
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;

    #[test]
    fn idle_mode_uses_base_color_with_heartbeat() {
        // Idle state should inherit the base color while enforcing a heartbeat pattern.
        let state = SystemState {
            usb_active: false,
            usb_initialized: true,
            usb_suspended: false,
            idle_mode: true,
            calibration_active: false,
            throttle_hold_enable: false,
            vt_enable: false,
        };

        let base = determine_base_mode(state);
        assert_eq!(base, StatusMode::NormalFlash);

        let descriptor = descriptor_for(StatusMode::Idle, base);
        let LedEffect::Heartbeat { .. } = descriptor.effect else {
            panic!("Idle should always use heartbeat effect");
        };
        assert_eq!(descriptor.color, COLOR_GREEN);
    }

    #[test]
    fn power_mode_uses_fast_heartbeat() {
        // Power mode descriptor should use the fast heartbeat cadence and green.
        let descriptor = descriptor_for(StatusMode::Power, StatusMode::Normal);
        if let LedEffect::Heartbeat { period_ms } = descriptor.effect {
            assert_eq!(period_ms, HEARTBEAT_POWER_MS);
            assert_eq!(descriptor.color, COLOR_GREEN);
        } else {
            panic!("Power mode must use heartbeat effect");
        }
    }

    #[test]
    fn calibration_has_priority_over_idle() {
        // Calibration activity should override idle when both flags are set.
        let state = SystemState {
            usb_active: true,
            usb_initialized: true,
            usb_suspended: false,
            idle_mode: true,
            calibration_active: true,
            throttle_hold_enable: false,
            vt_enable: false,
        };

        let base = determine_base_mode(state);
        assert_eq!(base, StatusMode::ActivityFlash);
    }

    #[test]
    fn heartbeat_effect_fades() {
        // Heartbeat should ramp up and down around the target color.
        let base = StatusMode::Normal;
        let descriptor = descriptor_for(StatusMode::Idle, base);
        let LedEffect::Heartbeat { period_ms } = descriptor.effect else {
            panic!("Idle should use heartbeat effect");
        };

        let start = descriptor.effect.color_for(descriptor.color, 0);
        let quarter = descriptor.effect.color_for(descriptor.color, period_ms / 4);
        let half = descriptor.effect.color_for(descriptor.color, period_ms / 2);
        let end = descriptor.effect.color_for(descriptor.color, period_ms);

        assert_eq!(start.g, 0);
        assert!(quarter.g > start.g);
        assert_eq!(half, descriptor.color);
        assert_eq!(end.g, 0);
    }

    #[test]
    fn heartbeat_pause_keeps_led_off() {
        // Added pause should hold the LED dark between breaths while allowing a clean restart.
        let base = StatusMode::Normal;
        let descriptor = descriptor_for(StatusMode::Idle, base);
        let LedEffect::Heartbeat { period_ms } = descriptor.effect else {
            panic!("Idle should use heartbeat effect");
        };

        let pause_sample = descriptor
            .effect
            .color_for(descriptor.color, period_ms + HEARTBEAT_PAUSE_MS / 2);
        assert_eq!(pause_sample, COLOR_OFF);

        let cycle = period_ms + HEARTBEAT_PAUSE_MS;
        let restart = descriptor
            .effect
            .color_for(descriptor.color, cycle + period_ms / 4);
        assert!(restart.g > 0);
    }

    #[test]
    fn blink_effect_toggles() {
        // Blink descriptor should alternate between the color and off state.
        let descriptor = descriptor_for(StatusMode::NormalFlash, StatusMode::NormalFlash);
        let LedEffect::Blink { period_ms } = descriptor.effect else {
            panic!("NormalFlash should use blink effect");
        };

        let on = descriptor.effect.color_for(descriptor.color, 0);
        let off = descriptor.effect.color_for(descriptor.color, period_ms / 2);
        assert_eq!(on, descriptor.color);
        assert_eq!(off, COLOR_OFF);
    }

    #[test]
    fn low_brightness_preserves_color_mix() {
        // Low brightness scaling must not drop any non-zero channel from the blend.
        let base = COLOR_ORANGE;
        let dimmed = scale_color(base, 50);
        assert!(dimmed.r > 0);
        assert!(dimmed.g > 0);
        assert_eq!(dimmed.b, 0);
        assert!(dimmed.r <= base.r);
        assert!(dimmed.g <= base.g);
        assert_eq!(dimmed.r as u16 * base.g as u16, dimmed.g as u16 * base.r as u16);
    }

    #[test]
    fn zero_brightness_turns_off_led() {
        // Zero brightness should fully blank the LED regardless of the base color.
        let base = COLOR_BLUE;
        let off = scale_color(base, 0);
        assert_eq!(off, COLOR_OFF);
    }

    #[test]
    fn determine_base_mode_before_usb() {
        // Before USB comes up the controller should stay in Power mode.
        let state = SystemState {
            usb_active: false,
            usb_initialized: false,
            usb_suspended: false,
            idle_mode: false,
            calibration_active: false,
            throttle_hold_enable: false,
            vt_enable: false,
        };

        assert_eq!(determine_base_mode(state), StatusMode::Power);
    }

    #[test]
    fn usb_suspend_takes_priority() {
        // USB suspend should trump other status priorities.
        let state = SystemState {
            usb_active: true,
            usb_initialized: true,
            usb_suspended: true,
            idle_mode: false,
            calibration_active: true, // Even with calibration active
            throttle_hold_enable: false,
            vt_enable: false,
        };

        assert_eq!(determine_base_mode(state), StatusMode::Suspended);
    }
}
