//! Shows dmenu with a list of windows of all workspaces.
//! Windows can be focused by selecting entries.
//!
//! All command line arguments are forwarded to dmenu.
//!
//! Node: dmenu supports selection of entries without closing
//! using Ctrl+Return. But because the output of dmenu is buffered by default,
//! this does not work correctly. A workaround is using a shell script
//! to start dmenu:
//!
//! ```
//!  #!/bin/sh
//!  stdbuf -o 0 dmenu /usr/bin/dmenu "$@"
//! ```

#![feature(old_io, exit_status)]

extern crate "rustc-serialize" as serialize;
extern crate i3;

use std::env;
use std::old_io as io;
use std::old_io::BufferedReader;
use std::old_io::process;
use std::old_io::process::Command;
use std::collections::BTreeMap;
use serialize::json::Json;

use i3::ipc::{I3Proto, get_socket_path, open_socket};
use i3::tree::{iter_nodes, get_focus, get_conid,
               get_node_string, num_children, get_type};

/// Creates a mapping between titles and container ids for all windows.
fn get_windows(tree: &Json) -> BTreeMap<String, i64> {
    let mut windows = BTreeMap::new();
    let mut insts = BTreeMap::<String, usize>::new();
    // List workspaces.
    for (n, wsname) in iter_nodes(tree)
            .filter(|&(n, _)| get_type(n).unwrap_or("") == "workspace")
            .map(|(n, _)| (n, get_node_string(n, "name").unwrap_or(""))) {
        // List windows of the workspace.
        for (n, _) in iter_nodes(n).filter(|&(n, d)| d > 0 && num_children(n) == 0) {
            let conid = match get_conid(n) {
                Some(conid) => conid,
                // Skip when node has no id.
                None => { continue; }
            };
            let mut key = format!("[{}] {}", wsname,
                                  get_node_string(n, "name").unwrap_or(""));
            // Disambiguate titles.
            if match insts.get_mut(&key) {
                Some(n) => {
                    // Append sequential number.
                    *n += 1;
                    key = format!("{} ({})", key, *n);
                    false
                }
                None => true
            } {
                insts.insert(key.clone(), 1);
            }
            windows.insert(key, conid);
        }
    }
    windows
}

/// Focuses a window by container id.
fn focus_con(sock: &mut I3Proto, conid: i64) -> io::IoResult<()> {
    let cmd = format!("[con_id={}] focus", conid);
    try!(sock.request_command(&cmd[..]));
    Ok(())
}

/// Starts dmenu with the windows found by `get_windows`
/// and handles the selection.
fn winmenu(sock: &mut I3Proto, windows: &BTreeMap<String, i64>,
           args: &[String], previous: Option<i64>) -> bool
{
    let mut dmenu = Command::new("dmenu").args(args).stderr(process::InheritFd(2)).spawn().unwrap();
    // Send the window titles to dmenu's stdin.
    let mut p = dmenu.stdin.take().unwrap();
    for win in windows.keys() {
        p.write_str(&win[..]).unwrap();
        p.write_str("\n").unwrap();
    }
    drop(p);
    // Handle the output.
    let mut read = BufferedReader::with_capacity(512, dmenu.stdout.take().unwrap());
    let mut cur_id = previous; // current focus
    for line in read.lines() {
        let line = match line {
            Ok(l) => l,
            Err(e) => { println!("{}", e); continue; }
        };
        let line = line.trim_right_matches('\n');
        if line.len() == 0 { continue; }
        // Try to find the selection.
        match windows.get(line) {
            Some(&conid) => {
                println!("selected: {}", conid);
                cur_id = Some(conid);
                focus_con(sock, conid).unwrap();
            }
            None => {
                println!("window `{}` not found", line);
            }
        }
    }
    // Restore `previous` focus if dmenu was aborted.
    let success = match dmenu.wait() {
        Ok(s) => s.success(),
        Err(e) => {
            println!("{}", e);
            false
        }
    };
    if !success {
        match previous {
            Some(conid) => {
                if cur_id != previous {
                    println!("dmenu aborted, restoring focus");
                    focus_con(sock, conid).unwrap()
                } else {
                    println!("dmenu aborted, focus unchanged");
                }
            }
            None => {
                println!("dmenu aborted");
            }
        }
    }
    success
}

fn main() {
    // Open the socket and get the layout tree.
    let path = get_socket_path().unwrap();
    let mut sock = open_socket(&path).unwrap();
    let tree = sock.request_tree().unwrap();

    // Save the initially selected window.
    let previous = get_focus(&tree).and_then(|n| get_conid(n));
    // Find all windows.
    let windows = get_windows(&tree);
    // Start dmenu. Forward all command line arguments.
    if !winmenu(&mut sock, &windows, &env::args().skip(1).collect::<Vec<_>>(), previous) {
        env::set_exit_status(1);
    }
}