rs_peekaboo 0.3.4

use crate::automation::{Target, parse_point, split_keys};
use crate::cache;
use crate::{
    Bounds, CommandResult, ComputerUseMode, Direction, ImageMode, Peekaboo, PeekabooConfig, Result,
    UiNode, VisionDetections,
};
use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};
use clap_complete::{generate, shells};
use serde_json::{Value, json};
use std::collections::BTreeMap;
use std::io;
use std::path::PathBuf;
use std::time::Duration;

#[derive(Parser, Debug)]
#[command(name = "rs-peekaboo")]
#[command(about = "Rust-native cross-platform computer-use CLI")]
pub struct Cli {
    #[arg(long, global = true, visible_alias = "json-output")]
    pub json: bool,
    // ponytail: not global=true to avoid shadowing subcommand --mode (ImageMode).
    // Env var RS_PEEKABOO_MODE works without global propagation.
    #[arg(long, env = "RS_PEEKABOO_MODE")]
    pub mode: Option<String>,
    /// Prefer background/AX actions that avoid focus steal when possible.
    #[arg(long, global = true, env = "RS_PEEKABOO_BACKGROUND")]
    pub background: bool,
    #[command(subcommand)]
    pub command: Commands,
}

#[derive(Subcommand, Debug)]
pub enum Commands {
    See(SeeArgs),
    Image(ImageArgs),
    Vision(VisionArgs),
    List(ListArgs),
    Click(ClickArgs),
    Type(TypeArgs),
    Press(PressArgs),
    Hotkey(HotkeyArgs),
    Paste(PasteArgs),
    Scroll(ScrollArgs),
    Swipe(DragArgs),
    Drag(DragArgs),
    Move(MoveArgs),
    SetValue(SetValueArgs),
    PerformAction(PerformActionArgs),
    Window(WindowArgs),
    App(AppArgs),
    Open(OpenArgs),
    Menu(MenuArgs),
    Clipboard(ClipboardArgs),
    Permissions(PermissionsArgs),
    Shell(ShellArgs),
    Run(RunArgs),
    Sleep(SleepArgs),
    Clean(CleanArgs),
    Doctor(DoctorArgs),
    Mcp(McpArgs),
    Tools(ToolsArgs),
    Completions(CompletionsArgs),
}

#[derive(Args, Debug)]
pub struct SeeArgs {
    #[arg(long)]
    pub app: Option<String>,
    #[arg(long, default_value = "screen")]
    pub mode: String,
    #[arg(long)]
    pub path: Option<PathBuf>,
    #[arg(long)]
    pub retina: bool,
    #[arg(long)]
    pub tree: bool,
    #[arg(long)]
    pub flat: bool,
    #[arg(long)]
    pub focused: bool,
    #[arg(long)]
    pub at: Option<String>,
}

#[derive(Args, Debug)]
pub struct ImageArgs {
    #[arg(long, default_value = "screen")]
    pub mode: String,
    #[arg(long)]
    pub path: Option<PathBuf>,
    #[arg(long)]
    pub retina: bool,
    /// Capture primary window of this app (window-scoped).
    #[arg(long)]
    pub app: Option<String>,
}

#[derive(Args, Debug)]
pub struct VisionArgs {
    #[arg(long)]
    pub import: Option<PathBuf>,
    #[arg(long)]
    pub screenshot: Option<PathBuf>,
    #[arg(long)]
    pub json: bool,
}

#[derive(Args, Debug)]
pub struct ListArgs {
    #[command(subcommand)]
    pub kind: ListKind,
}

#[derive(Subcommand, Debug)]
pub enum ListKind {
    Apps,
    Windows,
    Screens,
    Menubar,
    Permissions,
}

#[derive(Args, Debug)]
pub struct ClickArgs {
    pub target: Option<String>,
    #[arg(long)]
    pub on: Option<String>,
    #[arg(long)]
    pub query: Option<String>,
    #[arg(long)]
    pub coords: Option<String>,
    #[arg(long)]
    pub snapshot: Option<String>,
    /// Click by stable snapshot element index.
    #[arg(long)]
    pub index: Option<u32>,
    #[arg(long, default_value = "left")]
    pub button: String,
    #[arg(long, default_value_t = 1)]
    pub count: u32,
    /// Prefer AX/background action without focus steal.
    #[arg(long)]
    pub background: bool,
}

