lemur/src/input.rs

use std::{
    cmp::Ordering,
    collections::{hash_map::Entry, HashMap, HashSet},
    fmt::Display,
    str::FromStr,
    sync::{Arc, Mutex, RwLock},
};

use anyhow::anyhow;
use egui::{Key, KeyboardShortcut, Modifiers};
use gilrs::{ev::Code, Axis, Button, Gamepad, GamepadId};
use serde::{Deserialize, Serialize};
use winit::keyboard::{KeyCode, PhysicalKey};

use crate::{
    emulator::{SimId, VBKey},
    persistence::Persistence,
};

#[derive(Clone, PartialEq, Eq, Hash)]
struct DeviceId(u16, u16);

impl FromStr for DeviceId {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut ids = s.split("-");
        let vendor_id: u16 = ids.next().ok_or(anyhow!("missing vendor id"))?.parse()?;
        let product_id: u16 = ids.next().ok_or(anyhow!("missing product id"))?.parse()?;
        Ok(Self(vendor_id, product_id))
    }
}

impl Display for DeviceId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{}-{}", self.0, self.1))
    }
}

#[derive(Clone)]
pub struct GamepadInfo {
    pub id: GamepadId,
    pub name: String,
    device_id: DeviceId,
    pub bound_to: Option<SimId>,
}

pub trait Mappings {
    fn mapping_names(&self) -> HashMap<VBKey, Vec<String>>;
    fn clear_mappings(&mut self, key: VBKey);
    fn clear_all_mappings(&mut self);
    fn use_default_mappings(&mut self);
}

#[derive(Serialize, Deserialize)]
pub struct GamepadMapping {
    buttons: HashMap<Code, VBKey>,
    axes: HashMap<Code, (VBKey, VBKey)>,
    default_buttons: HashMap<Code, VBKey>,
    default_axes: HashMap<Code, (VBKey, VBKey)>,
}

impl GamepadMapping {
    fn for_gamepad(gamepad: &Gamepad) -> Self {
        let mut default_buttons = HashMap::new();
        let mut default_button = |btn: Button, key: VBKey| {
            if let Some(code) = gamepad.button_code(btn) {
                default_buttons.insert(code, key);
            }
        };
        default_button(Button::South, VBKey::A);
        default_button(Button::West, VBKey::B);
        default_button(Button::RightTrigger, VBKey::RT);
        default_button(Button::LeftTrigger, VBKey::LT);
        default_button(Button::Start, VBKey::STA);
        default_button(Button::Select, VBKey::SEL);

        let mut default_axes = HashMap::new();
        let mut default_axis = |axis: Axis, neg: VBKey, pos: VBKey| {
            if let Some(code) = gamepad.axis_code(axis) {
                default_axes.insert(code, (neg, pos));
            }
        };
        default_axis(Axis::LeftStickX, VBKey::LL, VBKey::LR);
        default_axis(Axis::LeftStickY, VBKey::LD, VBKey::LU);
        default_axis(Axis::RightStickX, VBKey::RL, VBKey::RR);
        default_axis(Axis::RightStickY, VBKey::RD, VBKey::RU);
        default_axis(Axis::DPadX, VBKey::LL, VBKey::LR);
        default_axis(Axis::DPadY, VBKey::LD, VBKey::LU);

        Self {
            buttons: default_buttons.clone(),
            axes: default_axes.clone(),
            default_buttons,
            default_axes,
        }
    }

    pub fn add_button_mapping(&mut self, key: VBKey, code: Code) {
        let entry = self.buttons.entry(code).or_insert(VBKey::empty());
        *entry = entry.union(key);
    }

    pub fn add_axis_neg_mapping(&mut self, key: VBKey, code: Code) {
        let entry = self
            .axes
            .entry(code)
            .or_insert((VBKey::empty(), VBKey::empty()));
        entry.0 = entry.0.union(key);
    }

    pub fn add_axis_pos_mapping(&mut self, key: VBKey, code: Code) {
        let entry = self
            .axes
            .entry(code)
            .or_insert((VBKey::empty(), VBKey::empty()));
        entry.1 = entry.1.union(key);
    }

    fn save_mappings(&self) -> PersistedGamepadMapping {
        fn flatten<V: Copy>(values: &HashMap<Code, V>) -> Vec<(Code, V)> {
            values.iter().map(|(k, v)| (*k, *v)).collect()
        }
        PersistedGamepadMapping {
            buttons: flatten(&self.buttons),
            axes: flatten(&self.axes),
            default_buttons: flatten(&self.default_buttons),
            default_axes: flatten(&self.default_axes),
        }
    }

