lesavka/client/src/input/inputs.rs

// client/src/input/inputs.rs

use anyhow::{Context, Result, bail};
use evdev::{AbsoluteAxisCode, Device, EventType, KeyCode, RelativeAxisCode};
use std::collections::HashSet;
use tokio::{
    sync::broadcast::Sender,
    time::{Duration, interval},
};
use tracing::{debug, info, warn};

use lesavka_common::lesavka::{KeyboardReport, MouseReport};

use super::{keyboard::KeyboardAggregator, mouse::MouseAggregator};
use crate::layout::{Layout, apply as apply_layout};
use tokio::sync::mpsc::UnboundedSender;

pub struct InputAggregator {
    kbd_tx: Sender<KeyboardReport>,
    mou_tx: Sender<MouseReport>,
    dev_mode: bool,
    released: bool,
    magic_active: bool,
    pending_release: bool,
    pending_kill: bool,
    pending_keys: HashSet<KeyCode>,
    paste_tx: Option<UnboundedSender<String>>,
    keyboards: Vec<KeyboardAggregator>,
    mice: Vec<MouseAggregator>,
}

impl InputAggregator {
    pub fn new(
        dev_mode: bool,
        kbd_tx: Sender<KeyboardReport>,
        mou_tx: Sender<MouseReport>,
        paste_tx: Option<UnboundedSender<String>>,
    ) -> Self {
        Self {
            kbd_tx,
            mou_tx,
            dev_mode,
            released: false,
            magic_active: false,
            pending_release: false,
            pending_kill: false,
            pending_keys: HashSet::new(),
            paste_tx,
            keyboards: Vec::new(),
            mice: Vec::new(),
        }
    }

    /// Called once at startup: enumerates input devices,
    /// classifies them, and constructs a aggregator struct per type.
    pub fn init(&mut self) -> Result<()> {
        let paths = std::fs::read_dir("/dev/input").context("Failed to read /dev/input")?;

        let mut found_any = false;

        for entry in paths {
            let entry = entry?;
            let path = entry.path();

            // skip anything that isn't "event*"
            if !path
                .file_name()
                .map_or(false, |f| f.to_string_lossy().starts_with("event"))
            {
                continue;
            }

            // ─── open the event node read-write *without* unsafe ──────────
            let mut dev = match Device::open(&path) {
                Ok(d) => d,
                Err(e) => {
                    warn!("❌ open {}: {e}", path.display());
                    continue;
                }
            };

            // non-blocking so fetch_events never stalls the whole loop
            dev.set_nonblocking(true)
                .with_context(|| format!("set_non_blocking {:?}", path))?;

            match classify_device(&dev) {
                DeviceKind::Keyboard => {
                    dev.grab()
                        .with_context(|| format!("grabbing keyboard {path:?}"))?;
                    info!(
                        "🤏🖱️ Grabbed keyboard {:?}",
                        dev.name().unwrap_or("UNKNOWN")
                    );

                    // pass dev_mode to aggregator
                    // let kbd_agg = KeyboardAggregator::new(dev, self.dev_mode);
                    let kbd_agg = KeyboardAggregator::new(
                        dev,
                        self.dev_mode,
                        self.kbd_tx.clone(),
                        self.paste_tx.clone(),
                    );
                    self.keyboards.push(kbd_agg);
                    found_any = true;
                    continue;
                }
                DeviceKind::Mouse => {
                    dev.grab()
                        .with_context(|| format!("grabbing mouse {path:?}"))?;
                    info!("🤏⌨️ Grabbed mouse {:?}", dev.name().unwrap_or("UNKNOWN"));

                    // let mouse_agg = MouseAggregator::new(dev);
                    let mouse_agg = MouseAggregator::new(dev, self.dev_mode, self.mou_tx.clone());
                    self.mice.push(mouse_agg);
                    found_any = true;
                    continue;
                }
                DeviceKind::Other => {
                    debug!(
                        "Skipping non-kbd/mouse device: {:?}",
                        dev.name().unwrap_or("UNKNOWN")
                    );
                    continue;
                }
            }
        }

        if !found_any {
            bail!("No suitable keyboard/mouse devices found or none grabbed.");
        }

        Ok(())
    }