#[derive(Args, Debug)]
pub struct TypeArgs {
    pub text: Option<String>,
    #[arg(long)]
    pub clear: bool,
    #[arg(long = "return")]
    pub press_return: bool,
    #[arg(long)]
    pub delay: Option<u64>,
}

#[derive(Args, Debug)]
pub struct PressArgs {
    pub key: String,
    #[arg(long, default_value_t = 1)]
    pub count: u32,
    #[arg(long)]
    pub delay: Option<u64>,
}

#[derive(Args, Debug)]
pub struct HotkeyArgs {
    pub keys: String,
}

#[derive(Args, Debug)]
pub struct PasteArgs {
    pub text: String,
}

#[derive(Args, Debug)]
pub struct ScrollArgs {
    #[arg(long, default_value = "down")]
    pub direction: String,
    #[arg(long, default_value_t = 3)]
    pub amount: u32,
}

#[derive(Args, Debug)]
pub struct DragArgs {
    #[arg(long)]
    pub from: String,
    #[arg(long)]
    pub to: String,
    #[arg(long, default_value_t = 250)]
    pub duration: u64,
}

#[derive(Args, Debug)]
pub struct MoveArgs {
    pub target: Option<String>,
    #[arg(long)]
    pub coords: Option<String>,
    #[arg(long)]
    pub snapshot: Option<String>,
}

#[derive(Args, Debug)]
pub struct SetValueArgs {
    #[arg(long)]
    pub on: String,
    #[arg(long)]
    pub value: String,
    #[arg(long)]
    pub snapshot: Option<String>,
}

#[derive(Args, Debug)]
pub struct PerformActionArgs {
    #[arg(long)]
    pub on: String,
    #[arg(long)]
    pub action: String,
    #[arg(long)]
    pub snapshot: Option<String>,
}

#[derive(Args, Debug)]
pub struct WindowArgs {
    #[command(subcommand)]
    pub action: WindowAction,
}

#[derive(Subcommand, Debug)]
pub enum WindowAction {
    List,
    Focus(AppOnly),
    Close(AppOnly),
    Minimize(AppOnly),
    Move(WindowBounds),
    Resize(WindowBounds),
    SetBounds(WindowBounds),
}

#[derive(Args, Debug)]
pub struct AppOnly {
    #[arg(long)]
    pub app: String,
}

#[derive(Args, Debug)]
pub struct WindowBounds {
    #[arg(long)]
    pub app: String,
    #[arg(long, default_value_t = 0)]
    pub x: i64,
    #[arg(long, default_value_t = 0)]
    pub y: i64,
    #[arg(long, default_value_t = 0)]
    pub width: i64,
    #[arg(long, default_value_t = 0)]
    pub height: i64,
}

#[derive(Args, Debug)]
pub struct AppArgs {
    #[command(subcommand)]
    pub action: AppAction,
}

#[derive(Subcommand, Debug)]
pub enum AppAction {
    List,
    Launch(AppOnly),
    Quit(AppOnly),
    Hide(AppOnly),
    Unhide(AppOnly),
    Switch(AppOnly),
}

#[derive(Args, Debug)]
pub struct OpenArgs {
    pub target: String,
    #[arg(long)]
    pub app: Option<String>,
    #[arg(long)]
    pub no_focus: bool,
}

#[derive(Args, Debug)]
pub struct MenuArgs {
    #[command(subcommand)]
    pub action: MenuAction,
}

#[derive(Subcommand, Debug)]
pub enum MenuAction {
    List(MenuList),
    ListAll(MenuList),
    Click(MenuClick),
}

#[derive(Args, Debug)]
pub struct MenuList {
    #[arg(long)]
    pub app: String,
}

