PicoBot/src/main.rs

use clap::{CommandFactory, Parser, Subcommand};

#[derive(Subcommand)]
enum ServiceCommand {
    /// Install and enable the systemd user service
    Install,
    /// Stop the systemd user service
    Stop,
    /// Show the systemd user service status
    Status,
    /// Start the systemd user service
    Start,
    /// Restart the systemd user service
    Restart,
    /// Stop, disable, and remove the systemd user service
    Uninstall,
}

#[derive(Parser)]
#[command(name = "picobot")]
#[command(about = "A CLI chatbot", long_about = None)]
#[command(version)]
enum Command {
    /// Connect to gateway
    Chat {
        /// Gateway WebSocket URL (e.g., ws://127.0.0.1:19876/ws)
        #[arg(long)]
        gateway_url: Option<String>,
        /// One-time pairing code; saves the issued client token locally
        #[arg(long)]
        pair_code: Option<String>,
    },
    /// Send one prompt through the gateway, print the final response, and exit
    Run {
        /// Prompt text; when omitted, read it from stdin
        prompt: Vec<String>,
        /// Gateway WebSocket or HTTP URL
        #[arg(long)]
        gateway_url: Option<String>,
        /// Maximum time to wait for the turn, in seconds
        #[arg(long, default_value_t = 300)]
        timeout: u64,
        /// Print the terminal turn as one JSON object
        #[arg(long)]
        json: bool,
        /// Print phase and tool progress to stderr
        #[arg(long)]
        verbose: bool,
    },
    /// Start gateway server
    Gateway {
        /// Host to bind to
        #[arg(long)]
        host: Option<String>,
        /// Port to listen on
        #[arg(long)]
        port: Option<u16>,
    },
    /// Reload a running gateway's configuration
    Reload {
        /// Gateway WebSocket or HTTP URL
        #[arg(long)]
        gateway_url: Option<String>,
    },
    /// Check PicoBot runtime and configured external dependencies
    Health {
        /// Print the structured report as JSON
        #[arg(long)]
        json: bool,
    },
    /// Generate a one-time browser pairing code from the local gateway
    Pair {
        /// Gateway WebSocket or HTTP URL
        #[arg(long)]
        gateway_url: Option<String>,
        /// Revoke every paired browser and CLI token before issuing the code
        #[arg(long)]
        revoke_all: bool,
    },
    /// Manage the PicoBot systemd user service
    Service {
        #[command(subcommand)]
        command: ServiceCommand,
    },
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut cmd = Command::command();

    // If no arguments, print help
    if std::env::args().len() <= 1 {
        cmd.print_help()?;
        println!();
        return Ok(());
    }

    match Command::parse() {
        Command::Chat {
            gateway_url,
            pair_code,
        } => {
            let config = picobot::config::Config::load_default().ok();
            let url = gateway_url
                .or_else(|| config.as_ref().map(|c| c.client.gateway_url.clone()))
                .unwrap_or_else(|| "ws://127.0.0.1:19876/ws".to_string());
            picobot::client::run(&url, pair_code.as_deref()).await?;
        }
        Command::Run {
            prompt,
            gateway_url,
            timeout,
            json,
            verbose,
        } => {
            if timeout == 0 {
                return Err("--timeout must be greater than zero".into());
            }
            let config = picobot::config::Config::load_default().ok();
            let url = gateway_url
                .or_else(|| config.as_ref().map(|c| c.client.gateway_url.clone()))
                .unwrap_or_else(|| "ws://127.0.0.1:19876/ws".to_string());
            let prompt = picobot::client::read_run_prompt(prompt)?;
            picobot::client::run_once(
                &url,
                prompt,
                picobot::client::RunOptions {
                    timeout: std::time::Duration::from_secs(timeout),
                    json,
                    verbose,
                },
            )
            .await?;
        }
        Command::Gateway { host, port } => {
            picobot::gateway::run(host, port).await?;
        }
        Command::Reload { gateway_url } => {
            let config = picobot::config::Config::load_default().ok();
            let url = gateway_url
                .or_else(|| config.as_ref().map(|c| c.client.gateway_url.clone()))
                .unwrap_or_else(|| "ws://127.0.0.1:19876/ws".to_string());
            println!("{}", picobot::client::reload_gateway(&url).await?);
        }
        Command::Health { json } => {
            let report = match picobot::config::Config::load_default() {
                Ok(config) => picobot::health::HealthService::new(config).check().await,
                Err(error) => picobot::health::HealthReport::configuration_error(error.to_string()),
            };
            if json {
                println!("{}", serde_json::to_string_pretty(&report)?);
            } else {
                println!("{}", report.render_text());
            }
            if !report.is_usable() {
                std::process::exit(1);
            }
        }
        Command::Pair {
            gateway_url,
            revoke_all,
        } => {
            let config = picobot::config::Config::load_default().ok();
            let url = gateway_url
                .or_else(|| config.as_ref().map(|c| c.client.gateway_url.clone()))
                .unwrap_or_else(|| "ws://127.0.0.1:19876/ws".to_string());
            let mut endpoint = gateway_api_url(&url, "/api/auth/code")?;
            if revoke_all {
                endpoint.query_pairs_mut().append_pair("revoke_all", "true");
            }
            let admin_token_path = picobot::config::get_user_config_dir().join("web_admin_token");
            let admin_token = std::fs::read_to_string(&admin_token_path).map_err(|error| {
                format!(
                    "cannot read local gateway admin token {}: {error}",
                    admin_token_path.display()
                )
            })?;
            let response = reqwest::Client::new()
                .post(endpoint)
                .header(
                    picobot::gateway::auth::ADMIN_TOKEN_HEADER,
                    admin_token.trim(),
                )
                .send()
                .await?;
            let status = response.status();
            let body: serde_json::Value = response.json().await?;
            if !status.is_success() {
                return Err(body
                    .get("error")
                    .and_then(serde_json::Value::as_str)
                    .unwrap_or("failed to generate pairing code")
                    .to_string()
                    .into());
            }
            let code = body
                .get("pairing_code")
                .and_then(serde_json::Value::as_str)
                .ok_or("gateway did not return a pairing code")?;
            println!("Pairing code: {code}");
            println!("Expires in 5 minutes and can be used once.");
            if revoke_all {
                println!("All existing paired devices were revoked.");
            }
        }
        Command::Service { command } => {
            let command = match command {
                ServiceCommand::Install => picobot::service::ServiceCommand::Install,
                ServiceCommand::Stop => picobot::service::ServiceCommand::Stop,
                ServiceCommand::Status => picobot::service::ServiceCommand::Status,
                ServiceCommand::Start => picobot::service::ServiceCommand::Start,
                ServiceCommand::Restart => picobot::service::ServiceCommand::Restart,
                ServiceCommand::Uninstall => picobot::service::ServiceCommand::Uninstall,
            };
            picobot::service::execute(command).await?;
        }
    }
    Ok(())
}

fn gateway_api_url(
    gateway_url: &str,
    path: &str,
) -> Result<reqwest::Url, Box<dyn std::error::Error>> {
    let mut url = reqwest::Url::parse(gateway_url)?;
    let scheme = match url.scheme() {
        "ws" => "http",
        "wss" => "https",
        "http" => "http",
        "https" => "https",
        other => return Err(format!("unsupported gateway URL scheme: {other}").into()),
    };
    url.set_scheme(scheme)
        .map_err(|_| "failed to set gateway URL scheme")?;
    url.set_path(path);
    url.set_query(None);
    url.set_fragment(None);
    Ok(url)
}