fixscreen/src/lib.rs

use std::{
    collections::HashMap,
    fmt,
    fs::File,
    io::{self, BufRead, BufReader},
    path::Path,
    process::{Command, Output},
};

use regex::Regex;
use thiserror::Error;
use xrandr::{XHandle, XrandrError};

pub fn swap_workspaces(cfg: &Config) -> Result<Vec<Result<Output, CommandError>>, CommandError> {
    let monitors = XHandle::open()?.monitors()?;

    // TODO(wathiede): if I ever get two monitors with the same resolution, use EDID to find a
    // better key.
    let map: HashMap<_, _> = monitors
        .iter()
        .map(|m| ((m.width_px as usize, m.height_px as usize), m.name.clone()))
        .collect();
    // TODO(wathiede):
    // i3-msg commands to move them.
    // Example:
    // i3-msg '[workspace="trustee"]' move workspace to output DP-0
    Ok(cfg
        .screens
        .iter()
        .map(|s| {
            let key = match s.orientation {
                Orientation::None | Orientation::Invert => {
                    (s.resolution.width, s.resolution.height)
                }
                Orientation::Left | Orientation::Right => (s.resolution.height, s.resolution.width),
            };
            if let Some(m) = map.get(&key) {
                println!("Moving {:?} to {}", s.workspaces, m);
                run_move_workspace_cmd(&s.workspaces, m)
            } else {
                Vec::new()
            }
        })
        .flatten()
        .collect())
}