#[derive(Args, Debug)]
pub struct MenuClick {
    #[arg(long)]
    pub app: String,
    #[arg(long)]
    pub menu: String,
    #[arg(long)]
    pub item: String,
}

#[derive(Args, Debug)]
pub struct ClipboardArgs {
    #[command(subcommand)]
    pub action: ClipboardAction,
}

#[derive(Subcommand, Debug)]
pub enum ClipboardAction {
    Read,
    Write(ClipboardWrite),
}

#[derive(Args, Debug)]
pub struct ClipboardWrite {
    pub text: String,
}

#[derive(Args, Debug)]
pub struct PermissionsArgs {
    #[command(subcommand)]
    pub action: Option<PermissionAction>,
}

#[derive(Subcommand, Debug)]
pub enum PermissionAction {
    Status,
    Grant,
}

#[derive(Args, Debug)]
pub struct ShellArgs {
    pub command: String,
    #[arg(long)]
    pub cwd: Option<PathBuf>,
}

#[derive(Args, Debug)]
pub struct RunArgs {
    pub file: PathBuf,
}

#[derive(Args, Debug)]
pub struct SleepArgs {
    pub duration: f64,
}

#[derive(Args, Debug)]
pub struct CleanArgs {
    #[arg(long)]
    pub all_snapshots: bool,
    #[arg(long)]
    pub snapshot: Option<String>,
}

#[derive(Args, Debug)]
pub struct DoctorArgs;

#[derive(Args, Debug)]
pub struct McpArgs;

#[derive(Args, Debug)]
pub struct ToolsArgs;

#[derive(Args, Debug)]
pub struct CompletionsArgs {
    #[arg(value_enum)]
    pub shell: CompletionShell,
}

#[derive(Copy, Clone, Debug, ValueEnum)]
pub enum CompletionShell {
    Bash,
    Zsh,
    Fish,
}

