lemur/src/input.rs

use std::{
    collections::{hash_map::Entry, HashMap},
    sync::{Arc, RwLock},
};

use gilrs::{ev::Code, Axis, Button, Gamepad, GamepadId};
use winit::keyboard::{KeyCode, PhysicalKey};

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

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

pub struct GamepadMapping {
    buttons: HashMap<Code, VBKey>,
    axes: HashMap<Code, (VBKey, VBKey)>,
}

impl GamepadMapping {
    fn for_gamepad(gamepad: &Gamepad) -> Self {
        let mut buttons = HashMap::new();
        let mut default_button = |btn: Button, key: VBKey| {
            if let Some(code) = gamepad.button_code(btn) {
                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 axes = HashMap::new();
        let mut default_axis = |axis: Axis, neg: VBKey, pos: VBKey| {
            if let Some(code) = gamepad.axis_code(axis) {
                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, axes }
    }
}

#[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);
    }

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

    pub fn keyboard_mapping_names(&self) -> HashMap<VBKey, 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
            .into_iter()
            .map(|(k, mut v)| {
                v.sort();
                (k, v.join(", "))
            })
            .collect()
    }
}

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

impl MappingProvider {
    pub fn new() -> Self {
        let mut mappings = HashMap::new();

        let mut p1_mappings = InputMapping::default();
        let p2_mappings = InputMapping::default();
        let mut default_key = |code, key| {
            p1_mappings.add_keyboard_mapping(key, PhysicalKey::Code(code));
        };
        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);

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

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

    pub fn map_gamepad(&self, sim_id: SimId, 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 gamepad_mapping = match lock.entry(device_id) {
            Entry::Occupied(entry) => entry.get().clone(),
            Entry::Vacant(entry) => {
                let mappings = GamepadMapping::for_gamepad(gamepad);
                entry.insert(Arc::new(RwLock::new(mappings))).clone()
            }
        };
        drop(lock);
        self.for_sim(sim_id)
            .write()
            .unwrap()
            .gamepads
            .insert(gamepad.id(), gamepad_mapping);
    }
}