    fn from_mappings(mappings: &PersistedGamepadMapping) -> Self {
        fn unflatten<V: Copy>(values: &[(Code, V)]) -> HashMap<Code, V> {
            values.iter().map(|(k, v)| (*k, *v)).collect()
        }
        Self {
            buttons: unflatten(&mappings.buttons),
            axes: unflatten(&mappings.axes),
            default_buttons: unflatten(&mappings.default_buttons),
            default_axes: unflatten(&mappings.default_axes),
        }
    }
}

impl Mappings for GamepadMapping {
    fn mapping_names(&self) -> HashMap<VBKey, Vec<String>> {
        let mut results: HashMap<VBKey, Vec<String>> = HashMap::new();
        for (axis, (left_keys, right_keys)) in &self.axes {
            for key in left_keys.iter() {
                results.entry(key).or_default().push(format!("-{axis}"));
            }
            for key in right_keys.iter() {
                results.entry(key).or_default().push(format!("+{axis}"));
            }
        }
        for (button, keys) in &self.buttons {
            for key in keys.iter() {
                results.entry(key).or_default().push(format!("{button}"));
            }
        }
        results
    }

    fn clear_mappings(&mut self, key: VBKey) {
        self.axes.retain(|_, (left, right)| {
            *left = left.difference(key);
            *right = right.difference(key);
            !(left.is_empty() && right.is_empty())
        });
        self.buttons.retain(|_, keys| {
            *keys = keys.difference(key);
            !keys.is_empty()
        });
    }

    fn clear_all_mappings(&mut self) {
        self.axes.clear();
        self.buttons.clear();
    }

    fn use_default_mappings(&mut self) {
        self.axes = self.default_axes.clone();
        self.buttons = self.default_buttons.clone();
    }
}

#[derive(Default)]
pub struct InputMapping {
    keys: HashMap<PhysicalKey, VBKey>,
    gamepads: HashMap<GamepadId, Arc<RwLock<GamepadMapping>>>,
}

impl InputMapping {
    pub fn map_keyboard(&self, key: &PhysicalKey) -> Option<VBKey> {
        self.keys.get(key).copied()
    }

    pub fn map_button(&self, id: &GamepadId, code: &Code) -> Option<VBKey> {
        let mappings = self.gamepads.get(id)?.read().unwrap();
        mappings.buttons.get(code).copied()
    }

    pub fn map_axis(&self, id: &GamepadId, code: &Code) -> Option<(VBKey, VBKey)> {
        let mappings = self.gamepads.get(id)?.read().unwrap();
        mappings.axes.get(code).copied()
    }

    pub fn add_keyboard_mapping(&mut self, key: VBKey, keyboard_key: PhysicalKey) {
        let entry = self.keys.entry(keyboard_key).or_insert(VBKey::empty());
        *entry = entry.union(key);
    }

    fn save_mappings(&self) -> PersistedKeyboardMapping {
        PersistedKeyboardMapping {
            keys: self.keys.iter().map(|(k, v)| (*k, *v)).collect(),
        }
    }

    fn restore_mappings(&mut self, persisted: &PersistedKeyboardMapping) {
        self.keys = persisted.keys.iter().map(|(k, v)| (*k, *v)).collect();
    }
}

impl Mappings for InputMapping {
    fn mapping_names(&self) -> HashMap<VBKey, Vec<String>> {
        let mut results: HashMap<VBKey, Vec<String>> = HashMap::new();
        for (keyboard_key, keys) in &self.keys {
            let name = match keyboard_key {
                PhysicalKey::Code(code) => format!("{code:?}"),
                k => format!("{:?}", k),
            };
            for key in keys.iter() {
                results.entry(key).or_default().push(name.clone());
            }
        }
        results
    }

    fn clear_mappings(&mut self, key: VBKey) {
        self.keys.retain(|_, keys| {
            *keys = keys.difference(key);
            !keys.is_empty()
        });
    }

    fn clear_all_mappings(&mut self) {
        self.keys.clear();
    }