pub fn execute(cli: Cli) -> Result<()> {
    let config = PeekabooConfig {
        mode: cli
            .mode
            .as_deref()
            .map(ComputerUseMode::parse)
            .transpose()?
            .unwrap_or_default(),
        background: cli.background,
    };
    let peekaboo = Peekaboo::with_config(config);
    let result = match cli.command {
        Commands::See(args) => {
            let snapshot = peekaboo.see(
                args.app.as_deref(),
                ImageMode::parse_or_err(&args.mode)?,
                args.path,
                args.retina,
            )?;
            if args.tree {
                CommandResult::ok(build_element_tree(
                    &snapshot
                        .elements
                        .into_iter()
                        .map(UiNode::from)
                        .collect::<Vec<_>>(),
                ))?
            } else if args.focused {
                let nodes: Vec<UiNode> = snapshot.elements.into_iter().map(UiNode::from).collect();
                let focused: Vec<&UiNode> =
                    nodes.iter().filter(|n| n.focused == Some(true)).collect();
                CommandResult::ok(json!({
                    "snapshot_id": snapshot.snapshot_id,
                    "elements": focused
                }))?
            } else if let Some(ref at) = args.at {
                let point = parse_point(at)?;
                let nodes: Vec<UiNode> = snapshot.elements.into_iter().map(UiNode::from).collect();
                let at_elem: Vec<&UiNode> = nodes
                    .iter()
                    .filter(|n| n.bounds.is_some_and(|b| b.contains(&point)))
                    .collect();
                CommandResult::ok(json!({
                    "snapshot_id": snapshot.snapshot_id,
                    "elements": at_elem
                }))?
            } else {
                CommandResult::ok(json!({
                    "snapshot_id": snapshot.snapshot_id,
                    "elements": snapshot.elements
                }))?
            }
        }
        Commands::Image(args) => {
            if let Some(app) = args.app.as_deref() {
                CommandResult::ok(peekaboo.image_app(app, args.path, args.retina)?)?
            } else {
                CommandResult::ok(peekaboo.image(
                    ImageMode::parse_or_err(&args.mode)?,
                    args.path,
                    args.retina,
                )?)?
            }
        }
        Commands::List(args) => CommandResult::ok(match args.kind {
            ListKind::Apps => peekaboo.list_apps()?,
            ListKind::Windows => peekaboo.list_windows()?,
            ListKind::Screens => peekaboo.list_screens()?,
            ListKind::Menubar => peekaboo.menu("list", "SystemUIServer", None, None)?,
            ListKind::Permissions => peekaboo.permissions(),
        })?,
        Commands::Click(args) => {
            let background = args.background || peekaboo.config.background;
            let (target, button, count) = click_request(args)?;
            CommandResult::ok(peekaboo.click_with_options(target, &button, count, background)?)?
        }
        Commands::Type(args) => CommandResult::ok(peekaboo.type_text(
            args.text.as_deref().unwrap_or_default(),
            args.clear,
            args.press_return,
            args.delay,
            None,
        )?)?,
        Commands::Press(args) => {
            CommandResult::ok(peekaboo.press(&args.key, args.count, args.delay)?)?
        }
        Commands::Hotkey(args) => {
            let keys = split_keys(&args.keys);
            CommandResult::ok(peekaboo.hotkey(&keys)?)?
        }
        Commands::Paste(args) => CommandResult::ok(peekaboo.paste(&args.text)?)?,
        Commands::Scroll(args) => CommandResult::ok(
            peekaboo.scroll(Direction::parse_or_err(&args.direction)?, args.amount)?,
        )?,
        Commands::Swipe(args) => CommandResult::ok(peekaboo.swipe(
            Target::Point(parse_point(&args.from)?),
            Target::Point(parse_point(&args.to)?),
            args.duration,
        )?)?,
        Commands::Drag(args) => CommandResult::ok(peekaboo.drag(
            Target::Point(parse_point(&args.from)?),
            Target::Point(parse_point(&args.to)?),
            args.duration,
        )?)?,
        Commands::Move(args) => CommandResult::ok(peekaboo.move_cursor(target_from_move(args)?)?)?,
        Commands::SetValue(args) => CommandResult::ok(peekaboo.set_value(
            Target::Query {
                query: args.on,
                snapshot: args.snapshot,
            },
            &args.value,
        )?)?,
        Commands::PerformAction(args) => CommandResult::ok(peekaboo.perform_action(
            Target::Query {
                query: args.on,
                snapshot: args.snapshot,
            },
            &args.action,
        )?)?,
        Commands::Window(args) => CommandResult::ok(window(&peekaboo, args.action)?)?,
        Commands::App(args) => CommandResult::ok(app(&peekaboo, args.action)?)?,
        Commands::Open(args) => {
            CommandResult::ok(peekaboo.open(&args.target, args.app.as_deref(), args.no_focus)?)?
        }
        Commands::Menu(args) => CommandResult::ok(menu(&peekaboo, args.action)?)?,
        Commands::Clipboard(args) => CommandResult::ok(match args.action {
            ClipboardAction::Read => json!({ "text": peekaboo.clipboard_read()? }),
            ClipboardAction::Write(write) => peekaboo.clipboard_write(&write.text)?,
        })?,
        Commands::Permissions(args) => CommandResult::ok(match args.action {
            Some(PermissionAction::Grant) => peekaboo.grant_permissions()?,
            _ => peekaboo.permissions(),
        })?,
        Commands::Shell(args) => {
            CommandResult::ok(peekaboo.shell(&args.command, args.cwd.as_deref())?)?
        }
        Commands::Run(args) => CommandResult::ok(peekaboo.run_file(&args.file)?)?,
        Commands::Sleep(args) => {
            let millis = (args.duration * 1000.0).max(0.0) as u64;
            std::thread::sleep(Duration::from_millis(millis));
            CommandResult::ok(json!({ "slept_ms": millis }))?
        }
        Commands::Clean(args) => CommandResult::ok(
            json!({ "removed": cache::clean_snapshots(args.all_snapshots, args.snapshot.as_deref())? }),
        )?,
        Commands::Doctor(_) => CommandResult::ok(peekaboo.doctor()?)?,
        Commands::Mcp(_) => return crate::mcp::serve(&peekaboo),
        Commands::Tools(_) => CommandResult::ok(tool_catalog())?,
        Commands::Completions(args) => return completions(args.shell),
        Commands::Vision(args) => {
            let mut map = serde_json::Map::new();
            if let Some(ref import_path) = args.import {
                let data = std::fs::read(import_path)?;
                let detections: VisionDetections = serde_json::from_slice(&data)?;
                let vision_nodes: Vec<UiNode> = detections
                    .elements
                    .iter()
                    .enumerate()
                    .map(|(i, el)| {
                        let role = el.role.clone().unwrap_or_else(|| "unknown".into());
                        UiNode {
                            id: format!("vision:{i}"),
                            backend: "vision".into(),
                            role,
                            subrole: None,
                            title: el.label.clone(),
                            label: el.label.clone(),
                            description: None,
                            value: None,
                            identifier: None,
                            app: String::new(),
                            pid: None,
                            window: None,
                            bounds: el.bounds,
                            enabled: None,
                            focused: None,
                            selected: None,
                            depth: None,
                            index_in_parent: Some(i as i32),
                            index: None,
                            parent_id: None,
                            children: None,
                            children_count: None,
                            actions: Vec::new(),
                            attributes: std::collections::HashMap::new(),
                            confidence: el.confidence,
                            source: vec!["vision".into()],
                            state: serde_json::json!({}),
                        }
                    })
                    .collect();
                map.insert("elements".into(), json!(vision_nodes));
                map.insert("source".into(), json!("import"));
            }
            if let Some(ref screenshot_path) = args.screenshot {
                peekaboo.image(ImageMode::Screen, Some(screenshot_path.clone()), false)?;
                map.insert(
                    "screenshot".into(),
                    json!(screenshot_path.to_string_lossy().to_string()),
                );
            }
            if args.json && map.is_empty() {
                map.insert("ok".into(), json!(true));
            }
            CommandResult::ok(Value::Object(map))?
        }
    };
    emit(result, cli.json)
}

