//! Keyboard state management and HID report generation.

use crate::hardware;
use crate::{NUMBER_OF_KEYS, KeyMatrix, KeyReport};
use crate::layout;
use crate::status::StatusSummary;
use usbd_human_interface_device::page::Keyboard;
use usb_device::device::UsbDeviceState;

/// Tracks USB lifecycle state (suspend/idle/activity) for the keyboard HID device.
pub struct UsbState {
    pub initialized: bool,
    pub active: bool,
    pub suspended: bool,
    pub wake_on_input: bool,
    pub idle_mode: bool,
    activity: bool,
    activity_elapsed_ms: u32,
}

impl UsbState {
    pub const fn new() -> Self {
        Self {
            initialized: false,
            active: false,
            suspended: false,
            wake_on_input: false,
            idle_mode: false,
            activity: false,
            activity_elapsed_ms: 0,
        }
    }

    pub fn on_poll(&mut self) {
        if !self.initialized {
            self.initialized = true;
        }
        if !self.active {
            self.mark_activity();
        }
        self.active = true;
    }

    pub fn mark_activity(&mut self) {
        self.activity = true;
        self.activity_elapsed_ms = 0;
        self.idle_mode = false;
    }

    pub fn handle_input_activity(&mut self) {
        self.mark_activity();
        if self.suspended && self.wake_on_input {
            self.wake_on_input = false;
        }
    }

    pub fn on_suspend_change(&mut self, state: UsbDeviceState) {
        let was_suspended = self.suspended;
        self.suspended = state == UsbDeviceState::Suspend;

        match (was_suspended, self.suspended) {
            (true, false) => {
                self.mark_activity();
                self.wake_on_input = false;
            }
            (false, true) => {
                self.idle_mode = true;
                self.activity = false;
                self.wake_on_input = true;
            }
            _ => {}
        }
    }

    pub fn advance_idle_timer(&mut self, interval_ms: u32) {
        if !self.activity {
            return;
        }
        self.activity_elapsed_ms = self.activity_elapsed_ms.saturating_add(interval_ms);
        if self.activity_elapsed_ms >= hardware::timers::IDLE_TIMEOUT_MS {
            self.activity = false;
            self.activity_elapsed_ms = 0;
            self.idle_mode = true;
        }
    }

    pub fn acknowledge_report(&mut self) {}
}

/// Captures per-key state transitions and the function layer active when it was pressed.
#[derive(Copy, Clone, Default, Debug, PartialEq, Eq)]
pub struct KeyboardButton {
    pub pressed: bool,
    pub previous_pressed: bool,
    pub fn_mode: u8,
}

/// Discrete states for the sticky modifier state machine.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub enum StickyState {
    Inactive,
    Armed,
    Latched,
}

impl StickyState {
    pub fn is_armed(self) -> bool {
        matches!(self, StickyState::Armed)
    }

    pub fn is_latched(self) -> bool {
        matches!(self, StickyState::Latched)
    }
}

/// Manages keyboard-wide state, layer selection, and HID report composition.
pub struct KeyboardState {
    buttons: [KeyboardButton; NUMBER_OF_KEYS],
    sticky_state: StickyState,
    sticky_key: Keyboard,
    caps_lock_active: bool,
    started: bool,
    usb: UsbState,
}

impl KeyboardState {
    pub fn new() -> Self {
        // Initialise button, sticky, and host communication state.
        Self {
            buttons: [KeyboardButton::default(); NUMBER_OF_KEYS],
            sticky_state: StickyState::Inactive,
            sticky_key: Keyboard::NoEventIndicated,
            caps_lock_active: false,
            started: false,
            usb: UsbState::new(),
        }
    }

    pub fn process_scan(
        &mut self,
        pressed_keys: KeyMatrix,
    ) -> KeyReport {
        // Update each button from the latest scan and build the keyboard report.
        let fn_mode = Self::fn_mode(&pressed_keys);

        for (index, pressed) in pressed_keys.iter().enumerate() {
            self.buttons[index].pressed = *pressed;
        }

        self.build_report(fn_mode)
    }

    pub fn update_caps_lock(&mut self, active: bool) {
        // Track the host LED state to drive the status indicator.
        self.caps_lock_active = active;
    }

    pub fn caps_lock_active(&self) -> bool {
        // Report whether the Caps Lock LED is currently active.
        self.caps_lock_active
    }

    pub fn mark_started(&mut self) {
        // Note that the HID interface has successfully exchanged reports.
        self.started = true;
        self.usb.mark_activity();
        self.usb.active = true;
        self.usb.initialized = true;
    }

    pub fn mark_stopped(&mut self) {
        // Reset the flag when USB communication fails.
        self.started = false;
        self.usb.active = false;
        self.usb.initialized = false;
        self.usb.activity = false;
        self.usb.idle_mode = false;
    }

    pub fn started(&self) -> bool {
        // Expose whether USB activity has been observed.
        self.started
    }

    pub fn sticky_state(&self) -> StickyState {
        // Current state of the sticky modifier toggle.
        self.sticky_state
    }

    fn toggle_sticky_state(&mut self) {
        // Advance through inactive → armed → latched lifecycle and clear when toggled off.
        self.sticky_state = match self.sticky_state {
            StickyState::Inactive => StickyState::Armed,
            StickyState::Armed | StickyState::Latched => {
                self.sticky_key = Keyboard::NoEventIndicated;
                StickyState::Inactive
            }
        };
    }

