swaysome/src/main.rs

extern crate byteorder;
extern crate clap;
extern crate serde_json;

use serde::{Deserialize, Serialize};

use clap::{Args, Parser, Subcommand};
use std::env;
use std::io::Cursor;
use std::io::{Read, Write};
use std::mem;
use std::os::unix::net::UnixStream;
use std::path::Path;

use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};

const RUN_COMMAND: u32 = 0;
const GET_WORKSPACES: u32 = 1;
// const SUBSCRIBE: u32 = 2;
const GET_OUTPUTS: u32 = 3;

#[derive(Parser, Debug)]
#[clap(author, version, about = "Better multimonitor handling for sway", long_about = None)]
#[clap(propagate_version = true)]
struct Cli {
    #[clap(subcommand)]
    command: Command,
}

#[derive(Subcommand, Debug)]
enum Command {
    #[clap(about = "Initialize the workspace groups for all the outputs")]
    Init(InitAction),

    #[clap(about = "Move the focused container to another workspace on the same workspace group")]
    Move(MoveAction),

    #[clap(about = "Focus to another workspace on the same workspace group")]
    Focus(FocusAction),

    #[clap(about = "Focus to workspace group")]
    FocusGroup(FocusAction),

    #[clap(about = "Focus to another workspace on all the outputs")]
    FocusAllOutputs(FocusAction),

    #[clap(about = "Move the focused container to the next output")]
    NextOutput,

    #[clap(about = "Move the focused container to the previous output")]
    PrevOutput,

    #[clap(about = "Move the focused container to the next group")]
    NextGroup,

    #[clap(about = "Move the focused container to the previous group")]
    PrevGroup,

    #[clap(about = "Rearrange already opened workspaces")]
    RearrangeWorkspaces,
}

#[derive(Args, Debug)]
struct InitAction {
    #[clap(value_name = "index", help = "The index to initialize with")]
    index: usize,
}

#[derive(Args, Debug)]
struct FocusAction {
    #[clap(value_name = "index", help = "The index to focus on")]
    index: usize,
}

#[derive(Args, Debug)]
struct MoveAction {
    #[clap(value_name = "index", help = "The index to move the container to")]
    index: usize,
}

fn get_stream() -> UnixStream {
    let socket_path = match env::var("I3SOCK") {
        Ok(val) => val,
        Err(_e) => {
            panic!("couldn't find i3/sway socket");
        }
    };

    let socket = Path::new(&socket_path);

    match UnixStream::connect(&socket) {
        Err(_) => panic!("couldn't connect to i3/sway socket"),
        Ok(stream) => stream,
    }
}

fn send_msg(mut stream: &UnixStream, msg_type: u32, payload: &str) {
    let payload_length = payload.len() as u32;

    let mut msg_prefix: [u8; 6 * mem::size_of::<u8>() + 2 * mem::size_of::<u32>()] =
        *b"i3-ipc00000000";

    msg_prefix[6..]
        .as_mut()
        .write_u32::<LittleEndian>(payload_length)
        .expect("Unable to write");

    msg_prefix[10..]
        .as_mut()
        .write_u32::<LittleEndian>(msg_type)
        .expect("Unable to write");

    let mut msg: Vec<u8> = msg_prefix[..].to_vec();
    msg.extend(payload.as_bytes());

    if stream.write_all(&msg[..]).is_err() {
        panic!("couldn't send message");
    }
}

fn send_command(stream: &UnixStream, command: &str) {
    eprint!("Sending command: '{}' - ", &command);
    send_msg(stream, RUN_COMMAND, command);
    check_success(stream);
}

fn read_msg(mut stream: &UnixStream) -> Result<String, &str> {
    let mut response_header: [u8; 14] = *b"uninitialized.";
    stream.read_exact(&mut response_header).unwrap();

    if &response_header[0..6] == b"i3-ipc" {
        let mut v = Cursor::new(vec![
            response_header[6],
            response_header[7],
            response_header[8],
            response_header[9],
        ]);
        let payload_length = v.read_u32::<LittleEndian>().unwrap();

        let mut payload = vec![0; payload_length as usize];
        stream.read_exact(&mut payload[..]).unwrap();
        let payload_str = String::from_utf8(payload).unwrap();
        Ok(payload_str)
    } else {
        eprint!("Not an i3-icp packet, emptying the buffer: ");
        let mut v = vec![];
        stream.read_to_end(&mut v).unwrap();
        eprintln!("{:?}", v);
        Err("Unable to read i3-ipc packet")
    }
}

fn check_success(stream: &UnixStream) {
    match read_msg(stream) {
        Ok(msg) => {
            let r: Vec<serde_json::Value> = serde_json::from_str(&msg).unwrap();
            match r[0]["success"] {
                serde_json::Value::Bool(true) => eprintln!("Command successful"),
                _ => panic!("Command failed: {:#?}", r),
            }
        }
        Err(_) => panic!("Unable to read response"),
    };
}