fn run_move_workspace_cmd(
    workspaces: &[String],
    monitor: &str,
) -> Vec<Result<Output, CommandError>> {
    workspaces
        .iter()
        .map(|workspace| {
            let s = format!(r#"[workspace="{workspace}"]"#);
            let args = vec!["i3-msg", &s, "move", "workspace", "to", "output", monitor];
            println!("{}", args.join(" "));
            if cfg!(debug_assertions) {
                Command::new("echo").args(args).output()
            } else {
                Command::new(&args[0]).args(&args[1..]).output()
            }
        })
        .map(|r| r.map_err(CommandError::from))
        .collect()
}

#[derive(Error, Debug)]
pub enum ParseError {
    #[error("error reading file")]
    Io(#[from] io::Error),
    #[error("invalid data: {0}")]
    Parse(String),
    #[error("xrandr: {0}")]
    XrandrError(#[from] XrandrError),
}

// Map monitor name to DFP connection.
pub type ScreenMapping = HashMap<String, String>;

pub fn screen_mapping_from_xorg_log<P: AsRef<Path>>(path: P) -> Result<ScreenMapping, ParseError> {
    let f = File::open(path)?;
    let f = BufReader::new(f);
    let re = Regex::new(r".*: (.*)\s+\(([^)]+)\): connected").unwrap();
    Ok(f.lines()
        .filter_map(|line| line.ok())
        .filter_map(|line| {
            if let Some(cap) = re.captures(&line) {
                Some((cap[1].to_string(), cap[2].to_string()))
            } else {
                None
            }
        })
        .collect())
}

#[derive(Error, Debug)]
pub enum CommandError {
    #[error("missing monitor: {0}")]
    MissingMonitor(String),
    #[error("error executing command")]
    Io(#[from] io::Error),
    #[error("xrandr error: {0}")]
    XrandrError(#[from] XrandrError),
}

#[derive(Clone, Debug, Default)]
pub struct Resolution {
    pub width: usize,
    pub height: usize,
}

#[derive(Clone, Debug, Default)]
pub struct Offset {
    pub x: isize,
    pub y: isize,
}

#[derive(Clone)]
pub enum Orientation {
    None,
    Right,
    Left,
    Invert,
}

impl Default for Orientation {
    fn default() -> Self {
        Orientation::None
    }
}
impl fmt::Debug for Orientation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Orientation::None => f.write_str("No Rotation"),
            Orientation::Right => f.write_str("Rotate Right"),
            Orientation::Left => f.write_str("Rotate Left"),
            Orientation::Invert => f.write_str("Invert"),
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct Screen {
    pub name: String,
    pub resolution: Resolution,
    pub offset: Offset,
    pub orientation: Orientation,
    pub primary: bool,
    pub workspaces: Vec<String>,
}

impl Screen {
    fn metamode(&self, map: &ScreenMapping) -> Result<String, CommandError> {
        let Resolution { width, height } = self.resolution;
        let Offset { x, y } = self.offset;
        let (_in_w, _in_h) = match self.orientation {
            Orientation::None | Orientation::Invert => (width, height),
            Orientation::Right | Orientation::Left => (height, width),
        };
        let rotation = match self.orientation {
            Orientation::None => "0",
            Orientation::Invert => "180",
            Orientation::Right => "270",
            Orientation::Left => "90",
        };

        let connection = match map.get(&self.name) {
            Some(connection) => connection,
            None => return Err(CommandError::MissingMonitor(self.name.to_string())),
        };

        Ok(format!(
            "{connection}: {w}x{h} {x:+}{y:+} {{ForceCompositionPipeline=On, Rotation={rotation}}}",
            connection = connection,
            w = width,
            h = height,
            x = x,
            y = y,
            rotation = rotation,
        ))
        /*
        Ok(format!("{connection}: {w}x{h} @{in_w}x{in_h} {x:+}{y:+} {{ForceCompositionPipeline=On, ViewPortIn={in_w}x{in_h}, ViewPortOut={w}x{h}+0+0, Rotation={rotation}}}",
        connection=connection,
        w=width,
        h=height,
        in_w=in_w,
        in_h=in_h,
        x=x,
        y=y,
        rotation=rotation,
        ))
        */
    }
}

#[derive(Debug, Default)]
pub struct Config {
    pub screens: Vec<Screen>,
}

pub fn run_cmd(screen_mapping: &ScreenMapping, cfg: &Config) -> Result<Output, CommandError> {
    let args = build_cmd_args(screen_mapping, cfg)?;
    if cfg!(debug_assertions) {
        Ok(Command::new("echo").args(args).output()?)
    } else {
        Ok(Command::new(&args[0]).args(&args[1..]).output()?)
    }
}

fn build_cmd_args(
    screen_mapping: &ScreenMapping,
    cfg: &Config,
) -> Result<Vec<String>, CommandError> {
    let metamode = cfg
        .screens
        .iter()
        .map(|c| c.metamode(screen_mapping))
        .collect::<Result<Vec<String>, CommandError>>()?
        .join(", ");
    Ok(vec![
        "nvidia-settings",
        "--assign",
        &format!("CurrentMetaMode={}", metamode),
    ]
    .iter()
    .map(|s| s.to_string())
    .collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn build_cmd() {
        let cfg = Config {
            screens: vec![
                Screen {
                    name: "DELL U2415".to_string(),
                    resolution: Resolution {
                        width: 1920,
                        height: 1200,
                    },
                    offset: Offset { x: 0, y: 0 },
                    orientation: Orientation::Right,
                    ..Default::default()
                },
                Screen {
                    name: "LG Electronics 34UM95".to_string(),
                    resolution: Resolution {
                        width: 3440,
                        height: 1440,
                    },
                    offset: Offset { x: 1200, y: 0 },
                    ..Default::default()
                },
                Screen {
                    name: "Lenovo Group Limited P27h-20".to_string(),
                    resolution: Resolution {
                        width: 2560,
                        height: 1440,
                    },
                    offset: Offset { x: 4640, y: 0 },
                    ..Default::default()
                },
            ],
        };
        let map = vec![
            ("Lenovo Group Limited P27h-20", "DFP-0.8"),
            ("LG Electronics 34UM95", "DFP-0.1"),
            ("DELL U2415", "DFP-5.8"),
        ]
        .iter()
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .collect();
        assert_eq!(
            build_cmd_args(&map, cfg).expect("failed build_cmd_args"),
            vec![
            "nvidia-settings",
            "--assign",
            "CurrentMetaMode=\
            DFP-5.8: 1920x1200 @1200x1920 +0+0 {ForceCompositionPipeline=On, ViewPortIn=1200x1920, ViewPortOut=1920x1200+0+0, Rotation=270}, \
            DFP-0.1: 3440x1440 @3440x1440 +1200+0 {ForceCompositionPipeline=On, ViewPortIn=3440x1440, ViewPortOut=3440x1440+0+0, Rotation=0}, \
            DFP-0.8: 2560x1440 @2560x1440 +4640+0 {ForceCompositionPipeline=On, ViewPortIn=2560x1440, ViewPortOut=2560x1440+0+0, Rotation=0}"

            ]
        );
    }

    #[test]
    fn parse() {
        let testdir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        let map = screen_mapping_from_xorg_log(&testdir.join("testdata/Xorg.0.log")).unwrap();

        assert_eq!(
            map,
            vec![
                ("Lenovo Group Limited P27h-20", "DFP-0.8"),
                ("LG Electronics 34UM95", "DFP-0.1"),
                ("DELL U2415", "DFP-5.8"),
            ]
            .iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect()
        );
    }
}