    fn use_default_mappings(&mut self) {
        self.keys.clear();
        let mut default_key = |code, key| {
            self.keys.insert(PhysicalKey::Code(code), key);
        };
        default_key(KeyCode::KeyA, VBKey::SEL);
        default_key(KeyCode::KeyS, VBKey::STA);
        default_key(KeyCode::KeyD, VBKey::B);
        default_key(KeyCode::KeyF, VBKey::A);
        default_key(KeyCode::KeyE, VBKey::LT);
        default_key(KeyCode::KeyR, VBKey::RT);
        default_key(KeyCode::KeyI, VBKey::RU);
        default_key(KeyCode::KeyJ, VBKey::RL);
        default_key(KeyCode::KeyK, VBKey::RD);
        default_key(KeyCode::KeyL, VBKey::RR);
        default_key(KeyCode::ArrowUp, VBKey::LU);
        default_key(KeyCode::ArrowLeft, VBKey::LL);
        default_key(KeyCode::ArrowDown, VBKey::LD);
        default_key(KeyCode::ArrowRight, VBKey::LR);
    }
}

#[derive(Clone)]
pub struct MappingProvider {
    persistence: Persistence,
    device_mappings: Arc<RwLock<HashMap<DeviceId, Arc<RwLock<GamepadMapping>>>>>,
    sim_mappings: HashMap<SimId, Arc<RwLock<InputMapping>>>,
    gamepad_info: Arc<RwLock<HashMap<GamepadId, GamepadInfo>>>,
}

impl MappingProvider {
    pub fn new(persistence: Persistence) -> Self {
        let mut sim_mappings = HashMap::new();
        let mut device_mappings = HashMap::new();

        let mut p1_mappings = InputMapping::default();
        let mut p2_mappings = InputMapping::default();

        if let Ok(persisted) = persistence.load_config::<PersistedInputMappings>("mappings") {
            p1_mappings.restore_mappings(&persisted.p1_keyboard);
            p2_mappings.restore_mappings(&persisted.p2_keyboard);

            for (device_id, mappings) in persisted.gamepads {
                let Ok(device_id) = device_id.parse::<DeviceId>() else {
                    continue;
                };
                let gamepad = GamepadMapping::from_mappings(&mappings);
                device_mappings.insert(device_id, Arc::new(RwLock::new(gamepad)));
            }
        } else {
            p1_mappings.use_default_mappings();
        }

        sim_mappings.insert(SimId::Player1, Arc::new(RwLock::new(p1_mappings)));
        sim_mappings.insert(SimId::Player2, Arc::new(RwLock::new(p2_mappings)));
        Self {
            persistence,
            device_mappings: Arc::new(RwLock::new(device_mappings)),
            gamepad_info: Arc::new(RwLock::new(HashMap::new())),
            sim_mappings,
        }
    }

    pub fn for_sim(&self, sim_id: SimId) -> &Arc<RwLock<InputMapping>> {
        self.sim_mappings.get(&sim_id).unwrap()
    }

    pub fn for_gamepad(&self, gamepad_id: GamepadId) -> Option<Arc<RwLock<GamepadMapping>>> {
        let lock = self.gamepad_info.read().unwrap();
        let device_id = lock.get(&gamepad_id)?.device_id.clone();
        drop(lock);
        let lock = self.device_mappings.read().unwrap();
        lock.get(&device_id).cloned()
    }

    pub fn handle_gamepad_connect(&self, gamepad: &Gamepad) {
        let device_id = DeviceId(
            gamepad.vendor_id().unwrap_or_default(),
            gamepad.product_id().unwrap_or_default(),
        );
        let mut lock = self.device_mappings.write().unwrap();
        let mappings = match lock.entry(device_id.clone()) {
            Entry::Vacant(entry) => {
                let mappings = GamepadMapping::for_gamepad(gamepad);
                entry.insert(Arc::new(RwLock::new(mappings)))
            }
            Entry::Occupied(entry) => entry.into_mut(),
        }
        .clone();
        drop(lock);
        let mut lock = self.gamepad_info.write().unwrap();
        let bound_to = SimId::values()
            .into_iter()
            .find(|sim_id| lock.values().all(|info| info.bound_to != Some(*sim_id)));
        if let Entry::Vacant(entry) = lock.entry(gamepad.id()) {
            let info = GamepadInfo {
                id: *entry.key(),
                name: gamepad.name().to_string(),
                device_id,
                bound_to,
            };
            entry.insert(info);
        }
        drop(lock);
        if let Some(sim_id) = bound_to {
            self.for_sim(sim_id)
                .write()
                .unwrap()
                .gamepads
                .insert(gamepad.id(), mappings);
        }
    }