#[derive(Serialize, Deserialize, Debug)]
struct Output {
    name: String,
    #[serde(default)]
    focused: bool,
    active: bool,
}

fn get_outputs(stream: &UnixStream) -> Vec<Output> {
    send_msg(stream, GET_OUTPUTS, "");
    let o = match read_msg(stream) {
        Ok(msg) => msg,
        Err(_) => panic!("Unable to get outputs"),
    };
    let mut outputs: Vec<Output> = serde_json::from_str(&o).unwrap();
    outputs.sort_by(|x, y| x.name.cmp(&y.name)); // sort_by_key doesn't work here (https://stackoverflow.com/a/47126516)
    outputs
}

#[derive(Serialize, Deserialize, Debug)]
struct Workspace {
    num: usize,
    output: String,
    visible: bool,
}

fn get_workspaces(stream: &UnixStream) -> Vec<Workspace> {
    send_msg(stream, GET_WORKSPACES, "");
    let ws = match read_msg(stream) {
        Ok(msg) => msg,
        Err(_) => panic!("Unable to get current workspace"),
    };
    let mut workspaces: Vec<Workspace> = serde_json::from_str(&ws).unwrap();
    workspaces.sort_by_key(|x| x.num);
    workspaces
}

fn get_current_output_index(stream: &UnixStream) -> usize {
    let outputs = get_outputs(stream);

    match outputs.iter().position(|x| x.focused) {
        Some(i) => i + 1,
        None => panic!("WTF! No focused output???"),
    }
}

fn get_current_output_name(stream: &UnixStream) -> String {
    let outputs = get_outputs(stream);

    let focused_output_index = match outputs.iter().find(|x| x.focused) {
        Some(i) => i.name.as_str(),
        None => panic!("WTF! No focused output???"),
    };

    focused_output_index.to_string()
}

fn get_current_workspace(stream: &UnixStream) -> Workspace {
    let outputs = get_outputs(stream);
    let workspaces = get_workspaces(stream);
    workspaces
        .into_iter()
        .find(|w| w.visible && outputs.iter().find(|o| o.name == w.output).unwrap().focused)
        .unwrap()
}

fn move_container_to_workspace(stream: &UnixStream, workspace_index: usize) {
    if workspace_index < 10 {
        move_container_to_workspace_relative(stream, workspace_index);
    } else {
        move_container_to_workspace_absolute(stream, workspace_index);
    }
}

fn move_container_to_workspace_absolute(stream: &UnixStream, workspace_index: usize) {
    let output_index = (workspace_index / 10) as usize;
    let outputs = get_outputs(stream);
    let workspaces = get_workspaces(stream);

    // If the workspace already exists
    match workspaces.iter().find(|w| w.num == workspace_index) {
        Some(_) => {
            let mut focus_cmd: String = "move container to workspace number ".to_string();
            focus_cmd.push_str(&workspace_index.to_string());
            send_command(stream, &focus_cmd);
        }
        None => {
            let target_group = workspace_index / 10;
            let target_screen_index = match workspaces.iter().find(|w| w.num / 10 == target_group) {
                // If other workspaces on the same group exists
                Some(other_workspace) => Some(
                    outputs
                        .iter()
                        .enumerate()
                        .find(|i| i.1.name == other_workspace.output)
                        .unwrap()
                        .0,
                ),
                None => {
                    // Or if the targeted output is currently connected
                    if output_index < outputs.len() {
                        Some(output_index)
                    } else {
                        None
                    }
                }
            };

            match target_screen_index {
                // If we have to send it to another screen
                Some(target_screen_index) => {
                    let target_output = &outputs[target_screen_index - 1];

                    let current_output_name = get_current_output_name(stream);

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&target_output.name);
                    send_command(stream, &focus_cmd);

                    let focused_workspace_index = get_current_workspace(stream).num;

                    let mut focus_cmd: String = "workspace ".to_string();
                    focus_cmd.push_str(&workspace_index.to_string());
                    send_command(stream, &focus_cmd);

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&current_output_name);
                    send_command(stream, &focus_cmd);

                    let mut focus_cmd: String = "move container to workspace ".to_string();
                    focus_cmd.push_str(&workspace_index.to_string());
                    send_command(stream, &focus_cmd);

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&target_output.name);
                    send_command(stream, &focus_cmd);

                    let mut focus_cmd: String = "workspace ".to_string();
                    focus_cmd.push_str(&focused_workspace_index.to_string());
                    send_command(stream, &focus_cmd);

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&current_output_name);
                    send_command(stream, &focus_cmd);
                }
                None => {
                    // Else, we send the container on the current output
                    let mut focus_cmd: String = "move container to workspace ".to_string();
                    focus_cmd.push_str(&workspace_index.to_string());
                    send_command(stream, &focus_cmd);
                }
            };
        }
    }
}

