//! WS2812 status LED driver adapted from the joystick firmware.

use rp2040_hal::{
    gpio::AnyPin,
    pio::{PIO, PIOExt, StateMachineIndex, UninitStateMachine},
};
use smart_leds::{RGB8, SmartLedsWrite};
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 };

#[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,
}

#[derive(Clone, Copy, Debug, Default)]
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,
    pub caps_lock_active: bool,
    pub sticky_armed: bool,
    pub sticky_latched: bool,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct StatusSummary {
    pub caps_lock_active: bool,
    pub sticky_armed: bool,
    pub sticky_latched: bool,
    pub usb_initialized: bool,
    pub usb_active: bool,
    pub usb_suspended: bool,
    pub idle_mode: bool,
}

impl StatusSummary {
    pub const fn new(
        caps_lock_active: bool,
        sticky_armed: bool,
        sticky_latched: bool,
        usb_initialized: bool,
        usb_active: bool,
        usb_suspended: bool,
        idle_mode: bool,
    ) -> Self {
        Self {
            caps_lock_active,
            sticky_armed,
            sticky_latched,
            usb_initialized,
            usb_active,
            usb_suspended,
            idle_mode,
        }
    }

    pub fn to_system_state(self) -> SystemState {
        SystemState {
            usb_active: self.usb_active,
            usb_initialized: self.usb_initialized,
            usb_suspended: self.usb_suspended,
            idle_mode: self.idle_mode,
            calibration_active: false,
            throttle_hold_enable: false,
            vt_enable: false,
            caps_lock_active: self.caps_lock_active,
            sticky_armed: self.sticky_armed,
            sticky_latched: self.sticky_latched,
        }
    }
}

#[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;

impl LedEffect {
    fn update_interval_ms(self) -> u32 {
        match self {
            LedEffect::Solid => 0,
            LedEffect::Blink { period_ms } => period_ms / 2,
            LedEffect::Heartbeat { .. } => 10,
        }
    }

    fn color_for(self, base: RGB8, elapsed_ms: u32) -> RGB8 {
        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 {
                    COLOR_OFF
                }
            }
            LedEffect::Heartbeat { period_ms } => {
                let period = period_ms.max(1);
                let phase = elapsed_ms % period;
                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 | 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 {
    let scale = brightness as u16;
    RGB8 {
        r: ((base.r as u16 * scale) / 255) as u8,
        g: ((base.g as u16 * scale) / 255) as u8,
        b: ((base.b as u16 * scale) / 255) as u8,
    }
}

fn determine_base_mode(system_state: SystemState) -> StatusMode {
    if system_state.usb_suspended {
        StatusMode::Suspended
    } else if system_state.caps_lock_active {
        StatusMode::Warning
    } else if system_state.sticky_latched {
        StatusMode::ActivityFlash
    } else if system_state.sticky_armed {
        StatusMode::Activity
    } 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
    }
}

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,
{
    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
    }

    pub fn update_from_system_state(&mut self, system_state: SystemState, current_time: u32) {
        let base_mode = determine_base_mode(system_state);
        let desired_mode = if system_state.idle_mode
            && system_state.usb_initialized
            && base_mode != StatusMode::Off
            && base_mode != StatusMode::Power
        {
            (StatusMode::Idle, base_mode)
        } else {
            (base_mode, base_mode)
        };

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

    pub fn set_mode(&mut self, mode: StatusMode, base_mode: StatusMode, current_time: u32) {
        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);
    }

    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
                && 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);
    }

    fn write_color(&mut self, color: RGB8) {
        let _ = self.ws2812_direct.write([color].iter().copied());
    }

    pub fn apply_summary(&mut self, summary: StatusSummary, current_time: u32) {
        self.update_from_system_state(summary.to_system_state(), current_time);
    }
}

impl<P, SM, I> StatusLed<P, SM, I>
where
    I: AnyPin<Function = P::PinFunction>,
    P: PIOExt,
    SM: StateMachineIndex,
{
    pub fn update(&mut self, mode: StatusMode) {
        self.set_mode(mode, mode, 0);
    }
}

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

    #[test]
    fn caps_lock_overrides_with_warning() {
        let state = SystemState {
            caps_lock_active: true,
            ..SystemState::default()
        };
        assert_eq!(determine_base_mode(state), StatusMode::Warning);
    }

    #[test]
    fn sticky_states_map_to_activity_modes() {
        let state = SystemState {
            sticky_latched: true,
            ..SystemState::default()
        };
        assert_eq!(determine_base_mode(state), StatusMode::ActivityFlash);

        let state = SystemState {
            sticky_armed: true,
            ..SystemState::default()
        };
        assert_eq!(determine_base_mode(state), StatusMode::Activity);
    }

    #[test]
    fn usb_not_active_flashes_normal() {
        let state = SystemState {
            usb_initialized: true,
            usb_active: false,
            ..SystemState::default()
        };
        assert_eq!(determine_base_mode(state), StatusMode::NormalFlash);
    }

    #[test]
    fn usb_suspend_takes_priority() {
        let state = SystemState {
            usb_suspended: true,
            caps_lock_active: true, // Even with caps lock active
            sticky_latched: true,   // Even with sticky latched
            ..SystemState::default()
        };
        assert_eq!(determine_base_mode(state), StatusMode::Suspended);
    }
}