    pub fn handle_gamepad_disconnect(&self, gamepad_id: GamepadId) {
        let mut lock = self.gamepad_info.write().unwrap();
        let Some(info) = lock.remove(&gamepad_id) else {
            return;
        };
        if let Some(sim_id) = info.bound_to {
            self.for_sim(sim_id)
                .write()
                .unwrap()
                .gamepads
                .remove(&gamepad_id);
        }
    }

    pub fn assign_gamepad(&self, gamepad_id: GamepadId, sim_id: SimId) {
        self.unassign_gamepad(gamepad_id);
        let mut lock = self.gamepad_info.write().unwrap();
        let Some(info) = lock.get_mut(&gamepad_id) else {
            return;
        };
        info.bound_to = Some(sim_id);
        let device_id = info.device_id.clone();
        drop(lock);
        let Some(device_mappings) = self
            .device_mappings
            .write()
            .unwrap()
            .get(&device_id)
            .cloned()
        else {
            return;
        };
        self.for_sim(sim_id)
            .write()
            .unwrap()
            .gamepads
            .insert(gamepad_id, device_mappings);
    }

    pub fn unassign_gamepad(&self, gamepad_id: GamepadId) {
        let mut lock = self.gamepad_info.write().unwrap();
        let Some(info) = lock.get_mut(&gamepad_id) else {
            return;
        };
        if let Some(sim_id) = info.bound_to {
            let mut sim_mapping = self.for_sim(sim_id).write().unwrap();
            sim_mapping.gamepads.remove(&gamepad_id);
        }
        info.bound_to = None;
    }

    pub fn gamepad_info(&self) -> Vec<GamepadInfo> {
        self.gamepad_info
            .read()
            .unwrap()
            .values()
            .cloned()
            .collect()
    }

    pub fn save(&self) {
        let p1_keyboard = self.for_sim(SimId::Player1).read().unwrap().save_mappings();
        let p2_keyboard = self.for_sim(SimId::Player2).read().unwrap().save_mappings();
        let mut gamepads = HashMap::new();
        for (device_id, gamepad) in self.device_mappings.read().unwrap().iter() {
            let mapping = gamepad.read().unwrap().save_mappings();
            gamepads.insert(device_id.to_string(), mapping);
        }
        let persisted = PersistedInputMappings {
            p1_keyboard,
            p2_keyboard,
            gamepads,
        };
        let _ = self.persistence.save_config("mappings", &persisted);
    }
}

#[derive(Serialize, Deserialize)]
struct PersistedInputMappings {
    p1_keyboard: PersistedKeyboardMapping,
    p2_keyboard: PersistedKeyboardMapping,
    gamepads: HashMap<String, PersistedGamepadMapping>,
}

#[derive(Serialize, Deserialize)]
struct PersistedKeyboardMapping {
    keys: Vec<(PhysicalKey, VBKey)>,
}

#[derive(Serialize, Deserialize)]
struct PersistedGamepadMapping {
    buttons: Vec<(Code, VBKey)>,
    axes: Vec<(Code, (VBKey, VBKey))>,
    default_buttons: Vec<(Code, VBKey)>,
    default_axes: Vec<(Code, (VBKey, VBKey))>,
}

#[derive(Serialize, Deserialize)]
pub struct Shortcut {
    pub shortcut: KeyboardShortcut,
    pub command: Command,
}

#[derive(Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum Command {
    OpenRom,
    Quit,
    FrameAdvance,
    Reset,
    PauseResume,
    // if you update this, update Command::all and add a default
}

impl Command {
    pub fn all() -> [Self; 5] {
        [
            Self::OpenRom,
            Self::Quit,
            Self::PauseResume,
            Self::Reset,
            Self::FrameAdvance,
        ]
    }

    pub fn name(self) -> &'static str {
        match self {
            Self::OpenRom => "Open ROM",
            Self::Quit => "Exit",
            Self::PauseResume => "Pause/Resume",
            Self::Reset => "Reset",
            Self::FrameAdvance => "Frame Advance",
        }
    }
}

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

struct Shortcuts {
    all: Vec<(Command, KeyboardShortcut)>,
    by_command: HashMap<Command, KeyboardShortcut>,
}