fn target_from_click(args: ClickArgs) -> Result<Target> {
    if let Some(coords) = args.coords {
        return Ok(Target::Point(parse_point(&coords)?));
    }
    if let Some(index) = args.index {
        return Ok(Target::Query {
            query: format!("index={index}"),
            snapshot: args.snapshot,
        });
    }
    let query = args
        .query
        .or(args.on)
        .or(args.target)
        .ok_or(crate::PeekabooError::MissingArgument("target"))?;
    Ok(Target::Query {
        query,
        snapshot: args.snapshot,
    })
}

fn click_request(args: ClickArgs) -> Result<(Target, String, u32)> {
    let button = args.button.clone();
    let count = args.count;
    Ok((target_from_click(args)?, button, count))
}

fn target_from_move(args: MoveArgs) -> Result<Target> {
    if let Some(coords) = args.coords {
        return Ok(Target::Point(parse_point(&coords)?));
    }
    let query = args
        .target
        .ok_or(crate::PeekabooError::MissingArgument("target"))?;
    Ok(Target::Query {
        query,
        snapshot: args.snapshot,
    })
}

fn window(peekaboo: &Peekaboo, action: WindowAction) -> Result<Value> {
    match action {
        WindowAction::List => peekaboo.window("list", None, None, None),
        WindowAction::Focus(args) => peekaboo.window("focus", Some(&args.app), None, None),
        WindowAction::Close(args) => peekaboo.window("close", Some(&args.app), None, None),
        WindowAction::Minimize(args) => peekaboo.window("minimize", Some(&args.app), None, None),
        WindowAction::Move(args) => {
            let app = args.app.clone();
            peekaboo.window("move", Some(&app), None, Some(bounds(args)))
        }
        WindowAction::Resize(args) => {
            let app = args.app.clone();
            peekaboo.window("resize", Some(&app), None, Some(bounds(args)))
        }
        WindowAction::SetBounds(args) => {
            let app = args.app.clone();
            peekaboo.window("set-bounds", Some(&app), None, Some(bounds(args)))
        }
    }
}

