lyra-engine/lyra-game/src/input/action.rs

use std::{collections::HashMap, ops::Deref, hash::{Hash, DefaultHasher, Hasher}, fmt::Debug};

use glam::Vec2;
use lyra_ecs::World;
use lyra_reflect::Reflect;

use crate::{plugin::Plugin, game::GameStages, EventQueue};

use super::{Button, InputButtons, KeyCode, MouseButton, MouseMotion};

pub trait ActionLabel: Debug {
    /// Returns a unique hash of the label.
    fn label_hash(&self) -> u64;
}

impl<T: Hash + Debug> ActionLabel for T {
    fn label_hash(&self) -> u64 {
        let mut s = DefaultHasher::new();
        self.hash(&mut s);
        s.finish()
    }
}

#[derive(Clone)]
pub struct ActionLabelWrapper(String, u64);

impl Debug for ActionLabelWrapper {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

impl ActionLabel for ActionLabelWrapper {
    fn label_hash(&self) -> u64 {
        self.1
    }
}

impl<A: ActionLabel + Hash> From<A> for ActionLabelWrapper {
    fn from(value: A) -> Self {
        let lbl = format!("{:?}", value);
        Self(lbl, value.label_hash())
    }
}

/// Some commonly used action labels.
/// 
/// The built-in systems uses these labels
#[derive(Clone, Copy, Hash, Debug)]
pub enum CommonActionLabel {
    MoveForwardBackward,
    MoveLeftRight,
    MoveUpDown,
    LookLeftRight,
    LookUpDown,
    LookRoll,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum GamepadFormat {
    DualAxis,
    Joystick,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum GamepadButton {
    FaceBottom,
    FaceLeft,
    FaceRight,
    FaceTop,
    VirtualConfirm,
    VirtualDeny,
    LThumbstick,
    RThumbstick,
    DPadUp,
    DPadDown,
    DPadLeft,
    DPadRight,
    LShoulder,
    RShoulder,
    LTrigger,
    RTrigger,
    Special,
    LSpecial,
    RSpecial,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum GamepadAxis {
    LThumbstickX,
    LThumbstickY,
    RThumbstickX,
    RThumbstickY,
    LTrigger,
    RTrigger,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum GamepadInput {
    Button(GamepadButton),
    Axis(GamepadAxis),
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum MouseAxis {
    X,
    Y,
    ScrollWheel,
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum MouseInput {
    Button(MouseButton),
    Axis(MouseAxis),
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum ActionSource {
    Keyboard(KeyCode),
    Gamepad(GamepadFormat, GamepadInput),
    Mouse(MouseInput)
}

impl ActionSource {
    /// Convert this Source into a Binding. Uses a default Binding modifier value of `1.0`.
    pub fn into_binding(self) -> Binding {
        Binding::from_source(self)
    }

    /// Convert this Source into a Binding, setting the modifier.
    pub fn into_binding_modifier(self, modifier: f32) -> Binding {
        Binding::from_source_modifier(self, modifier)
    }
}

#[derive(Clone, Debug)]
pub struct Binding {
    pub source: ActionSource,
    pub modifier: f32,
}

impl Binding {
    /// Create a binding from a Source. Uses a default value of `1.0` as the modifier.
    pub fn from_source(source: ActionSource) -> Self {
        // TODO: Some types of input will have a default modifier, mainly axis'
        Self {
            source,
            modifier: 1.0,
        }
    }

    /// Create a binding from a Source, with a modifier
    pub fn from_source_modifier(source: ActionSource, modifier: f32) -> Self {
        Self {
            source,
            modifier,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ActionState {
    /// No input is being given by the user for the button/axis.
    Idle,
    /// The button is pressed
    Pressed(f32),
    /// The button was just pressed
    JustPressed(f32),
    /// The button was just released
    JustReleased,
    
    //Released,
    /// The axis moved
    Axis(f32),
    /// Some other state with a modifier
    Other(f32),
}

#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub enum ActionKind {
    Button,
    Axis
}

#[derive(Clone, Debug)]
pub struct Action {
    pub kind: ActionKind,
    pub state: ActionState,
}

impl Action {
    pub fn new(kind: ActionKind) -> Self {
        Self {
            kind,
            state: ActionState::Idle,
        }
    }
}

#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct LayoutId(pub u32);

impl From<u32> for LayoutId {
    fn from(value: u32) -> Self {
        Self(value)
    }
}

#[derive(Clone, Default)]
pub struct Layout {
    mappings: HashMap<ActionMappingId, ActionMapping>,
    active_mapping: ActionMappingId,
}

impl Layout {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_mapping(&mut self, mapping: ActionMapping) -> &mut Self {
        self.mappings.insert(mapping.id, mapping);

        self
    }
}

#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct ActionMappingId(pub u32);

impl From<u32> for ActionMappingId {
    fn from(value: u32) -> Self {
        Self(value)
    }
}

#[derive(Clone)]
pub struct ActionMapping {
    layout: LayoutId,
    id: ActionMappingId,
    action_binds: HashMap<u64, Vec<Binding>>,
}

impl ActionMapping {
    pub fn new(layout: LayoutId, id: ActionMappingId) -> Self {
        Self {
            layout,
            id,
            action_binds: HashMap::new(),
        }
    }

    pub fn builder(layout: LayoutId, id: ActionMappingId) -> ActionMappingBuilder {
        ActionMappingBuilder::new(ActionMapping::new(layout, id))
    }

    /// Creates a binding for the action.
    /// 
    /// If the action is not in this layout, this will panic!
    /// 
    /// Parameters:
    /// * `action` - The label corresponding to the action in this Layout.
    /// * `bind` - The Binding to add to the Action.
    pub fn bind<L>(&mut self, action: L, bindings: &[Binding]) -> &mut Self
    where
        L: ActionLabel
    {
        let mut bindings = bindings.to_vec();

        let action_binds = self.action_binds.entry(action.label_hash()).or_default();
        action_binds.append(&mut bindings);

        self
    }

    pub fn finish(self) -> Self {
        self
    }
}

pub struct ActionMappingBuilder {
    mapping: ActionMapping,
}

impl ActionMappingBuilder {
    fn new(mapping: ActionMapping) -> Self {
        Self {
            mapping,
        }
    }

    pub fn bind<L>(mut self, action: L, bindings: &[Binding]) -> Self
    where
        L: ActionLabel
    {
        let mut bindings = bindings.to_vec();

        let action_binds = self.mapping.action_binds.entry(action.label_hash()).or_default();
        action_binds.append(&mut bindings);

        self
    }

    pub fn finish(self) -> ActionMapping {
        self.mapping
    }
}

#[derive(Clone, Default, Reflect)]
pub struct ActionHandler {
    #[reflect(skip)] // TODO: dont just skip all these
    pub actions: HashMap<u64, Action>,
    #[reflect(skip)]
    pub layouts: HashMap<LayoutId, Layout>,
    #[reflect(skip)]
    pub current_layout: LayoutId,
    #[reflect(skip)]
    pub current_mapping: ActionMappingId,
}

impl ActionHandler {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn builder() -> ActionHandlerBuilder {
        ActionHandlerBuilder::default()
    }

    pub fn add_layout(&mut self, id: LayoutId) {
        self.layouts.insert(id, Layout::new());
    }

    pub fn action<L>(&self, label: L) -> Option<&Action>
    where
        L: ActionLabel
    {
        self.actions.get(&label.label_hash())
    }

    pub fn add_action<L>(&mut self, label: L, action: Action)
    where
        L: ActionLabel
    {
        self.actions.insert(label.label_hash(), action);
    }

    pub fn add_mapping(&mut self, mapping: ActionMapping) {
        let layout = self.layouts.get_mut(&mapping.layout).unwrap();
        layout.add_mapping(mapping);
    }

    /// Returns true if the action is pressed (or was just pressed).
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn is_action_pressed<L>(&self, action: L) -> bool
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        matches!(action.state, ActionState::Pressed(_) | ActionState::JustPressed(_))
    }

    /// Returns true if the action was just pressed.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn was_action_just_pressed<L>(&self, action: L) -> bool
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        matches!(action.state, ActionState::JustPressed(_))
    }

    /// Returns true if the action was just released.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn was_action_just_released<L>(&self, action: L) -> bool
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        matches!(action.state, ActionState::JustReleased)
    }

    /// Returns an action's state.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn get_action_state<L>(&self, action: L) -> ActionState
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        action.state
    }

    /// Returns the action's modifier if it is pressed (or was just pressed).
    /// Returns `None` if the action's state is not `ActionState::Pressed` or `ActionState::JustPressed`.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn get_pressed_modifier<L>(&self, action: L) -> Option<f32>
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        match action.state {
            ActionState::Pressed(v) | ActionState::JustPressed(v) => Some(v),
            _ => None,
        }
    }

    /// Returns the action's modifier if it was just pressed.
    /// Returns `None` if the action's state is not `ActionState::JustPressed`.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn get_just_pressed_modifier<L>(&self, action: L) -> Option<f32>
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        match action.state {
            ActionState::JustPressed(v) => Some(v),
            _ => None,
        }
    }
    
    /// Returns the action's modifier if its an updated axis.
    /// Returns `None` if the action's state is not `ActionState::Axis`.
    /// 
    /// This will panic if the action name does not correspond to an action.
    pub fn get_axis_modifier<L>(&self, action: L) -> Option<f32>
    where
        L: ActionLabel
    {
        let action = self.actions.get(&action.label_hash())
            .unwrap_or_else(|| panic!("Action {action:?} was not found"));

        match action.state {
            ActionState::Axis(v) => Some(v),
            _ => None,
        }
    }
}

#[derive(Default)]
pub struct ActionHandlerBuilder {
    handler: ActionHandler,
}

impl ActionHandlerBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn add_layout(mut self, id: LayoutId) -> Self {
        self.handler.layouts.insert(id, Layout::new());

        self
    }

    pub fn add_action<L>(mut self, label: L, action: Action) -> Self
    where
        L: ActionLabel
    {
        self.handler.actions.insert(label.label_hash(), action);

        self
    }

    pub fn add_mapping(mut self, mapping: ActionMapping) -> Self {
        let layout = self.handler.layouts.get_mut(&mapping.layout).unwrap();
        layout.add_mapping(mapping);

        self
    }

    pub fn finish(self) -> ActionHandler {
        self.handler
    }
}

fn actions_system(world: &mut World) -> anyhow::Result<()> {
    let keys = world.try_get_resource::<InputButtons<KeyCode>>()
        .map(|r| r.deref().clone());
    let mouse_events = world
        .try_get_resource::<EventQueue>()
        .and_then(|q| q.read_events::<MouseMotion>());
    //let mouse = world.try_get_resource()

    // clear the states of all axises each frame
    {
        let mut handler = world.try_get_resource_mut::<ActionHandler>()
            .expect("No Input Action handler was created in the world!");

        let layout = handler.layouts.get(&handler.current_layout).expect("No active layout");
        let mapping = layout.mappings.get(&layout.active_mapping).expect("No active mapping");
        
        for (action, _) in mapping.action_binds.clone().iter() {
            let action = handler.actions.get_mut(action).expect("Action name for binding is invalid!");
            if action.kind == ActionKind::Axis {
                action.state = ActionState::Axis(0.0);
            }
        }
    }
    
    let motion_avg = if let Some(mouse_events) = mouse_events {
        let count = mouse_events.len();
        let mut sum = Vec2::ZERO;
        for mm in mouse_events {
            sum += mm.delta;
        }
        Some(sum / count as f32)
    } else { None };

    let mut handler = world.try_get_resource_mut::<ActionHandler>()
        .expect("No Input Action handler was created in the world!");

    let layout = handler.layouts.get(&handler.current_layout).expect("No active layout");
    let mapping = layout.mappings.get(&layout.active_mapping).expect("No active mapping");

    for (action_lbl, binds) in mapping.action_binds.clone().iter() {
        let action = handler.actions.get_mut(action_lbl).expect("Action name for binding is invalid!");
        let mut new_state = None;

        for bind in binds.iter() {
            match bind.source {
                ActionSource::Keyboard(key) => if let Some(keys) = &keys {
                    // JustPressed needs to be first, since is_pressed includes buttons that
                    // were just pressed.
                    match action.kind {
                        ActionKind::Button => {
                            if keys.was_just_pressed(key) {
                                new_state = Some(ActionState::JustPressed(bind.modifier));
                            } else if keys.is_pressed(key) {
                                new_state = Some(ActionState::Pressed(bind.modifier));
                            }
                        },
                        ActionKind::Axis => {
                            if keys.is_pressed(key) {
                                new_state = Some(ActionState::Axis(bind.modifier));
                            }
                        }
                    }
                    
                },
                ActionSource::Gamepad(_, _) => todo!(),
                ActionSource::Mouse(m) => match m {
                    MouseInput::Button(_) => todo!(),
                    MouseInput::Axis(a) => if let Some(motion_avg) = motion_avg {
                        match a {
                            MouseAxis::X => {
                                new_state = Some(ActionState::Axis(motion_avg.x));
                            },
                            MouseAxis::Y => {
                                new_state = Some(ActionState::Axis(motion_avg.y));
                            },
                            MouseAxis::ScrollWheel => todo!(),
                        }
                    },
                },
            }
        }

        if let Some(new_state) = new_state {
            action.state = new_state;
        } else if action.kind == ActionKind::Button {
            match action.state {
                ActionState::Idle => {},
                ActionState::JustReleased => action.state = ActionState::Idle,
                _ => action.state = ActionState::JustReleased,
            }
        }
    }

    Ok(())
}

#[derive(Default, Clone, Copy)]
pub struct InputActionPlugin;

impl Plugin for InputActionPlugin {
    fn setup(&self, game: &mut crate::game::Game) {
        game.add_system_to_stage(GameStages::PreUpdate, "input_actions", actions_system, &[]);
    }
}