owlen/crates/owlen-cli/src/mcp.rs

use std::collections::{HashMap, HashSet};

use anyhow::{Result, anyhow};
use clap::{Args, Subcommand, ValueEnum};
use owlen_core::config::{self as core_config, Config, McpConfigScope, McpServerConfig};
use owlen_tui::config as tui_config;

#[derive(Debug, Subcommand)]
pub enum McpCommand {
    /// Add or update an MCP server in the selected scope
    Add(AddArgs),
    /// List MCP servers across scopes
    List(ListArgs),
    /// Remove an MCP server from a scope
    Remove(RemoveArgs),
}

pub fn run_mcp_command(command: McpCommand) -> Result<()> {
    match command {
        McpCommand::Add(args) => handle_add(args),
        McpCommand::List(args) => handle_list(args),
        McpCommand::Remove(args) => handle_remove(args),
    }
}

#[derive(Debug, Clone, Copy, ValueEnum, Default)]
pub enum ScopeArg {
    User,
    #[default]
    Project,
    Local,
}

impl From<ScopeArg> for McpConfigScope {
    fn from(value: ScopeArg) -> Self {
        match value {
            ScopeArg::User => McpConfigScope::User,
            ScopeArg::Project => McpConfigScope::Project,
            ScopeArg::Local => McpConfigScope::Local,
        }
    }
}

#[derive(Debug, Args)]
pub struct AddArgs {
    /// Logical name used to reference the server
    pub name: String,
    /// Command or endpoint invoked for the server
    pub command: String,
    /// Transport mechanism (stdio, http, websocket)
    #[arg(long, default_value = "stdio")]
    pub transport: String,
    /// Configuration scope to write the server into
    #[arg(long, value_enum, default_value_t = ScopeArg::Project)]
    pub scope: ScopeArg,
    /// Environment variables (KEY=VALUE) passed to the server process
    #[arg(long = "env")]
    pub env: Vec<String>,
    /// Additional arguments appended when launching the server
    #[arg(trailing_var_arg = true, value_name = "ARG")]
    pub args: Vec<String>,
}

#[derive(Debug, Args, Default)]
pub struct ListArgs {
    /// Restrict output to a specific configuration scope
    #[arg(long, value_enum)]
    pub scope: Option<ScopeArg>,
    /// Display only the effective servers (after precedence resolution)
    #[arg(long)]
    pub effective_only: bool,
}

#[derive(Debug, Args)]
pub struct RemoveArgs {
    /// Name of the server to remove
    pub name: String,
    /// Optional explicit scope to remove from
    #[arg(long, value_enum)]
    pub scope: Option<ScopeArg>,
}

fn handle_add(args: AddArgs) -> Result<()> {
    let mut config = load_config()?;
    let scope: McpConfigScope = args.scope.into();
    let mut env_map = HashMap::new();
    for pair in &args.env {
        let (key, value) = pair
            .split_once('=')
            .ok_or_else(|| anyhow!("Environment pairs must use KEY=VALUE syntax: '{}'", pair))?;
        if key.trim().is_empty() {
            return Err(anyhow!("Environment variable name cannot be empty"));
        }
        env_map.insert(key.trim().to_string(), value.to_string());
    }

    let server = McpServerConfig {
        name: args.name.clone(),
        command: args.command.clone(),
        args: args.args.clone(),
        transport: args.transport.to_lowercase(),
        env: env_map,
        oauth: None,
    };

    config.add_mcp_server(scope, server.clone(), None)?;
    if matches!(scope, McpConfigScope::User) {
        tui_config::save_config(&config)?;
    }

    if let Some(path) = core_config::mcp_scope_path(scope, None) {
        println!(
            "Registered MCP server '{}' in {} scope ({})",
            server.name,
            scope,
            path.display()
        );
    } else {
        println!(
            "Registered MCP server '{}' in {} scope.",
            server.name, scope
        );
    }

    Ok(())
}

fn handle_list(args: ListArgs) -> Result<()> {
    let mut config = load_config()?;
    config.refresh_mcp_servers(None)?;

    let scoped = config.scoped_mcp_servers();
    if scoped.is_empty() {
        println!("No MCP servers configured.");
        return Ok(());
    }

    let filter_scope = args.scope.map(|scope| scope.into());
    let effective = config.effective_mcp_servers();
    let mut active = HashSet::new();
    for server in effective {
        active.insert((
            server.name.clone(),
            server.command.clone(),
            server.transport.to_lowercase(),
        ));
    }

    println!(
        "{:<2} {:<8} {:<20} {:<10} Command",
        "", "Scope", "Name", "Transport"
    );
    for entry in scoped {
        if filter_scope
            .as_ref()
            .is_some_and(|target_scope| entry.scope != *target_scope)
        {
            continue;
        }

        let payload = format_command_line(&entry.config.command, &entry.config.args);
        let key = (
            entry.config.name.clone(),
            entry.config.command.clone(),
            entry.config.transport.to_lowercase(),
        );
        let marker = if active.contains(&key) { "*" } else { " " };

        if args.effective_only && marker != "*" {
            continue;
        }

        println!(
            "{} {:<8} {:<20} {:<10} {}",
            marker, entry.scope, entry.config.name, entry.config.transport, payload
        );
    }

    let scoped_resources = config.scoped_mcp_resources();
    if !scoped_resources.is_empty() {
        println!();
        println!("{:<2} {:<8} {:<30} Title", "", "Scope", "Resource");
        let effective_keys: HashSet<(String, String)> = config
            .effective_mcp_resources()
            .iter()
            .map(|res| (res.server.clone(), res.uri.clone()))
            .collect();

        for entry in scoped_resources {
            if filter_scope
                .as_ref()
                .is_some_and(|target_scope| entry.scope != *target_scope)
            {
                continue;
            }

            let key = (entry.config.server.clone(), entry.config.uri.clone());
            let marker = if effective_keys.contains(&key) {
                "*"
            } else {
                " "
            };
            if args.effective_only && marker != "*" {
                continue;
            }

            let reference = format!("@{}:{}", entry.config.server, entry.config.uri);
            let title = entry.config.title.as_deref().unwrap_or("—");

            println!("{} {:<8} {:<30} {}", marker, entry.scope, reference, title);
        }
    }

    Ok(())
}

fn handle_remove(args: RemoveArgs) -> Result<()> {
    let mut config = load_config()?;
    let scope_hint = args.scope.map(|scope| scope.into());
    let result = config.remove_mcp_server(scope_hint, &args.name, None)?;

    match result {
        Some(scope) => {
            if matches!(scope, McpConfigScope::User) {
                tui_config::save_config(&config)?;
            }

            if let Some(path) = core_config::mcp_scope_path(scope, None) {
                println!(
                    "Removed MCP server '{}' from {} scope ({})",
                    args.name,
                    scope,
                    path.display()
                );
            } else {
                println!("Removed MCP server '{}' from {} scope.", args.name, scope);
            }
        }
        None => {
            println!("No MCP server named '{}' was found.", args.name);
        }
    }

    Ok(())
}

fn load_config() -> Result<Config> {
    let mut config = tui_config::try_load_config().unwrap_or_default();
    config.refresh_mcp_servers(None)?;
    Ok(config)
}

fn format_command_line(command: &str, args: &[String]) -> String {
    if args.is_empty() {
        command.to_string()
    } else {
        format!("{} {}", command, args.join(" "))
    }
}