bevy_extended_ui 1.6.0-beta.2

use crate::CurrentWidgetState;
use crate::widgets::{BindToID, IgnoreParentState, UIGenID, UIWidgetState};
use bevy::prelude::*;
use std::collections::HashMap;

/// Represents the `BoundStateIndex` data structure used by the extended UI system.
#[derive(Resource, Default)]
pub struct BoundStateIndex {
    by_widget: HashMap<usize, Vec<Entity>>,
    by_entity: HashMap<Entity, usize>,
}

/// Plugin that manages widget focus and state propagation.
pub struct StateService;

impl Plugin for StateService {
    /// Registers widget state systems.
    fn build(&self, app: &mut App) {
        app.init_resource::<BoundStateIndex>();
        app.register_type::<Pickable>();
        app.add_systems(
            PostUpdate,
            (refresh_bound_state_index, update_widget_states).chain(),
        );
        app.add_systems(
            Update,
            (
                internal_state_check.run_if(resource_changed::<CurrentWidgetState>),
                handle_tab_focus,
                unfocus_disabled,
            ),
        );
    }
}

/// Handles `remove_from_bound_state_index` in the extended UI workflow.
fn remove_from_bound_state_index(index: &mut BoundStateIndex, entity: Entity, widget_id: usize) {
    let should_remove = if let Some(entries) = index.by_widget.get_mut(&widget_id) {
        entries.retain(|current| *current != entity);
        entries.is_empty()
    } else {
        false
    };

    if should_remove {
        index.by_widget.remove(&widget_id);
    }
}

/// Handles `refresh_bound_state_index` in the extended UI workflow.
fn refresh_bound_state_index(
    mut index: ResMut<BoundStateIndex>,
    query: Query<(Entity, &BindToID), Or<(Added<BindToID>, Changed<BindToID>)>>,
    mut removed: RemovedComponents<BindToID>,
) {
    for entity in removed.read() {
        if let Some(previous_id) = index.by_entity.remove(&entity) {
            remove_from_bound_state_index(&mut index, entity, previous_id);
        }
    }

    for (entity, bind_to) in query.iter() {
        let widget_id = bind_to.0;

        if let Some(previous_id) = index.by_entity.insert(entity, widget_id) {
            if previous_id != widget_id {
                remove_from_bound_state_index(&mut index, entity, previous_id);
            }
        }

        let entries = index.by_widget.entry(widget_id).or_default();
        if !entries.contains(&entity) {
            entries.push(entity);
        }
    }
}

/// Synchronizes the widget state from parent UI elements to child elements linked via [`BindToID`].
///
/// This system propagates UI states such as `hovered`, `focused`, `readonly`, `disabled`, and `checked`
/// from widgets that have a [`UIGenID`] to other UI elements bound to the same ID.
///
/// # Parameters
/// - `main_query`: Retrieves all UI widgets with a [`UIGenID`] whose [`UIWidgetState`] has changed.
/// - `inner_query`: Finds all UI elements that are bound via [`BindToID`], excluding those with their
///   own `UIGenID` or an explicit [`IgnoreParentState`].
///
/// # Purpose
/// Enables state propagation for compound widgets like checkbox containers, input groups, or radio button groups.
///
/// # Example
/// If a container with ID `#group` is focused, and an internal widget is bound to it, the inner widget
/// will also be marked as focused.
pub fn update_widget_states(
    main_query: Query<(&UIGenID, &UIWidgetState), (Changed<UIWidgetState>, With<UIGenID>)>,
    index: Option<Res<BoundStateIndex>>,
    mut inner_query: Query<
        (Entity, &BindToID, &mut UIWidgetState),
        (Without<UIGenID>, Without<IgnoreParentState>),
    >,
) {
    for (id, state) in main_query.iter() {
        if let Some(index) = index.as_ref() {
            if let Some(targets) = index.by_widget.get(&id.get()) {
                for entity in targets {
                    let Ok((_, _, mut inner_state)) = inner_query.get_mut(*entity) else {
                        continue;
                    };

                    inner_state.hovered = state.hovered;
                    inner_state.focused = state.focused;
                    inner_state.readonly = state.readonly;
                    inner_state.disabled = state.disabled;
                    inner_state.checked = state.checked;
                }
                continue;
            }
        }

        for (_, bind_to, mut inner_state) in inner_query.iter_mut() {
            if bind_to.0 != id.get() {
                continue;
            }

            inner_state.hovered = state.hovered;
            inner_state.focused = state.focused;
            inner_state.readonly = state.readonly;
            inner_state.disabled = state.disabled;
            inner_state.checked = state.checked;
        }
    }
}