fn bounds(args: WindowBounds) -> Bounds {
    Bounds {
        x: args.x,
        y: args.y,
        width: args.width,
        height: args.height,
    }
}

fn app(peekaboo: &Peekaboo, action: AppAction) -> Result<Value> {
    match action {
        AppAction::List => peekaboo.app("list", None),
        AppAction::Launch(args) => peekaboo.app("launch", Some(&args.app)),
        AppAction::Quit(args) => peekaboo.app("quit", Some(&args.app)),
        AppAction::Hide(args) => peekaboo.app("hide", Some(&args.app)),
        AppAction::Unhide(args) => peekaboo.app("unhide", Some(&args.app)),
        AppAction::Switch(args) => peekaboo.app("switch", Some(&args.app)),
    }
}

fn menu(peekaboo: &Peekaboo, action: MenuAction) -> Result<Value> {
    match action {
        MenuAction::List(args) => peekaboo.menu("list", &args.app, None, None),
        MenuAction::ListAll(args) => peekaboo.menu("list-all", &args.app, None, None),
        MenuAction::Click(args) => {
            peekaboo.menu("click", &args.app, Some(&args.menu), Some(&args.item))
        }
    }
}

fn tool_catalog() -> Value {
    json!([
        "see",
        "image",
        "vision",
        "list",
        "click",
        "type",
        "press",
        "hotkey",
        "paste",
        "scroll",
        "swipe",
        "drag",
        "move",
        "set-value",
        "perform-action",
        "window",
        "app",
        "open",
        "menu",
        "clipboard",
        "permissions",
        "shell",
        "run",
        "sleep",
        "clean",
        "doctor",
        "mcp",
        "tools",
        "completions"
    ])
}

fn completions(shell: CompletionShell) -> Result<()> {
    let mut command = Cli::command();
    let mut stdout = io::stdout();
    match shell {
        CompletionShell::Bash => generate(shells::Bash, &mut command, "rs-peekaboo", &mut stdout),
        CompletionShell::Zsh => generate(shells::Zsh, &mut command, "rs-peekaboo", &mut stdout),
        CompletionShell::Fish => generate(shells::Fish, &mut command, "rs-peekaboo", &mut stdout),
    };
    Ok(())
}

fn build_element_tree(nodes: &[UiNode]) -> Value {
    let mut by_app: BTreeMap<String, Vec<&UiNode>> = BTreeMap::new();
    for node in nodes {
        by_app.entry(node.app.clone()).or_default().push(node);
    }
    let tree: Vec<Value> = by_app
        .into_iter()
        .map(|(app, nodes)| {
            let mut by_window: BTreeMap<String, Vec<&UiNode>> = BTreeMap::new();
            for node in &nodes {
                let key = node.window.clone().unwrap_or_default();
                by_window.entry(key).or_default().push(node);
            }
            let windows: Vec<Value> = by_window
                .into_iter()
                .map(|(window, nodes)| {
                    json!({
                        "window": window,
                        "elements": nodes
                    })
                })
                .collect();
            json!({
                "app": app,
                "windows": windows
            })
        })
        .collect();
    json!(tree)
}

fn emit(result: CommandResult, json_output: bool) -> Result<()> {
    if json_output {
        println!("{}", serde_json::to_string_pretty(&result)?);
    } else if let Some(data) = &result.data {
        print_human(data);
    }
    Ok(())
}

fn print_human(value: &Value) {
    match value {
        Value::String(text) => println!("{text}"),
        _ => println!(
            "{}",
            serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string())
        ),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn click_command_should_preserve_button_and_count() {
        let args = ClickArgs {
            target: None,
            on: None,
            query: None,
            coords: Some("10,20".to_string()),
            snapshot: None,
            index: None,
            button: "right".to_string(),
            count: 3,
            background: false,
        };

        let (_, button, count) = click_request(args).expect("click args should parse");

        assert_eq!(button, "right");
        assert_eq!(count, 3);
    }
}