fn move_container_to_workspace_relative(stream: &UnixStream, workspace_index: usize) {
    let current_workspace_index: usize = get_current_workspace(stream).num;
    let focused_output_index = current_workspace_index / 10;

    let mut cmd: String = "move container to workspace number ".to_string();
    let full_ws_name = format!("{}{}", focused_output_index, workspace_index)
        .parse::<i32>()
        .unwrap();
    cmd.push_str(&full_ws_name.to_string());
    send_command(stream, &cmd);
}

fn focus_to_workspace(stream: &UnixStream, workspace_index: usize) {
    if workspace_index < 10 {
        focus_to_workspace_relative(stream, workspace_index);
    } else {
        focus_to_workspace_absolute(stream, workspace_index);
    }
}

fn focus_to_workspace_absolute(stream: &UnixStream, workspace_index: usize) {
    let output_index = (workspace_index / 10) as usize;
    let outputs = get_outputs(stream);
    let workspaces = get_workspaces(stream);

    // If the workspace already exists
    match workspaces.iter().find(|w| w.num == workspace_index) {
        Some(_) => {
            let mut focus_cmd: String = "workspace number ".to_string();
            focus_cmd.push_str(&workspace_index.to_string());
            send_command(stream, &focus_cmd);
        }
        None => {
            let target_group = workspace_index / 10;
            let target_screen_index = match workspaces.iter().find(|w| w.num / 10 == target_group) {
                // If other workspaces on the same group exists
                Some(other_workspace) => Some(
                    outputs
                        .iter()
                        .enumerate()
                        .find(|i| i.1.name == other_workspace.output)
                        .unwrap()
                        .0,
                ),
                None => {
                    // Or if the targeted output is currently connected
                    if output_index < outputs.len() {
                        Some(output_index)
                    } else {
                        None
                    }
                }
            };

            match target_screen_index {
                // If we have to send it to another screen
                Some(target_screen_index) => {
                    let target_output = &outputs[target_screen_index - 1];

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&target_output.name);
                    send_command(stream, &focus_cmd);
                }
                None => {}
            };
            // Then we focus the workspace
            let mut focus_cmd: String = "workspace number ".to_string();
            focus_cmd.push_str(&workspace_index.to_string());
            send_command(stream, &focus_cmd);
        }
    }
}

fn focus_to_workspace_relative(stream: &UnixStream, workspace_index: usize) {
    let current_workspace_index: usize = get_current_workspace(stream).num;
    let focused_output_index = current_workspace_index / 10;

    let mut cmd: String = "workspace number ".to_string();
    let full_ws_name = format!("{}{}", focused_output_index, &workspace_index)
        .parse::<i32>()
        .unwrap();
    cmd.push_str(&full_ws_name.to_string());
    send_command(stream, &cmd);
}

fn focus_to_group(stream: &UnixStream, group_index: usize) {
    let outputs = get_outputs(stream);
    let workspaces = get_workspaces(stream);

    let current_workspace_index: usize = get_current_workspace(stream).num;
    let contextual_target_workspace_index =
        current_workspace_index - (current_workspace_index / 10) * 10;

    let target_workspace_index = group_index * 10 + contextual_target_workspace_index;

    // If the workspace already exists
    match workspaces.iter().find(|w| w.num == target_workspace_index) {
        Some(_) => {
            let mut focus_cmd: String = "workspace number ".to_string();
            focus_cmd.push_str(&target_workspace_index.to_string());
            send_command(stream, &focus_cmd);
        }
        None => {
            let target_screen_index = match workspaces.iter().find(|w| w.num / 10 == group_index) {
                // If other workspaces on the same group exists
                Some(other_workspace) => Some(
                    outputs
                        .iter()
                        .enumerate()
                        .find(|i| i.1.name == other_workspace.output)
                        .unwrap()
                        .0
                        + 1,
                ),
                None => {
                    // Or if the targeted output is currently connected
                    if group_index <= outputs.len() {
                        Some(group_index)
                    } else {
                        None
                    }
                }
            };

            match target_screen_index {
                // If we have to send it to another screen
                Some(target_screen_index) => {
                    let target_output = &outputs[target_screen_index - 1];

                    let mut focus_cmd: String = "focus output ".to_string();
                    focus_cmd.push_str(&target_output.name);
                    send_command(stream, &focus_cmd);
                }
                None => {}
            };
            // Then we focus the workspace
            let mut focus_cmd: String = "workspace number ".to_string();
            focus_cmd.push_str(&target_workspace_index.to_string());
            send_command(stream, &focus_cmd);
        }
    }
}

