use core::num::NonZeroU32;
use std::collections::BTreeSet;

use serde::{Deserialize, Serialize};

use crate::input::InputSnapshot;
use crate::widget_id::WidgetId;

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct FocusScopeId(NonZeroU32);

impl FocusScopeId {
    pub const ROOT: Self = match NonZeroU32::new(1) {
        Some(n) => Self(n),
        None => panic!("ROOT scope id is 1"),
    };

    #[must_use]
    pub const fn get(self) -> NonZeroU32 {
        self.0
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FocusScopeKind {
    Window,
    Modal,
    Roving,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct FocusScope {
    id: FocusScopeId,
    kind: FocusScopeKind,
}

const FIRST_CHILD_SCOPE: NonZeroU32 = match NonZeroU32::new(2) {
    Some(n) => n,
    None => panic!("first child scope id is 2"),
};

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum InputModality {
    Keyboard,
    Pointer,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RovingDirection {
    Next,
    Prev,
    First,
    Last,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum FocusRequest {
    Set(WidgetId),
    Clear,
    Advance,
    Retreat,
    Rove(RovingDirection),
}

#[derive(Clone, Debug, PartialEq)]
pub struct FocusManager {
    focused: Option<WidgetId>,
    modality: InputModality,
    scopes: Vec<FocusScope>,
    next_scope: NonZeroU32,
    tab_stops: Vec<(WidgetId, FocusScopeId)>,
    tab_stop_ids: BTreeSet<WidgetId>,
    focusable: BTreeSet<WidgetId>,
    text_inputs: BTreeSet<WidgetId>,
    request: Option<FocusRequest>,
}

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

impl FocusManager {
    #[must_use]
    pub fn new() -> Self {
        let root = FocusScope {
            id: FocusScopeId::ROOT,
            kind: FocusScopeKind::Window,
        };
        Self {
            focused: None,
            modality: InputModality::Pointer,
            scopes: vec![root],
            next_scope: FIRST_CHILD_SCOPE,
            tab_stops: Vec::new(),
            tab_stop_ids: BTreeSet::new(),
            focusable: BTreeSet::new(),
            text_inputs: BTreeSet::new(),
            request: None,
        }
    }

    #[must_use]
    pub fn focused(&self) -> Option<WidgetId> {
        self.focused
    }

    #[must_use]
    pub fn focus_visible(&self) -> bool {
        matches!(self.modality, InputModality::Keyboard)
    }

    pub fn observe_input(&mut self, input: &InputSnapshot) {
        if !input.keys_pressed.is_empty() {
            self.modality = InputModality::Keyboard;
        } else if !input.buttons_pressed.is_empty() || !input.buttons_released.is_empty() {
            self.modality = InputModality::Pointer;
        }
    }

    #[must_use]
    pub fn current_scope(&self) -> FocusScopeId {
        self.scopes.last().map_or(FocusScopeId::ROOT, |s| s.id)
    }

    pub fn push_scope(&mut self, kind: FocusScopeKind) -> FocusScopeId {
        let id = FocusScopeId(self.next_scope);
        self.next_scope = match self.next_scope.checked_add(1) {
            Some(n) => n,
            None => panic!("focus scope id space exhausted"),
        };
        self.scopes.push(FocusScope { id, kind });
        id
    }

    pub fn pop_scope(&mut self) -> Option<FocusScopeId> {
        if self.scopes.len() <= 1 {
            return None;
        }
        self.scopes.pop().map(|s| s.id)
    }

    pub fn register_tab_stop(&mut self, id: WidgetId) {
        if !self.tab_stop_ids.insert(id) {
            return;
        }
        self.focusable.insert(id);
        let scope = self.current_scope();
        self.tab_stops.push((id, scope));
    }

    pub fn register_focusable(&mut self, id: WidgetId) {
        self.focusable.insert(id);
    }

    pub fn register_text_input(&mut self, id: WidgetId) {
        self.text_inputs.insert(id);
    }

    #[must_use]
    pub fn is_text_input_focused(&self) -> bool {
        self.focused
            .is_some_and(|id| self.text_inputs.contains(&id))
    }

    #[must_use]
    pub fn tab_stops(&self) -> &[(WidgetId, FocusScopeId)] {
        &self.tab_stops
    }

    pub fn focusable_ids(&self) -> impl Iterator<Item = WidgetId> + '_ {
        self.focusable.iter().copied()
    }

    pub fn request(&mut self, request: FocusRequest) {
        self.request = Some(request);
    }

    pub fn request_focus(&mut self, id: WidgetId) {
        self.request(FocusRequest::Set(id));
    }

    pub fn begin_frame(&mut self) {
        self.tab_stops.clear();
        self.tab_stop_ids.clear();
        self.focusable.clear();
        self.text_inputs.clear();
        self.scopes.truncate(1);
        self.next_scope = FIRST_CHILD_SCOPE;
    }

    pub fn end_frame(&mut self) {
        let request = self.request.take();
        let next = match request {
            None => self.focused,
            Some(FocusRequest::Set(id)) => Some(id),
            Some(FocusRequest::Clear) => None,
            Some(FocusRequest::Advance) => self.advance_within_scope(false),
            Some(FocusRequest::Retreat) => self.advance_within_scope(true),
            Some(FocusRequest::Rove(direction)) => self.rove_within_roving_scope(direction),
        };
        self.focused = self.validate_focus(next);
    }

    fn validate_focus(&self, candidate: Option<WidgetId>) -> Option<WidgetId> {
        candidate.filter(|id| self.focusable.contains(id))
    }

    fn nearest_scope(&self, kind: FocusScopeKind) -> Option<FocusScopeId> {
        self.scopes
            .iter()
            .rev()
            .find(|s| s.kind == kind)
            .map(|s| s.id)
    }

    fn stops_for_scope(&self, scope: FocusScopeId) -> Vec<WidgetId> {
        self.tab_stops
            .iter()
            .filter(|(_, s)| *s == scope)
            .map(|(id, _)| *id)
            .collect()
    }

    fn advance_within_scope(&self, reverse: bool) -> Option<WidgetId> {
        let trap = self
            .nearest_scope(FocusScopeKind::Modal)
            .unwrap_or(FocusScopeId::ROOT);
        let stops = self.stops_for_scope(trap);
        cycle(&stops, self.focused, reverse)
    }

    fn rove_within_roving_scope(&self, direction: RovingDirection) -> Option<WidgetId> {
        let scope = self.nearest_scope(FocusScopeKind::Roving)?;
        let stops = self.stops_for_scope(scope);
        match direction {
            RovingDirection::Next => cycle(&stops, self.focused, false),
            RovingDirection::Prev => cycle(&stops, self.focused, true),
            RovingDirection::First => stops.first().copied(),
            RovingDirection::Last => stops.last().copied(),
        }
    }
}

fn cycle(stops: &[WidgetId], current: Option<WidgetId>, reverse: bool) -> Option<WidgetId> {
    if stops.is_empty() {
        return None;
    }
    let len = stops.len();
    let next_index = match current.and_then(|id| stops.iter().position(|s| *s == id)) {
        None => {
            if reverse {
                len - 1
            } else {
                0
            }
        }
        Some(idx) => {
            if reverse {
                (idx + len - 1) % len
            } else {
                (idx + 1) % len
            }
        }
    };
    stops.get(next_index).copied()
}

#[cfg(test)]
mod tests {
    use super::{FocusManager, FocusRequest, FocusScopeKind, RovingDirection};
    use crate::input::{
        FrameInstant, InputSnapshot, KeyChar, KeyCode, KeyEvent, ModifierMask, PointerButton,
        PointerButtonMask,
    };
    use crate::widget_id::{WidgetId, WidgetKey};

    fn id(key: &'static str) -> WidgetId {
        WidgetId::ROOT.child(WidgetKey::new(key))
    }

    fn run_frame<F: FnOnce(&mut FocusManager)>(focus: &mut FocusManager, build: F) {
        focus.begin_frame();
        build(focus);
        focus.end_frame();
    }

    fn key_input() -> InputSnapshot {
        let mut input = InputSnapshot::idle(FrameInstant::ZERO);
        input.keys_pressed.push(KeyEvent::new(
            KeyCode::Char(KeyChar::from_char('a')),
            ModifierMask::NONE,
        ));
        input
    }

    fn pointer_button_input() -> InputSnapshot {
        let mut input = InputSnapshot::idle(FrameInstant::ZERO);
        input.buttons_pressed = PointerButtonMask::just(PointerButton::Primary);
        input
    }

    #[test]
    fn observe_input_flips_to_keyboard_on_keypress() {
        let mut focus = FocusManager::new();
        focus.observe_input(&pointer_button_input());
        assert!(!focus.focus_visible());
        focus.observe_input(&key_input());
        assert!(focus.focus_visible());
    }

    #[test]
    fn observe_input_flips_back_to_pointer_on_button() {
        let mut focus = FocusManager::new();
        focus.observe_input(&key_input());
        assert!(focus.focus_visible());
        focus.observe_input(&pointer_button_input());
        assert!(!focus.focus_visible());
    }

    #[test]
    fn observe_input_idle_keeps_modality() {
        let mut focus = FocusManager::new();
        focus.observe_input(&key_input());
        assert!(focus.focus_visible());
        focus.observe_input(&InputSnapshot::idle(FrameInstant::ZERO));
        assert!(focus.focus_visible(), "idle frames must not flip modality");
    }

    #[test]
    fn advance_walks_tab_stops_in_order() {
        let mut focus = FocusManager::new();
        let a = id("a");
        let b = id("b");
        let c = id("c");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.register_tab_stop(c);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(a));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.register_tab_stop(c);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(b));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.register_tab_stop(c);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(c));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.register_tab_stop(c);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(a));
    }