    pub fn status_summary(&self) -> StatusSummary {
        // Produce a condensed summary consumed by the status LED driver.
        let usb_active = self.usb.active && !self.usb.idle_mode;
        StatusSummary::new(
            self.caps_lock_active,
            matches!(self.sticky_state, StickyState::Armed),
            matches!(self.sticky_state, StickyState::Latched),
            self.usb.initialized,
            usb_active,
            self.usb.suspended,
            self.usb.idle_mode,
        )
    }

    fn build_report(&mut self, fn_mode: u8) -> KeyReport {
        // Translate layer-aware button state into the NKRO HID report payload.
        let mut report = [Keyboard::NoEventIndicated; NUMBER_OF_KEYS];
        let mut sticky_toggle_requested = false;

        for (index, button) in self.buttons.iter_mut().enumerate() {
            let changed = button.pressed != button.previous_pressed;
            let just_pressed = changed && button.pressed;

            if just_pressed {
                button.fn_mode = fn_mode;

                match (index as u8, fn_mode) {
                    (idx, layer) if idx == layout::STICKY_BUTTON[0] && layer == layout::STICKY_BUTTON[1] => {
                        sticky_toggle_requested = true;
                    }
                    (idx, layer) if idx == layout::OS_LOCK_BUTTON[0] && layer == layout::OS_LOCK_BUTTON[1] => {
                        report[36] = layout::OS_LOCK_BUTTON_KEYS[0];
                        report[37] = layout::OS_LOCK_BUTTON_KEYS[1];
                    }
                    _ => {}
                }
            }

            let layer_key = layout::MAP[button.fn_mode as usize][index];
            if layer_key == Keyboard::NoEventIndicated {
                button.previous_pressed = button.pressed;
                continue;
            }

            if self.sticky_state == StickyState::Armed && button.pressed {
                self.sticky_key = layer_key;
                self.sticky_state = StickyState::Latched;
            }

            if button.pressed {
                report[index] = layer_key;
            }

            button.previous_pressed = button.pressed;
        }

        if sticky_toggle_requested {
            self.toggle_sticky_state();
        }

        const STICKY_REPORT_INDEX: usize = 46;
        report[STICKY_REPORT_INDEX] = self.sticky_key;

        report
    }

    fn fn_mode(pressed_keys: &KeyMatrix) -> u8 {
        // Count the active FN buttons and clamp to the highest supported layer.
        let active_fn_keys = layout::FN_BUTTONS
            .iter()
            .filter(|key_index| pressed_keys[**key_index as usize])
            .count() as u8;

        active_fn_keys.min(2)
    }

    pub fn usb_state(&mut self) -> &mut UsbState {
        &mut self.usb
    }

    pub fn usb(&self) -> &UsbState {
        &self.usb
    }
}

impl Default for KeyboardState {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn fn_mode_caps_at_two() {
        // Ensure the layer helper never exceeds the maximum FN layer value.
        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[layout::FN_BUTTONS[2] as usize] = true;

        assert_eq!(KeyboardState::fn_mode(&pressed), 2);
    }

    #[test]
    fn sticky_button_transitions_between_states() {
        // Sticky button chord should arm, latch on next key, and clear when pressed again.
        let mut state = KeyboardState::new();

        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[layout::STICKY_BUTTON[0] as usize] = true;

        state.process_scan(pressed);
        assert_eq!(state.sticky_state(), StickyState::Armed);

        // Press another key to latch sticky
        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[0] = true;
        state.process_scan(pressed);
        assert_eq!(state.sticky_state(), StickyState::Latched);

        // Press sticky again to clear
        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[layout::STICKY_BUTTON[0] as usize] = true;
        state.process_scan(pressed);
        assert_eq!(state.sticky_state(), StickyState::Inactive);
    }

    #[test]
    fn status_summary_reflects_keyboard_state() {
        // Status summary must mirror the internal keyboard and USB flags.
        let mut state = KeyboardState::new();
        state.update_caps_lock(true);
        state.mark_started();
        let summary = state.status_summary();
        assert!(summary.caps_lock_active);
        assert!(summary.usb_active);
        assert!(summary.usb_initialized);
        assert!(!summary.sticky_armed);
        assert!(!summary.sticky_latched);
    }

    #[test]
    fn sticky_key_is_exposed_in_reports_after_latch() {
        // Ensure the latched sticky modifier key stays in the HID report stream until cleared.
        let mut state = KeyboardState::new();

        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[layout::STICKY_BUTTON[0] as usize] = true;
        state.process_scan(pressed);

        let mut pressed = [false; NUMBER_OF_KEYS];
        pressed[layout::FN_BUTTONS[0] as usize] = true;
        pressed[layout::FN_BUTTONS[1] as usize] = true;
        pressed[0] = true;
        let report = state.process_scan(pressed);

        assert_eq!(state.sticky_state(), StickyState::Latched);
        assert_eq!(report[0], layout::MAP[2][0]);
        assert_eq!(report[46], layout::MAP[2][0]);

        let pressed = [false; NUMBER_OF_KEYS];
        let report = state.process_scan(pressed);

        assert_eq!(report[46], layout::MAP[2][0]);
    }
}