fn focus_all_outputs_to_workspace(stream: &UnixStream, workspace_index: usize) {
    let current_output = get_current_output_name(stream);

    // Iterate on all outputs to focus on the given workspace
    let outputs = get_outputs(stream);
    for output in outputs.iter() {
        let mut cmd: String = "focus output ".to_string();
        cmd.push_str(output.name.as_str());
        send_command(stream, &cmd);

        focus_to_workspace(stream, workspace_index);
    }

    // Get back to currently focused output
    let mut cmd: String = "focus output ".to_string();
    cmd.push_str(&current_output);
    send_command(stream, &cmd);
}

fn move_container_to_next_output(stream: &UnixStream) {
    move_container_to_next_or_prev_output(stream, false);
}

fn move_container_to_prev_output(stream: &UnixStream) {
    move_container_to_next_or_prev_output(stream, true);
}

fn move_container_to_next_or_prev_output(stream: &UnixStream, go_to_prev: bool) {
    let outputs = get_outputs(stream);
    let focused_output_index = get_current_output_index(stream);

    let target_output = if go_to_prev {
        &outputs[(focused_output_index + outputs.len() - 1) % outputs.len() - 1]
    } else {
        &outputs[(focused_output_index + 1) % outputs.len() - 1]
    };

    let workspaces = get_workspaces(stream);
    let target_workspace = workspaces
        .iter()
        .find(|x| x.output == target_output.name && x.visible)
        .unwrap();

    // Move container to target workspace
    let mut cmd: String = "move container to workspace number ".to_string();
    cmd.push_str(&target_workspace.num.to_string());
    send_command(stream, &cmd);

    // Focus that workspace to follow the container
    let mut cmd: String = "workspace number ".to_string();
    cmd.push_str(&target_workspace.num.to_string());
    send_command(stream, &cmd);
}

fn focus_container_to_next_group(stream: &UnixStream) {
    focus_container_to_next_or_prev_group(stream, false);
}

fn focus_container_to_prev_group(stream: &UnixStream) {
    focus_container_to_next_or_prev_group(stream, true);
}

fn focus_container_to_next_or_prev_group(stream: &UnixStream, go_to_prev: bool) {
    let current_workspace_index: usize = get_current_workspace(stream).num;
    let focused_group_index = current_workspace_index / 10;

    if go_to_prev {
        focus_to_group(stream, focused_group_index - 1);
    } else {
        focus_to_group(stream, focused_group_index + 1);
    };
}

fn init_workspaces(stream: &UnixStream, workspace_index: usize) {
    let outputs = get_outputs(stream);

    let cmd_prefix: String = "focus output ".to_string();
    for output in outputs.iter().filter(|x| x.active).rev() {
        let mut cmd = cmd_prefix.clone();
        cmd.push_str(output.name.as_str());
        println!("{:?}", cmd.clone());
        send_command(stream, &cmd);

        let mut cmd: String = "workspace number ".to_string();
        let full_ws_name = format!("{}{}", get_current_output_index(stream), &workspace_index)
            .parse::<i32>()
            .unwrap();
        cmd.push_str(&full_ws_name.to_string());
        send_command(stream, &cmd);
    }
}

fn rearrange_workspaces(stream: &UnixStream) {
    let outputs = get_outputs(stream);
    let workspaces = get_workspaces(stream);

    let focus_cmd_prefix: String = "workspace number ".to_string();
    let move_cmd_prefix: String = "move workspace to ".to_string();
    for workspace in workspaces.iter() {
        let mut focus_cmd = focus_cmd_prefix.clone();
        focus_cmd.push_str(&workspace.num.to_string());
        send_command(stream, &focus_cmd);

        let output_index = workspace.num / 10;
        if output_index <= outputs.len() - 1 {
            let mut move_cmd = move_cmd_prefix.clone();
            move_cmd.push_str(&outputs[output_index.max(1) - 1].name);
            send_command(stream, &move_cmd);
        }
    }
}

fn main() {
    let cli = Cli::parse();
    let stream = get_stream();

    match &cli.command {
        Command::Init(action) => {
            init_workspaces(&stream, action.index);
        }
        Command::Move(action) => {
            move_container_to_workspace(&stream, action.index);
        }
        Command::Focus(action) => {
            focus_to_workspace(&stream, action.index);
        }
        Command::FocusGroup(action) => {
            focus_to_group(&stream, action.index);
        }
        Command::FocusAllOutputs(action) => {
            focus_all_outputs_to_workspace(&stream, action.index);
        }
        Command::NextOutput => {
            move_container_to_next_output(&stream);
        }
        Command::PrevOutput => {
            move_container_to_prev_output(&stream);
        }
        Command::NextGroup => {
            focus_container_to_next_group(&stream);
        }
        Command::PrevGroup => {
            focus_container_to_prev_group(&stream);
        }
        Command::RearrangeWorkspaces => {
            rearrange_workspaces(&stream);
        }
    }
}