    #[test]
    fn retreat_walks_in_reverse() {
        let mut focus = FocusManager::new();
        let a = id("a");
        let b = id("b");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request_focus(a);
        });
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request(FocusRequest::Retreat);
        });
        assert_eq!(focus.focused(), Some(b));
    }

    #[test]
    fn modal_trap_restricts_tab_traversal() {
        let mut focus = FocusManager::new();
        let outer = id("outer");
        let modal_a = id("modal_a");
        let modal_b = id("modal_b");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(outer);
            f.push_scope(FocusScopeKind::Modal);
            f.register_tab_stop(modal_a);
            f.register_tab_stop(modal_b);
            f.request_focus(modal_a);
        });
        assert_eq!(focus.focused(), Some(modal_a));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(outer);
            f.push_scope(FocusScopeKind::Modal);
            f.register_tab_stop(modal_a);
            f.register_tab_stop(modal_b);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(modal_b));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(outer);
            f.push_scope(FocusScopeKind::Modal);
            f.register_tab_stop(modal_a);
            f.register_tab_stop(modal_b);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(modal_a));
    }

    #[test]
    fn roving_within_composite_widget() {
        let mut focus = FocusManager::new();
        let toolbar_a = id("tool_a");
        let toolbar_b = id("tool_b");
        let toolbar_c = id("tool_c");
        run_frame(&mut focus, |f| {
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(toolbar_a);
            f.register_tab_stop(toolbar_b);
            f.register_tab_stop(toolbar_c);
            f.request_focus(toolbar_b);
        });
        run_frame(&mut focus, |f| {
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(toolbar_a);
            f.register_tab_stop(toolbar_b);
            f.register_tab_stop(toolbar_c);
            f.request(FocusRequest::Rove(RovingDirection::Next));
        });
        assert_eq!(focus.focused(), Some(toolbar_c));
        run_frame(&mut focus, |f| {
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(toolbar_a);
            f.register_tab_stop(toolbar_b);
            f.register_tab_stop(toolbar_c);
            f.request(FocusRequest::Rove(RovingDirection::First));
        });
        assert_eq!(focus.focused(), Some(toolbar_a));
    }

    #[test]
    fn focus_drops_when_widget_disappears() {
        let mut focus = FocusManager::new();
        let a = id("a");
        let b = id("b");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request_focus(b);
        });
        assert_eq!(focus.focused(), Some(b));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
        });
        assert_eq!(focus.focused(), None);
    }

    #[test]
    fn pop_scope_does_not_unseat_root() {
        let mut focus = FocusManager::new();
        assert!(focus.pop_scope().is_none());
        let inner = focus.push_scope(FocusScopeKind::Modal);
        assert_eq!(focus.pop_scope(), Some(inner));
        assert!(focus.pop_scope().is_none());
    }

    #[test]
    fn tab_skips_roving_stops_without_modal() {
        let mut focus = FocusManager::new();
        let host = id("toolbar");
        let rove_a = id("rove_a");
        let rove_b = id("rove_b");
        let trailing = id("after");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(host);
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(rove_a);
            f.register_tab_stop(rove_b);
            f.pop_scope();
            f.register_tab_stop(trailing);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(host));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(host);
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(rove_a);
            f.register_tab_stop(rove_b);
            f.pop_scope();
            f.register_tab_stop(trailing);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(trailing));
    }

    #[test]
    fn rove_returns_none_outside_roving_scope() {
        let mut focus = FocusManager::new();
        let a = id("a");
        let b = id("b");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request_focus(a);
        });
        assert_eq!(focus.focused(), Some(a));
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request(FocusRequest::Rove(RovingDirection::Next));
        });
        assert_eq!(focus.focused(), None);
    }

    #[test]
    fn begin_frame_resets_scope_ids() {
        let mut focus = FocusManager::new();
        focus.begin_frame();
        let first = focus.push_scope(FocusScopeKind::Modal);
        focus.end_frame();
        focus.begin_frame();
        let second = focus.push_scope(FocusScopeKind::Modal);
        focus.end_frame();
        assert_eq!(first, second, "scope ids reset each frame");
    }

    #[test]
    fn begin_frame_truncates_leaked_scopes() {
        let mut focus = FocusManager::new();
        focus.push_scope(FocusScopeKind::Modal);
        focus.push_scope(FocusScopeKind::Roving);
        focus.begin_frame();
        assert_eq!(focus.current_scope(), super::FocusScopeId::ROOT);
    }

    #[test]
    fn register_tab_stop_dedupes_per_frame() {
        let mut focus = FocusManager::new();
        let a = id("a");
        let b = id("b");
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request_focus(a);
        });
        assert_eq!(focus.tab_stops().len(), 2);
        run_frame(&mut focus, |f| {
            f.register_tab_stop(a);
            f.register_tab_stop(a);
            f.register_tab_stop(b);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(b));
    }

    #[test]
    fn programmatic_focus_holds_for_focusable_non_tab_stop() {
        let mut focus = FocusManager::new();
        let click_only = id("canvas");
        run_frame(&mut focus, |f| {
            f.register_focusable(click_only);
            f.request_focus(click_only);
        });
        assert_eq!(focus.focused(), Some(click_only));
    }

    #[test]
    fn programmatic_focus_drops_for_unknown_widget() {
        let mut focus = FocusManager::new();
        let ghost = id("ghost");
        run_frame(&mut focus, |f| {
            f.request_focus(ghost);
        });
        assert_eq!(focus.focused(), None);
    }

    #[test]
    fn is_text_input_focused_distinguishes_text_input_from_button() {
        let mut focus = FocusManager::new();
        let button = id("button");
        let input = id("input");
        run_frame(&mut focus, |f| {
            f.register_focusable(button);
            f.register_focusable(input);
            f.register_text_input(input);
            f.request_focus(button);
        });
        assert_eq!(focus.focused(), Some(button));
        assert!(
            !focus.is_text_input_focused(),
            "button focus must not match"
        );
        run_frame(&mut focus, |f| {
            f.register_focusable(button);
            f.register_focusable(input);
            f.register_text_input(input);
            f.request_focus(input);
        });
        assert_eq!(focus.focused(), Some(input));
        assert!(focus.is_text_input_focused());
    }

    #[test]
    fn is_text_input_focused_false_when_nothing_focused() {
        let focus = FocusManager::new();
        assert!(!focus.is_text_input_focused());
    }

    #[test]
    fn modal_traps_tab_with_nested_roving_excluded() {
        let mut focus = FocusManager::new();
        let modal_a = id("modal_a");
        let rove_x = id("rove_x");
        let modal_b = id("modal_b");
        run_frame(&mut focus, |f| {
            f.push_scope(FocusScopeKind::Modal);
            f.register_tab_stop(modal_a);
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(rove_x);
            f.pop_scope();
            f.register_tab_stop(modal_b);
            f.request_focus(modal_a);
        });
        run_frame(&mut focus, |f| {
            f.push_scope(FocusScopeKind::Modal);
            f.register_tab_stop(modal_a);
            f.push_scope(FocusScopeKind::Roving);
            f.register_tab_stop(rove_x);
            f.pop_scope();
            f.register_tab_stop(modal_b);
            f.request(FocusRequest::Advance);
        });
        assert_eq!(focus.focused(), Some(modal_b));
    }
}