    /// We spawn the sub-aggregators in a loop or using separate tasks.
    /// (For a real system: you'd spawn a separate task for each aggregator.)
    pub async fn run(&mut self) -> Result<()> {
        // Example approach: poll each aggregator in a simple loop
        let mut tick = interval(Duration::from_millis(10));
        let mut current = Layout::SideBySide;
        loop {
            let mut want_kill = false;
            for kbd in &mut self.keyboards {
                kbd.process_events();
                want_kill |= kbd.magic_kill();
            }
            let magic_now = self.keyboards.iter().any(|k| k.magic_grab());
            let magic_left = self.keyboards.iter().any(|k| k.magic_left());
            let magic_right = self.keyboards.iter().any(|k| k.magic_right());

            if magic_now && !self.magic_active {
                self.toggle_grab();
            }
            if (magic_left || magic_right) && self.magic_active {
                current = match current {
                    Layout::SideBySide => Layout::FullLeft,
                    Layout::FullLeft => Layout::FullRight,
                    Layout::FullRight => Layout::SideBySide,
                };
                apply_layout(current);
            }
            if want_kill && !self.pending_kill {
                warn!("🧙  magic chord - killing 🪄 AVADA KEDAVRA!!! 💥💀⚰️");
                for k in &mut self.keyboards {
                    k.send_empty_report();
                    k.set_send(false);
                }
                for m in &mut self.mice {
                    m.reset_state();
                    m.set_send(false);
                }
                self.pending_kill = true;
                self.capture_pending_keys();
            }

            if self.pending_release || self.pending_kill {
                let chord_released = if self.pending_keys.is_empty() {
                    !self
                        .keyboards
                        .iter()
                        .any(|k| k.magic_grab() || k.magic_kill())
                } else {
                    self.pending_keys
                        .iter()
                        .all(|key| !self.keyboards.iter().any(|k| k.has_key(*key)))
                };
                if chord_released {
                    for k in &mut self.keyboards {
                        k.set_grab(false);
                        k.reset_state();
                    }
                    for m in &mut self.mice {
                        m.set_grab(false);
                        m.reset_state();
                    }
                    self.released = true;
                    if self.pending_kill {
                        return Ok(());
                    }
                    self.pending_release = false;
                    self.pending_keys.clear();
                }
            }

            for mouse in &mut self.mice {
                mouse.process_events();
            }

            self.magic_active = magic_now;
            tick.tick().await;
        }
    }

    fn toggle_grab(&mut self) {
        if self.pending_release || self.pending_kill {
            return;
        }
        if self.released {
            tracing::info!("🧙  magic chord - restricting devices 🪄 IMPERIUS!!! 🎮🔒");
        } else {
            tracing::info!("🧙  magic chord - freeing devices 🪄 EXPELLIARMUS!!! 🔓🕊️");
        }
        if self.released {
            // switching to remote control
            for k in &mut self.keyboards {
                k.reset_state();
                k.set_send(true);
                k.set_grab(true);
            }
            for m in &mut self.mice {
                m.reset_state();
                m.set_send(true);
                m.set_grab(true);
            }
            self.released = false;
            self.pending_release = false;
        } else {
            // switching to local control: stop sending, keep grab until chord released
            for k in &mut self.keyboards {
                k.send_empty_report();
                k.set_send(false);
            }
            for m in &mut self.mice {
                m.reset_state();
                m.set_send(false);
            }
            self.pending_release = true;
            self.capture_pending_keys();
        }
    }

    fn capture_pending_keys(&mut self) {
        self.pending_keys.clear();
        for k in &self.keyboards {
            for key in k.pressed_keys_snapshot() {
                self.pending_keys.insert(key);
            }
        }
    }
}

/// The classification function
fn classify_device(dev: &Device) -> DeviceKind {
    let evbits = dev.supported_events();

    // Keyboard logic
    if evbits.contains(EventType::KEY) {
        if let Some(keys) = dev.supported_keys() {
            if keys.contains(KeyCode::KEY_A) || keys.contains(KeyCode::KEY_ENTER) {
                return DeviceKind::Keyboard;
            }
        }
    }

    // Mouse logic (relative)
    if evbits.contains(EventType::RELATIVE) {
        if let (Some(rel), Some(keys)) = (dev.supported_relative_axes(), dev.supported_keys()) {
            let has_xy =
                rel.contains(RelativeAxisCode::REL_X) && rel.contains(RelativeAxisCode::REL_Y);
            let has_btn = keys.contains(KeyCode::BTN_LEFT) || keys.contains(KeyCode::BTN_RIGHT);
            if has_xy && has_btn {
                return DeviceKind::Mouse;
            }
        }
    }
    // Touchpad logic (absolute)
    if evbits.contains(EventType::ABSOLUTE) {
        if let (Some(abs), Some(keys)) = (dev.supported_absolute_axes(), dev.supported_keys()) {
            let has_xy = (abs.contains(AbsoluteAxisCode::ABS_X)
                && abs.contains(AbsoluteAxisCode::ABS_Y))
                || (abs.contains(AbsoluteAxisCode::ABS_MT_POSITION_X)
                    && abs.contains(AbsoluteAxisCode::ABS_MT_POSITION_Y));
            let has_btn = keys.contains(KeyCode::BTN_TOUCH) || keys.contains(KeyCode::BTN_LEFT);
            if has_xy && has_btn {
                return DeviceKind::Mouse;
            }
        }
    }

    DeviceKind::Other
}

/// Internal enum for device classification
#[derive(Debug, Clone, Copy)]
enum DeviceKind {
    Keyboard,
    Mouse,
    Other,
}