Super ultrawide Niri

use std::collections::HashMap;
use std::io;

use niri_ipc::socket::Socket;
use niri_ipc::{Action, Event, Request, Response, Window, Workspace};

#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct WindowId(u64);

#[derive(Clone, Copy, PartialEq, Eq)]
struct WorkspaceId(u64);

struct OutputName<'a>(&'a str);

struct WindowState {
    workspace: Option<WorkspaceId>,
    width: f64,
}

fn main() -> io::Result<()> {
    let mut socket = Socket::connect()?;
    if !matches!(socket.send(Request::EventStream)?, Ok(Response::Handled)) {
        eprintln!("niri rejected event stream");
        std::process::exit(1);
    }

    let mut known: HashMap<WindowId, WindowState> = HashMap::new();
    let mut read_event = socket.read_events();
    loop {
        let result = match read_event()? {
            // A full snapshot of the current state. Just refresh our state.
            Event::WindowsChanged { windows } => {
                known = windows
                    .into_iter()
                    .map(|w| {
                        (
                            WindowId(w.id),
                            WindowState {
                                workspace: w.workspace_id.map(WorkspaceId),
                                width: w.layout.tile_size.0,
                            },
                        )
                    })
                    .collect();
                Ok(())
            }
            Event::WindowOpenedOrChanged { window } => {
                let workspace = window.workspace_id.map(WorkspaceId);
                let entry = WindowState {
                    workspace,
                    width: window.layout.tile_size.0,
                };
                let prev = known.insert(WindowId(window.id), entry);
                match prev {
                    // Don't center floats.
                    _ if window.is_floating => Ok(()),
                    // New window, try to re-center.
                    None => maybe_center_new(&window),
                    // Window changed workspace, try to re-center.
                    Some(state) if state.workspace != workspace => center_focused_if_fits(),
                    // Skip other things.
                    Some(_) => Ok(()),
                }
            }
            Event::WindowClosed { id } => {
                if known.remove(&WindowId(id)).is_some() {
                    center_focused_if_fits()
                } else {
                    Ok(())
                }
            }
            Event::WindowLayoutsChanged { changes } => {
                // Only re-center if the width was changed, otherwise our re-center will
                // loop back indefinitely.
                let mut resized = false;
                for (id, layout) in changes {
                    if let Some(state) = known.get_mut(&WindowId(id)) {
                        if (state.width - layout.tile_size.0).abs() > 0.5 {
                            state.width = layout.tile_size.0;
                            resized = true;
                        }
                    }
                }
                if resized {
                    center_focused_if_fits()
                } else {
                    Ok(())
                }
            }
            _ => Ok(()),
        };

        if let Err(e) = result {
            eprintln!("autocenter: {e}");
        }
    }
}

/// Center a newly created window if the workspace is focused and if there's surrounding free space left.
fn maybe_center_new(window: &Window) -> io::Result<()> {
    let Some(workspace_id) = window.workspace_id.map(WorkspaceId) else {
        return Ok(());
    };
    let Some(focused) = focused_workspace()? else {
        return Ok(());
    };
    if WorkspaceId(focused.id) == workspace_id {
        center_if_fits(&focused)?;
    }
    Ok(())
}

/// Center windows in the focused workspace if there's surrounding free space left.
fn center_focused_if_fits() -> io::Result<()> {
    if let Some(focused) = focused_workspace()? {
        center_if_fits(&focused)?;
    }
    Ok(())
}

/// Center windows in the workspace if there's surrounding free space left.
fn center_if_fits(workspace: &Workspace) -> io::Result<()> {
    let Some(output) = workspace.output.as_deref().map(OutputName) else {
        return Ok(());
    };
    let Some(width) = output_width(output)? else {
        return Ok(());
    };
    if workspace_width(WorkspaceId(workspace.id))? < f64::from(width) {
        center_visible_columns()?;
    }
    Ok(())
}

/// Issue a one-shot query to Niri, wait, and return the response.
fn query(request: Request) -> io::Result<Response> {
    match Socket::connect()?.send(request)? {
        Ok(response) => Ok(response),
        Err(msg) => Err(io::Error::other(msg)),
    }
}

/// Get the focused workspace.
fn focused_workspace() -> io::Result<Option<Workspace>> {
    match query(Request::Workspaces)? {
        Response::Workspaces(ws) => Ok(ws.into_iter().find(|w| w.is_focused)),
        _ => Ok(None),
    }
}

/// Get the width of an output (monitor).
fn output_width(name: OutputName<'_>) -> io::Result<Option<u32>> {
    match query(Request::Outputs)? {
        Response::Outputs(outputs) => Ok(outputs
            .get(name.0)
            .and_then(|o| o.logical.as_ref())
            .map(|l| l.width)),
        _ => Ok(None),
    }
}

/// Calculates the width of all columns in the workspace.
fn workspace_width(workspace_id: WorkspaceId) -> io::Result<f64> {
    let Response::Windows(windows) = query(Request::Windows)? else {
        return Ok(0.0);
    };
    let mut columns: HashMap<usize, f64> = HashMap::new();
    for w in windows {
        if w.workspace_id.map(WorkspaceId) != Some(workspace_id) {
            continue;
        }
        if let Some((col, _)) = w.layout.pos_in_scrolling_layout {
            let width = columns.entry(col).or_insert(0.0);
            *width = width.max(w.layout.tile_size.0);
        }
    }
    Ok(columns.values().sum())
}

/// Send a command to center the visible columns.
fn center_visible_columns() -> io::Result<()> {
    if let Err(msg) = Socket::connect()?.send(Request::Action(Action::CenterVisibleColumns {}))? {
        eprintln!("center-visible-columns rejected: {msg}");
    }
    Ok(())
}