impl Default for Shortcuts {
    fn default() -> Self {
        let mut shortcuts = Shortcuts {
            all: vec![],
            by_command: HashMap::new(),
        };
        shortcuts.set(
            Command::OpenRom,
            KeyboardShortcut::new(Modifiers::COMMAND, Key::O),
        );
        shortcuts.set(
            Command::Quit,
            KeyboardShortcut::new(Modifiers::COMMAND, Key::Q),
        );
        shortcuts.set(
            Command::PauseResume,
            KeyboardShortcut::new(Modifiers::NONE, Key::F5),
        );
        shortcuts.set(
            Command::Reset,
            KeyboardShortcut::new(Modifiers::SHIFT, Key::F5),
        );
        shortcuts.set(
            Command::FrameAdvance,
            KeyboardShortcut::new(Modifiers::NONE, Key::F6),
        );
        shortcuts
    }
}

impl Shortcuts {
    fn set(&mut self, command: Command, shortcut: KeyboardShortcut) {
        if self.by_command.insert(command, shortcut).is_some() {
            for (cmd, sht) in &mut self.all {
                if *cmd == command {
                    *sht = shortcut;
                    break;
                }
            }
        } else {
            self.all.push((command, shortcut));
        }
        self.all.sort_by(|l, r| order_shortcut(l.1, r.1));
    }

    fn unset(&mut self, command: Command) {
        if self.by_command.remove(&command).is_some() {
            self.all.retain(|(c, _)| *c != command);
        }
    }

    fn save(&self) -> PersistedShortcuts {
        let mut shortcuts = PersistedShortcuts { shortcuts: vec![] };
        for command in Command::all() {
            let shortcut = self.by_command.get(&command).copied();
            shortcuts.shortcuts.push((command, shortcut));
        }
        shortcuts
    }
}

fn order_shortcut(left: KeyboardShortcut, right: KeyboardShortcut) -> Ordering {
    left.logical_key.cmp(&right.logical_key).then_with(|| {
        specificity(left.modifiers)
            .cmp(&specificity(right.modifiers))
            .reverse()
    })
}

fn specificity(modifiers: egui::Modifiers) -> usize {
    let mut mods = 0;
    if modifiers.alt {
        mods += 1;
    }
    if modifiers.command || modifiers.ctrl {
        mods += 1;
    }
    if modifiers.shift {
        mods += 1;
    }
    mods
}

#[derive(Serialize, Deserialize)]
struct PersistedShortcuts {
    shortcuts: Vec<(Command, Option<KeyboardShortcut>)>,
}

#[derive(Clone)]
pub struct ShortcutProvider {
    persistence: Persistence,
    shortcuts: Arc<Mutex<Shortcuts>>,
}

impl ShortcutProvider {
    pub fn new(persistence: Persistence) -> Self {
        let mut shortcuts = Shortcuts::default();
        if let Ok(saved) = persistence.load_config::<PersistedShortcuts>("shortcuts") {
            for (command, shortcut) in saved.shortcuts {
                if let Some(shortcut) = shortcut {
                    shortcuts.set(command, shortcut);
                } else {
                    shortcuts.unset(command);
                }
            }
        }
        Self {
            persistence,
            shortcuts: Arc::new(Mutex::new(shortcuts)),
        }
    }

    pub fn shortcut_for(&self, command: Command) -> Option<KeyboardShortcut> {
        let lock = self.shortcuts.lock().unwrap();
        lock.by_command.get(&command).copied()
    }

    pub fn consume_all(&self, input: &mut egui::InputState) -> HashSet<Command> {
        let lock = self.shortcuts.lock().unwrap();
        lock.all
            .iter()
            .filter_map(|(command, shortcut)| input.consume_shortcut(shortcut).then_some(*command))
            .collect()
    }

    pub fn set(&self, command: Command, shortcut: KeyboardShortcut) {
        let updated = {
            let mut lock = self.shortcuts.lock().unwrap();
            lock.set(command, shortcut);
            lock.save()
        };
        let _ = self.persistence.save_config("shortcuts", &updated);
    }

    pub fn unset(&self, command: Command) {
        let updated = {
            let mut lock = self.shortcuts.lock().unwrap();
            lock.unset(command);
            lock.save()
        };
        let _ = self.persistence.save_config("shortcuts", &updated);
    }

    pub fn reset(&self) {
        let updated = {
            let mut lock = self.shortcuts.lock().unwrap();
            *lock = Shortcuts::default();
            lock.save()
        };
        let _ = self.persistence.save_config("shortcuts", &updated);
    }
}