niri-tag/daemon/manager.rs

use crate::socket::{create_niri_socket, query, tell};
use anyhow::{Result, anyhow};
use niri_ipc::{
    Action, Event, Reply, Request, Response, Window, WorkspaceReferenceArg,
    state::{EventStreamState, EventStreamStatePart},
};
use niri_tag::TagCmd;
use smol::{
    channel::{self},
    future,
    io::BufReader,
    net::unix::UnixStream,
};
use std::collections::HashMap;

pub struct NiriTag {
    tags: HashMap<u8, bool>,
    windows: HashMap<u64, u8>,
}

impl NiriTag {
    pub fn new() -> Self {
        Self {
            tags: HashMap::new(),
            windows: HashMap::new(),
        }
    }

    #[allow(unreachable_code)]
    pub async fn manage_tags(
        self,
        ev_rx: channel::Receiver<Event>,
        tag_rx: channel::Receiver<TagCmd>,
    ) -> Result<()> {
        let mut state = EventStreamState::default();
        let mut tags = NiriTag::new();
        let mut socket = create_niri_socket().await?;
        // base tag is always visible
        tags.tags.insert(0, true);

        async fn on_focused_window<O>(socket: &mut BufReader<UnixStream>, mut ok: O) -> Result<()>
        where
            O: FnMut(Option<Window>) -> Result<()>,
        {
            let q = query(socket, Request::FocusedWindow).await?;
            if let Reply::Ok(Response::FocusedWindow(win)) = q {
                ok(win)
            } else {
                Err(anyhow!(
                    "Invalid response from Niri when requesting FocusedWindow: {}",
                    if q.is_err() {
                        q.unwrap_err()
                    } else {
                        serde_json::to_string(&q.unwrap())?
                    }
                ))
            }
        }

        loop {
            let recvd: Receivable =
                future::or(async { ev_rx.recv().await.map(Receivable::Event) }, async {
                    tag_rx.recv().await.map(Receivable::TagCmd)
                })
                .await?;
            tracing::debug!("received {:?}", recvd);

            let res = match recvd {
                Receivable::Event(ev) => {
                    let _ = state.apply(ev.clone());
                    Ok(())
                }
                Receivable::TagCmd(cmd) => match cmd {
                    TagCmd::Add(t) => {
                        on_focused_window(&mut socket, |win| {
                            if let Some(win) = win {
                                let wid = win.id;
                                tags.windows.insert(wid, t);
                                Ok(())
                            } else {
                                Err(anyhow!("No focused window to tag"))
                            }
                        })
                        .await
                    }
                    TagCmd::Remove(_) => {
                        on_focused_window(&mut socket, |win| {
                            if let Some(win) = win {
                                let wid = win.id;
                                tags.windows.insert(wid, 0);
                                Ok(())
                            } else {
                                Err(anyhow!("No focused window to untag"))
                            }
                        })
                        .await
                    }
                    TagCmd::Toggle(t) => {
                        on_focused_window(&mut socket, |win| {
                            if let Some(win) = win {
                                let wid = win.id;
                                let toggle = if *tags.windows.get(&wid).unwrap_or(&0) == t {
                                    0
                                } else {
                                    t
                                };
                                tracing::debug!("toggling {} to tag {}", wid, toggle);
                                tags.windows.insert(wid, toggle);
                                Ok(())
                            } else {
                                Err(anyhow!("No focused window to untag"))
                            }
                        })
                        .await
                    }

                    TagCmd::Enable(t) => {
                        tags.tags.insert(t, true);
                        Ok(())
                    }
                    TagCmd::Disable(t) => {
                        tags.tags.insert(t, false);
                        Ok(())
                    }
                    TagCmd::ToggleWs(t) => {
                        tracing::debug!("toggling tag {}", t);
                        tags.tags.insert(t, !tags.tags.get(&t).unwrap_or(&false));
                        Ok(())
                    }
                },
            };
            match res {
                Ok(()) => (),
                Err(e) => tracing::error!("error occurred in manager loop: {}", e),
            }
            // TODO: react selectively instead of brute forcing window state
            // use Event::*;
            // match ev {
            //     WorkspaceActivated { .. } => (),
            //     WorkspacesChanged { .. } => (),
            //     WorkspaceUrgencyChanged { .. } => (),
            //     WindowsChanged { .. } => (),
            //     WindowOpenedOrChanged { .. } => (),
            //     WindowUrgencyChanged { .. } => (),
            //     WindowClosed { .. } => (),
            //     _ => (),
            // }
            for (&wid, window) in state.windows.windows.iter() {
                let (active, inactive): (Vec<_>, Vec<_>) = state
                    .workspaces
                    .workspaces
                    .iter()
                    .map(|(wsid, ws)| (wsid, ws.is_active))
                    .partition(|(_, a)| *a);
                if let Some(wsid) = window.workspace_id {
                    if let Some(&window_tag) = tags.windows.get(&wid) {
                        if let Some(&tag_enabled) = tags.tags.get(&window_tag) {
                            if tag_enabled && inactive.contains(&(&wsid, false)) {
                                tell(
                                    &mut socket,
                                    Request::Action(Action::MoveWindowToWorkspace {
                                        window_id: Some(wid),
                                        reference: WorkspaceReferenceArg::Index(0),
                                        focus: false,
                                    }),
                                )
                                .await?;
                                tracing::debug!("making visible {}", wid);
                            } else if !tag_enabled && active.contains(&(&wsid, true)) {
                                let hidden = *inactive.first().unwrap().0;
                                tell(
                                    &mut socket,
                                    Request::Action(Action::MoveWindowToWorkspace {
                                        window_id: Some(wid),
                                        reference: WorkspaceReferenceArg::Id(hidden),
                                        focus: false,
                                    }),
                                )
                                .await?;
                                tracing::debug!("making hidden {}", wid);
                            }
                        } else {
                            tags.windows.insert(wid, 0);
                        }
                    }
                }
            }
        }
        tracing::error!("Manager loop ended");
        unreachable!("Manager loop ended");
        Ok(())
    }
}

#[derive(Debug)]
enum Receivable {
    Event(Event),
    TagCmd(TagCmd),
}