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);

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

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 }
    }

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

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()
        });
    }
}

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

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()
        });
    }
}

#[derive(Clone)]
pub struct MappingProvider {
    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() -> 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())),
            gamepad_info: 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 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;
        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) {
            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;
        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()
    }
}