yabai.rs

use std::string::String;

use serde::de::DeserializeOwned;
use serde::Deserialize;
use smallvec::SmallVec;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixStream;
use tokio::task::JoinHandle;

pub const QWW: &str = "query\0--windows\0--window\0\0";
// pub const QSS: &str = "query\0--spaces\0--space\0\0";
pub const STACKED_BORDER: &str = "config\x00active_window_border_color\x000xffdab061\0\0";
pub const UNSTACKED_BORDER: &str = "config\x00active_window_border_color\x000xff93abbc\0\0";
pub const UNFOCUSED_BORDER: &str = "config\x00active_window_border_color\x000xff555555\0\0";
pub const ADDR: &str = "/tmp/yabai_adi.socket";
pub const Y3_ADDR: &str = "/tmp/y3_rust.socket";

#[derive(thiserror::Error, Debug)]
pub enum CmdError {
    #[error("Command failed: {0}")]
    CommandFailed(String),
    #[error(transparent)]
    Other(#[from] std::io::Error),
}

#[inline]
async fn connect_write(args: &[&str]) -> Result<UnixStream, CmdError> {
    let mut socket = UnixStream::connect(ADDR).await?;
    let mut buf: SmallVec<[u8; 63]> = SmallVec::new();
    for arg in args {
        buf.extend_from_slice(arg.as_bytes());
        buf.extend_from_slice(&[0]);
    }
    buf.extend_from_slice(&[0]);
    socket.write_all(&buf).await?;
    Ok(socket)
}

#[inline]
pub async fn send(args: &[&str]) -> Result<(), CmdError> {
    let mut socket = connect_write(args).await?;
    // let mut buf = [0; 63];
    let mut buf = String::new();
    let n = socket.read_to_string(&mut buf).await?;
    if n > 0 {
        return Err(CmdError::CommandFailed(buf));
        // return Err(CmdError::CommandFailed(from_utf8(&buf[1..n - 1]).unwrap().to_string()));
    }
    Ok(())
}

#[inline]
pub async fn send_async(args: &[&str]) -> Result<JoinHandle<Result<(), CmdError>>, CmdError> {
    let mut socket = connect_write(args).await?;
    // Decouple reading the reply, which only happens on a failed command and adds up to 50ms tail latency
    let err_handle = tokio::spawn(async move {
        // let mut buf = [0; 63];
        let mut buf = String::new();
        let n = socket.read_to_string(&mut buf).await?;
        if n > 0 {
            return Err(CmdError::CommandFailed(buf));
            // return Err(CmdError::CommandFailed(from_utf8(&buf[1..n - 1]).unwrap().to_compact_string()));
        }
        Ok(())
    });
    Ok(err_handle)
}

// TODO: do I need query trait for yabai data structures?
#[inline]
async fn query<T: DeserializeOwned>(args: &[&str]) -> anyhow::Result<T> {
    // println!("running query");
    let mut socket = connect_write(args).await?;
    let mut buf = vec![0; 4096];
    let n = socket.read(&mut buf).await?;
    let v = serde_json::from_reader(&buf[..n])?;
    Ok(v)
}

#[derive(Clone, Deserialize, Debug, Default)]
#[serde(rename_all = "kebab-case")]
pub struct Window {
    pub(crate) id: u32,
    pub(crate) stack_index: u8,
    pub(crate) frame: Frame,
    has_focus: bool,
    #[serde(skip)]
    pub(crate) id_below: u32,
    // #[serde(skip)]
    // input_dir: String,
}

impl Window {
    pub async fn get_source() -> anyhow::Result<Self> {
        let windows: Vec<Self> = query(&["query", "--windows", "--space"]).await?;
        let mut source_w = Self::default();
        for w in windows.clone() {
            if w.has_focus {
                source_w = w;
                break;
            }
        }
        for w in windows {
            if w.frame == source_w.frame && w.stack_index == (source_w.stack_index - 1) {
                source_w.id_below = w.id;
                break;
            }
        }
        Ok(source_w)
    }

    pub async fn get_target(input_dir: &str) -> anyhow::Result<Self> {
        let w: Self = query(&["query", "--windows", "--window", input_dir]).await?;
        Ok(w)
    }

    /*
    pub async fn rotate(input_dir: &str) -> anyhow::Result<()> {
        let s = Space::get_source().await?;
        if self.id == s.first_window {
            insert_warp("east", input_dir).await?;
        } else if self.sw_id == s.last_window {
            insert_warp("west", input_dir).await?;
        }
        Ok(())
    }
    */

    /*
    async fn unstack_async(&self, input_dir: &str) -> anyhow::Result<()> {
        let h1 = send_async(&["window", &self.id.to_compact_string(), "--toggle", "float"]).await?;
        let h2 = send_async(&["window", &self.id_below.to_compact_string(), "--insert", input_dir,]).await?;
        let h3 = send_async(&["window", &self.id.to_compact_string(), "--toggle", "float"]).await?;
        let h4 = send_async(&[UNSTACKED_BORDER]).await?;
        h1.await.unwrap()?;
        h2.await.unwrap()?;
        h3.await.unwrap()?;
        h4.await.unwrap()?;
        Ok(())
    }
    */
}

#[derive(Clone, Deserialize, Debug, Default, PartialEq)]
pub struct Frame {
    x: f32,
    pub(crate) y: f32,
    pub(crate) w: f32,
    pub(crate) h: f32,
}

#[derive(Deserialize, Debug, Default)]
pub struct Display {
    pub(crate) id: u8,
    pub(crate) frame: Frame,
}

impl Display {
    pub async fn get_source() -> anyhow::Result<Self> {
        let d: Self = query(&["query", "--displays", "--display"]).await?;
        Ok(d)
    }

    pub async fn get_target(input_dir: &str) -> anyhow::Result<Self> {
        let d: Self = query(&["query", "--displays", "--display", input_dir])
            .await
            .unwrap_or_default();
        Ok(d)
    }
}

#[derive(Deserialize, Debug, Default)]
#[serde(rename_all = "kebab-case")]
pub struct Space {
    pub(crate) is_visible: bool,
    pub(crate) first_window: u32,
    pub(crate) last_window: u32,
}

impl Space {
    pub async fn get_source() -> anyhow::Result<Self> {
        let s: Self = query(&["query", "--spaces", "--space"]).await?;
        Ok(s)
    }

    pub async fn get_target_all(input_dir: &str) -> anyhow::Result<Vec<Self>> {
        let spaces: Vec<Self> = query(&["query", "--spaces", "--display", input_dir]).await?;
        Ok(spaces)
    }

    /*
    async fn rotate(input_dir: &str) -> anyhow::Result<()> {
    let s = Space::get_source().await?;
    if sw_id == s.first_window {
    insert_warp("east", input_dir).await?;
    } else if self.sw_id == s.last_window {
    insert_warp("west", input_dir).await?;
    }
    Ok(())
    }
    */
}