bevy_extended_ui 1.7.0

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 (widget_id, widget_state) in main_query.iter() {
        if let Some(bound_index) = index.as_ref() {
            if let Some(bound_entities) = bound_index.by_widget.get(&widget_id.get()) {
                for bound_entity in bound_entities {
                    let Ok((_, _, mut bound_state)) = inner_query.get_mut(*bound_entity) else {
                        continue;
                    };
                    copy_widget_state(&mut bound_state, widget_state);
                }
                continue;
            }
        }

        for (_, bind_to, mut bound_state) in inner_query.iter_mut() {
            if bind_to.0 != widget_id.get() {
                continue;
            }
            copy_widget_state(&mut bound_state, widget_state);
        }
    }
}

fn copy_widget_state(target: &mut UIWidgetState, source: &UIWidgetState) {
    target.hovered = source.hovered;
    target.focused = source.focused;
    target.readonly = source.readonly;
    target.disabled = source.disabled;
    target.checked = source.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 widget_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 focusable_widgets: Vec<_> = widget_query
        .iter_mut()
        .filter(|(_, state, _)| !state.disabled)
        .collect();

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

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

    let mut focused_index = None;
    for (index, (_, state, _)) in focusable_widgets.iter().enumerate() {
        if state.focused {
            focused_index = Some(index);
            break;
        }
    }

    match focused_index {
        Some(current_index) => {
            focusable_widgets[current_index].1.focused = false;
            let next_index = if reverse {
                (current_index + focusable_count - 1) % focusable_count
            } else {
                (current_index + 1) % focusable_count
            };
            focusable_widgets[next_index].1.focused = true;
            current_state.widget_id = focusable_widgets[next_index].2.get();
        }
        None => {
            let next_index = if reverse { focusable_count - 1 } else { 0 };
            focusable_widgets[next_index].1.focused = true;
            current_state.widget_id = focusable_widgets[next_index].2.get();
        }
    }
}

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