/// Clears the `focused` state from all widgets except the currently focused one.
///
/// Ensures that only a single UI widget is marked as focused at any given time.
/// The focused widget ID is tracked in the [`CurrentWidgetState`] resource.
///
/// # Parameters
/// - `current_state_element`: The current global widget focus state.
/// - `query`: All UI widgets with a [`UIGenID`] and a mutable [`UIWidgetState`].
///
/// # Behavior
/// If the current widget ID is `0` (none), the system does nothing.
/// Otherwise, it clears `focused = false` on all widgets except the one with the matching ID.
fn internal_state_check(
    current_state_element: Res<CurrentWidgetState>,
    mut query: Query<(&mut UIWidgetState, &UIGenID), With<UIGenID>>,
) {
    for (mut state, gen_id) in query.iter_mut() {
        if gen_id.get() == current_state_element.widget_id {
            continue;
        }
        state.focused = false;
    }
}

/// Handles keyboard tab navigation between focusable UI widgets.
///
/// This system detects when the Tab key is pressed and moves the focus to the next available widget,
/// based on sorted [`UIGenID`] values. Shift+Tab navigates in reverse.
///
/// # Parameters
/// - `keys`: The current keyboard input state.
/// - `mut current_state`: The global [`CurrentWidgetState`] resource tracking focused widget ID.
/// - `widgets`: A list of all focusable widgets that can receive focus.
///
/// # Behavior
/// - Widgets are sorted by `UIGenID.0`.
/// - Pressing `Tab` sets focus to the next widget in order.
/// - Pressing `Shift+Tab` sets focus to the previous widget in order.
/// - The focus wraps around if reaching the end or beginning.
///
/// # Requirements
/// All focusable widgets must have unique, non-zero `UIGenID` values.
///
/// # Example
/// When the user presses Tab while focused on widget `#2`, focus will move to widget `#3`.
fn handle_tab_focus(
    mut query: Query<(Entity, &mut UIWidgetState, &UIGenID)>,
    keyboard: Res<ButtonInput<KeyCode>>,
    mut current_state: ResMut<CurrentWidgetState>,
) {
    if !keyboard.just_pressed(KeyCode::Tab) {
        return;
    }
    let reverse = keyboard.pressed(KeyCode::ShiftLeft) || keyboard.pressed(KeyCode::ShiftRight);

    let mut elems: Vec<_> = query
        .iter_mut()
        .filter(|(_, state, _)| !state.disabled)
        .collect();

    elems.sort_by_key(|(_, _, id)| id.get());

    let len = elems.len();
    if len == 0 {
        return;
    }

    let mut focused_idx = None;
    for (i, (_, state, _)) in elems.iter().enumerate() {
        if state.focused {
            focused_idx = Some(i);
            break;
        }
    }

    match focused_idx {
        Some(i) => {
            elems[i].1.focused = false;
            let next = if reverse {
                (i + len - 1) % len
            } else {
                (i + 1) % len
            };
            elems[next].1.focused = true;
            current_state.widget_id = elems[next].2.get();
        }
        None => {
            let next = if reverse { len - 1 } else { 0 };
            elems[next].1.focused = true;
            current_state.widget_id = elems[next].2.get();
        }
    }
}

/// Clears focus from widgets that became disabled.
fn unfocus_disabled(mut q: Query<&mut UIWidgetState, Changed<UIWidgetState>>) {
    for mut s in &mut q {
        if s.disabled && s.focused {
            s.focused = false;
        }
    }
}