opencode-provider-manager 0.1.7-beta.3

//! TUI (Terminal User Interface) for OpenCode Provider Manager.

use anyhow::{Context, Result};
use app::import::{ImportMergeMode, import_source};
use app::state::AppState;
use clap::{Parser, Subcommand};
use config_core::{ConfigLayer, OpenCodeConfig};
use opencode_provider_manager::{app, config_core};
use serde::Serialize;
use std::path::PathBuf;
use std::process;

mod event;
mod tui_app;
mod ui;

/// Command-line arguments.
#[derive(Parser, Debug)]
#[command(name = "opm", about = "OpenCode Provider Manager", version)]
struct Args {
    /// Subcommand to run (defaults to TUI if not specified)
    #[command(subcommand)]
    command: Option<Commands>,

    /// Start with a specific config layer view (for TUI mode, defaults to project).
    #[arg(long, value_name = "LAYER", global = true)]
    layer: Option<String>,

    /// Path to a custom opencode.json config file (for TUI mode).
    #[arg(long, value_name = "PATH", global = true)]
    config: Option<String>,

    /// Start in split view mode (for TUI mode).
    #[arg(long, global = true)]
    split: bool,
}

/// Available subcommands.
#[derive(Subcommand, Debug)]
enum Commands {
    /// Launch the TUI (default behavior).
    Tui {
        /// Start with a specific config layer view (defaults to project).
        #[arg(long, value_name = "LAYER")]
        layer: Option<String>,

        /// Path to a custom opencode.json config file.
        #[arg(long, value_name = "PATH")]
        config: Option<String>,

        /// Start in split view mode.
        #[arg(long)]
        split: bool,
    },

    /// List configured providers as JSON.
    ListProviders {
        /// Which config layer to read from (defaults to merged).
        #[arg(long, value_name = "LAYER", default_value = "merged")]
        layer: String,
    },

    /// Show config as JSON.
    ShowConfig {
        /// Which config layer to show (defaults to merged).
        #[arg(long, value_name = "LAYER", default_value = "merged")]
        layer: String,
    },

    /// Validate config files.
    Validate,

    /// Import JSON/JSONC/TOML/YAML config/provider/model snippets from a file, directory, URL, or inline text.
    Import {
        /// File path, directory path, GitHub URL, raw URL, or inline snippet to import.
        #[arg(long, value_name = "SOURCE")]
        input: String,

        /// Target config layer to modify.
        #[arg(long, value_name = "LAYER", default_value = "project")]
        layer: String,

        /// Merge imported config into the target layer, or replace that layer.
        #[arg(long, value_name = "MODE", default_value = "merge")]
        mode: String,

        /// Provider ID hint for provider/model fragments that do not contain an ID.
        #[arg(long, value_name = "ID")]
        provider_id: Option<String>,

        /// Preview import summary without saving.
        #[arg(long)]
        dry_run: bool,
    },

    /// Manage oh-my-openagent agent configurations.
    AgentConfig {
        /// Subcommand for agent config management.
        #[command(subcommand)]
        command: AgentConfigCommands,
    },
}

/// Subcommands for agent-config management.
#[derive(Subcommand, Debug)]
enum AgentConfigCommands {
    /// Show agent config as JSON.
    Show {
        /// Which config layer to show (defaults to merged).
        #[arg(long, value_name = "LAYER", default_value = "merged")]
        layer: String,

        /// Path to a custom oh-my-opencode.json config file.
        #[arg(long, value_name = "PATH")]
        config: Option<String>,
    },

    /// Validate agent config files.
    Validate {
        /// Path to a custom oh-my-opencode.json config file.
        #[arg(long, value_name = "PATH")]
        config: Option<String>,
    },

    /// List all configured agents.
    ListAgents {
        /// Which config layer to read from (defaults to merged).
        #[arg(long, value_name = "LAYER", default_value = "merged")]
        layer: String,

        /// Path to a custom oh-my-opencode.json config file.
        #[arg(long, value_name = "PATH")]
        config: Option<String>,
    },

    /// List available models from opencode.
    ListAvailableModels,

    /// Set model for an agent.
    SetModel {
        /// Agent ID (e.g., build, plan, sisyphus).
        #[arg(value_name = "AGENT")]
        agent: String,

        /// Model ID in provider/model format.
        #[arg(long, value_name = "MODEL")]
        model: String,

        /// Fallback model IDs (comma-separated).
        #[arg(long, value_name = "MODELS")]
        fallback: Option<String>,

        /// Target config layer (defaults to project).
        #[arg(long, value_name = "LAYER", default_value = "project")]
        layer: String,

        /// Path to a custom oh-my-opencode.json config file.
        #[arg(long, value_name = "PATH")]
        config: Option<String>,
    },
}

/// Provider info for JSON output.
#[derive(Serialize)]
struct ProviderInfo {
    id: String,
    name: Option<String>,
}

fn main() {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    let args = Args::parse();

    let global_config = args.config.clone();

    // Determine which command to run
    let result = match args.command {
        None => {
            // No subcommand - run TUI with top-level args
            run_tui_blocking(args.layer, args.config, args.split)
        }
        Some(Commands::Tui {
            layer,
            config,
            split,
        }) => {
            // Explicit TUI subcommand
            run_tui_blocking(layer, config.or(global_config), split)
        }
        Some(Commands::ListProviders { layer }) => {
            // CLI: list providers
            run_list_providers(&layer, args.config.as_deref())
        }
        Some(Commands::ShowConfig { layer }) => {
            // CLI: show config
            run_show_config(&layer, args.config.as_deref())
        }
        Some(Commands::Validate) => {
            // CLI: validate configs
            run_validate(args.config.as_deref())
        }
        Some(Commands::Import {
            input,
            layer,
            mode,
            provider_id,
            dry_run,
        }) => run_import(
            &input,
            &layer,
            &mode,
            provider_id.as_deref(),
            args.config.as_deref(),
            dry_run,
        ),
        Some(Commands::AgentConfig { command }) => match command {
            AgentConfigCommands::Show { layer, config } => {
                run_agent_show(&layer, config.as_deref())
            }
            AgentConfigCommands::Validate { config } => run_agent_validate(config.as_deref()),
            AgentConfigCommands::ListAgents { layer, config } => {
                run_agent_list(&layer, config.as_deref())
            }
            AgentConfigCommands::ListAvailableModels => run_agent_list_available_models(),
            AgentConfigCommands::SetModel {
                agent,
                model,
                fallback,
                layer,
                config,
            } => run_agent_set_model(
                &agent,
                &model,
                fallback.as_deref(),
                &layer,
                config.as_deref(),
            ),
        },
    };

    // Handle errors — use Alternate ({:#}) to print the full error chain,
    // not just the outermost context message.
    if let Err(e) = result {
        eprintln!("Error: {e:#}");
        process::exit(1);
    }
}

fn run_tui_blocking(layer: Option<String>, config: Option<String>, split: bool) -> Result<()> {
    tokio::runtime::Runtime::new()
        .context("Failed to initialize async runtime")?
        .block_on(run_tui(layer, config, split))
}

/// Run the TUI application.
async fn run_tui(layer: Option<String>, config: Option<String>, split: bool) -> Result<()> {
    // Initialize app state
    let mut state = AppState::new().context("Failed to initialize app state")?;

    apply_custom_config_path(&mut state, config.as_deref())?;

    // Apply explicit layer selection
    if let Some(ref layer_str) = layer {
        match layer_str.to_lowercase().as_str() {
            "global" => state.edit_layer = config_core::ConfigLayer::Global,
            "project" => state.edit_layer = config_core::ConfigLayer::Project,
            "custom" => state.edit_layer = config_core::ConfigLayer::Custom,
            other => {
                return Err(anyhow::anyhow!(
                    "Invalid --layer '{}'. Must be one of: global, project, custom",
                    other
                ));
            }
        }
    }

    // Load configs
    state.load_configs().context("Failed to load configs")?;

    // Run TUI
    let terminal = ratatui::init();
    let result = tui_app::run(terminal, state, split).await;
    ratatui::restore();

    result
}

fn load_state(config: Option<&str>) -> Result<AppState> {
    let mut state = AppState::new().context("Failed to initialize app state")?;
    apply_custom_config_path(&mut state, config)?;
    state.load_configs().context("Failed to load configs")?;
    Ok(state)
}

fn parse_config_layer(layer: &str) -> Result<ConfigLayer> {
    match layer.to_lowercase().as_str() {
        "global" => Ok(ConfigLayer::Global),
        "project" => Ok(ConfigLayer::Project),
        "custom" => Ok(ConfigLayer::Custom),
        other => Err(anyhow::anyhow!(
            "Invalid layer '{}'. Must be one of: global, project, custom",
            other
        )),
    }
}

fn parse_import_mode(mode: &str) -> Result<ImportMergeMode> {
    match mode.to_lowercase().as_str() {
        "merge" => Ok(ImportMergeMode::Merge),
        "replace" => Ok(ImportMergeMode::Replace),
        other => Err(anyhow::anyhow!(
            "Invalid import mode '{}'. Must be one of: merge, replace",
            other
        )),
    }
}

fn apply_custom_config_path(state: &mut AppState, config: Option<&str>) -> Result<()> {
    let Some(path_str) = config else {
        return Ok(());
    };

    let config_path = PathBuf::from(path_str);
    let ext = config_path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("");
    if !matches!(ext, "json" | "jsonc") {
        return Err(anyhow::anyhow!(
            "Invalid --config path: file must have .json or .jsonc extension, got '{}'",
            path_str
        ));
    }

    let canonical = if config_path.exists() {
        config_path
            .canonicalize()
            .context("Failed to resolve config path")?
    } else if let Some(parent) = config_path.parent() {
        if parent.as_os_str().is_empty() {
            config_path.clone()
        } else if parent.exists() {
            let file_name = config_path
                .file_name()
                .ok_or_else(|| anyhow::anyhow!("Invalid --config path: missing file name"))?;
            parent
                .canonicalize()
                .context("Failed to resolve config directory")?
                .join(file_name)
        } else {
            config_path.clone()
        }
    } else {
        config_path.clone()
    };

    state.paths.custom = Some(canonical);
    Ok(())
}

/// List providers as JSON.
fn run_list_providers(layer_str: &str, config: Option<&str>) -> Result<()> {
    let state = load_state(config)?;

    // Get the appropriate config based on layer
    let config = get_config_for_layer(&state, layer_str)
        .with_context(|| format!("Invalid layer: {}", layer_str))?;

    // Build provider info list
    let providers: Vec<ProviderInfo> = config
        .provider
        .as_ref()
        .map(|providers| {
            providers
                .iter()
                .map(|(id, provider)| ProviderInfo {
                    id: id.clone(),
                    name: provider.name.clone(),
                })
                .collect()
        })
        .unwrap_or_default();

    // Output as JSON
    let json = serde_json::to_string_pretty(&providers)
        .context("Failed to serialize providers to JSON")?;
    println!("{}", json);

    Ok(())
}

/// Show config as JSON (with sensitive values redacted).
fn run_show_config(layer_str: &str, config: Option<&str>) -> Result<()> {
    let state = load_state(config)?;

    // Get the appropriate config based on layer
    let config = get_config_for_layer(&state, layer_str)
        .with_context(|| format!("Invalid layer: {}", layer_str))?;

    // Deep-clone and redact sensitive fields before serialization
    let mut redacted = config.clone();
    redact_sensitive_values(&mut redacted);

    // Output as pretty JSON
    let json =
        serde_json::to_string_pretty(&redacted).context("Failed to serialize config to JSON")?;
    println!("{}", json);

    Ok(())
}

/// Key names that should be redacted from JSON output.
const SENSITIVE_KEYS: &[&str] = &[
    "apiKey",
    "apikey",
    "key",
    "secret",
    "token",
    "password",
    "credential",
    "privateKey",
    "private_key",
    "accessToken",
    "access_token",
    "refreshToken",
    "refresh_token",
];

/// Recursively redact sensitive string values in a config.
fn redact_sensitive_values(config: &mut config_core::OpenCodeConfig) {
    if let Some(ref mut providers) = config.provider {
        for provider in providers.values_mut() {
            if let Some(ref mut options) = provider.options {
                for (key, value) in options.iter_mut() {
                    if SENSITIVE_KEYS.contains(&key.as_str()) && value.is_string() {
                        *value = serde_json::Value::String("***".to_string());
                    }
                }
            }
        }
    }
}

/// Validate config files.
fn run_validate(config: Option<&str>) -> Result<()> {
    let state = load_state(config)?;

    let mut has_errors = false;

    // Validate global config if present
    if let Some(ref global) = state.global_config {
        if let Err(e) = config_core::validate_config(global) {
            eprintln!("Global config error: {}", e);
            has_errors = true;
        } else {
            println!("Global config: OK");
        }
    } else {
        println!("Global config: not found");
    }

    // Validate custom config if present
    if let Some(ref custom) = state.custom_config {
        if let Err(e) = config_core::validate_config(custom) {
            eprintln!("Custom config error: {}", e);
            has_errors = true;
        } else {
            println!("Custom config: OK");
        }
    } else if state.paths.custom.is_some() {
        println!("Custom config: not found");
    }

    // Validate project config if present
    if let Some(ref project) = state.project_config {
        if let Err(e) = config_core::validate_config(project) {
            eprintln!("Project config error: {}", e);
            has_errors = true;
        } else {
            println!("Project config: OK");
        }
    } else {
        println!("Project config: not found");
    }

    // Validate merged config
    if let Err(e) = config_core::validate_config(&state.merged_config) {
        eprintln!("Merged config error: {}", e);
        has_errors = true;
    } else {
        println!("Merged config: OK");
    }

    if has_errors {
        process::exit(1);
    }

    Ok(())
}

fn run_import(
    input: &str,
    layer_str: &str,
    mode_str: &str,
    provider_id: Option<&str>,
    custom_config: Option<&str>,
    dry_run: bool,
) -> Result<()> {
    let mut state = load_state(custom_config)?;
    let layer = parse_config_layer(layer_str)?;
    let mode = parse_import_mode(mode_str)?;
    let summary = import_source(&mut state, input, provider_id, layer, mode)?;

    println!(
        "Imported {} provider(s), {} model(s): {}",
        summary.provider_count,
        summary.model_count,
        if summary.provider_ids.is_empty() {
            "(none)".to_string()
        } else {
            summary.provider_ids.join(", ")
        }
    );

    if dry_run {
        println!("Dry run: not saved");
        return Ok(());
    }

    state.save(layer)?;
    println!("Saved to {layer_str} layer");
    Ok(())
}

/// Get config for a specific layer.
fn get_config_for_layer<'a>(state: &'a AppState, layer: &str) -> Result<&'a OpenCodeConfig> {
    match layer.to_lowercase().as_str() {
        "merged" => Ok(&state.merged_config),
        "global" => state
            .global_config
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Global config not found")),
        "project" => state
            .project_config
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Project config not found")),
        "custom" => state
            .custom_config
            .as_ref()
            .ok_or_else(|| anyhow::anyhow!("Custom config not found")),
        _ => Err(anyhow::anyhow!(
            "Invalid layer '{}'. Must be one of: merged, global, project, custom",
            layer
        )),
    }
}

// ---------------------------------------------------------------------------
// Agent-config CLI commands
// ---------------------------------------------------------------------------

/// Agent info for JSON output.
#[derive(Serialize)]
struct AgentInfo {
    id: String,
    name: Option<String>,
    model: Option<String>,
    mode: Option<String>,
    disabled: Option<bool>,
}

/// Build an AgentConfigManager, optionally overriding the config path.
fn build_agent_manager(custom_config: Option<&str>) -> Result<omo_config::AgentConfigManager> {
    use omo_config::AgentConfigManager;
    use std::path::PathBuf;

    let mut manager =
        AgentConfigManager::new().context("Failed to initialize agent config manager")?;

    if let Some(path_str) = custom_config {
        let path = PathBuf::from(path_str);
        if !path.exists() {
            return Err(anyhow::anyhow!(
                "Custom agent config not found: {}",
                path.display()
            ));
        }
        manager.project_path = Some(path);
    }

    manager.load_all().context("Failed to load agent configs")?;
    Ok(manager)
}

/// Parse agent config layer string.
fn parse_agent_layer(layer: &str) -> Result<(Option<omo_config::ConfigLayer>, bool)> {
    match layer.to_lowercase().as_str() {
        "merged" => Ok((None, true)),
        "global" => Ok((Some(omo_config::ConfigLayer::Global), false)),
        "project" => Ok((Some(omo_config::ConfigLayer::Project), false)),
        _ => Err(anyhow::anyhow!(
            "Invalid layer '{}'. Must be one of: merged, global, project",
            layer
        )),
    }
}

/// Show agent config as JSON.
fn run_agent_show(layer_str: &str, custom_config: Option<&str>) -> Result<()> {
    let manager = build_agent_manager(custom_config)?;
    let (layer, is_merged) = parse_agent_layer(layer_str)?;

    let config = if is_merged {
        let global = manager.global_config.clone().unwrap_or_default();
        let project = manager.project_config.clone().unwrap_or_default();
        omo_config::merge_agent_configs(&[global, project])
    } else {
        manager.load_layer(layer.unwrap())?.unwrap_or_default()
    };

    let json = serde_json::to_string_pretty(&config).context("Failed to serialize agent config")?;
    println!("{}", json);
    Ok(())
}

/// Validate agent config files.
fn run_agent_validate(custom_config: Option<&str>) -> Result<()> {
    let manager = build_agent_manager(custom_config)?;
    let mut has_errors = false;

    // Try to fetch available models for enhanced validation
    let available_models = fetch_available_models().ok();
    if available_models.is_some() {
        println!("Using 'opencode models' for model availability validation");
    }

    let models_ref = available_models.as_ref();

    if let Some(ref global) = manager.global_config {
        match omo_config::validate_agent_config_with_models(global, models_ref) {
            Ok(()) => println!("Global agent config: OK"),
            Err(e) => {
                eprintln!("Global agent config error: {}", e);
                has_errors = true;
            }
        }
    } else {
        println!("Global agent config: not found");
    }

    if let Some(ref project) = manager.project_config {
        match omo_config::validate_agent_config_with_models(project, models_ref) {
            Ok(()) => println!("Project agent config: OK"),
            Err(e) => {
                eprintln!("Project agent config error: {}", e);
                has_errors = true;
            }
        }
    } else {
        println!("Project agent config: not found");
    }

    // Validate merged config
    let global = manager.global_config.clone().unwrap_or_default();
    let project = manager.project_config.clone().unwrap_or_default();
    let merged = omo_config::merge_agent_configs(&[global, project]);
    match omo_config::validate_agent_config_with_models(&merged, models_ref) {
        Ok(()) => println!("Merged agent config: OK"),
        Err(e) => {
            eprintln!("Merged agent config error: {}", e);
            has_errors = true;
        }
    }

    if has_errors {
        process::exit(1);
    }
    Ok(())
}

/// List all configured agents.
fn run_agent_list(layer_str: &str, custom_config: Option<&str>) -> Result<()> {
    let manager = build_agent_manager(custom_config)?;
    let (layer, is_merged) = parse_agent_layer(layer_str)?;

    let config = if is_merged {
        let global = manager.global_config.clone().unwrap_or_default();
        let project = manager.project_config.clone().unwrap_or_default();
        omo_config::merge_agent_configs(&[global, project])
    } else {
        manager.load_layer(layer.unwrap())?.unwrap_or_default()
    };

    let mut agents: Vec<AgentInfo> = Vec::new();

    if let Some(ref agent_configs) = config.agents {
        let mut push_agent = |id: &str, agent: &omo_config::AgentDefinition| {
            agents.push(AgentInfo {
                id: id.to_string(),
                name: agent.display_name.clone(),
                model: agent.model.clone(),
                mode: agent
                    .mode
                    .as_ref()
                    .map(|m| format!("{:?}", m).to_lowercase()),
                disabled: agent.disable,
            });
        };

        if let Some(ref build) = agent_configs.build {
            push_agent("build", build);
        }
        if let Some(ref plan) = agent_configs.plan {
            push_agent("plan", plan);
        }
        if let Some(ref sisyphus) = agent_configs.sisyphus {
            push_agent("sisyphus", sisyphus);
        }
        if let Some(ref hephaestus) = agent_configs.hephaestus {
            push_agent("hephaestus", hephaestus);
        }
        if let Some(ref prometheus) = agent_configs.prometheus {
            push_agent("prometheus", prometheus);
        }
        if let Some(ref oracle) = agent_configs.oracle {
            push_agent("oracle", oracle);
        }
        if let Some(ref librarian) = agent_configs.librarian {
            push_agent("librarian", librarian);
        }
        if let Some(ref explore) = agent_configs.explore {
            push_agent("explore", explore);
        }
        if let Some(ref multimodal_looker) = agent_configs.multimodal_looker {
            push_agent("multimodal-looker", multimodal_looker);
        }
        if let Some(ref metis) = agent_configs.metis {
            push_agent("metis", metis);
        }
        if let Some(ref momus) = agent_configs.momus {
            push_agent("momus", momus);
        }
        if let Some(ref atlas) = agent_configs.atlas {
            push_agent("atlas", atlas);
        }

        for (id, agent) in &agent_configs.custom {
            agents.push(AgentInfo {
                id: id.clone(),
                name: agent.display_name.clone(),
                model: agent.model.clone(),
                mode: agent
                    .mode
                    .as_ref()
                    .map(|m| format!("{:?}", m).to_lowercase()),
                disabled: agent.disable,
            });
        }
    }

    let json = serde_json::to_string_pretty(&agents).context("Failed to serialize agent list")?;
    println!("{}", json);
    Ok(())
}

/// Run `opencode models` and parse output into a set of available model IDs.
fn fetch_available_models() -> Result<std::collections::HashSet<String>> {
    // Try different ways to invoke opencode CLI for cross-platform compatibility
    let output = if cfg!(target_os = "windows") {
        // On Windows, try opencode.cmd first, then fallback to cmd /c
        std::process::Command::new("opencode.cmd")
            .args(["models"])
            .output()
            .or_else(|_| {
                std::process::Command::new("cmd")
                    .args(["/c", "opencode", "models"])
                    .output()
            })
            .context("Failed to run 'opencode models'. Is opencode CLI installed and in PATH?")?
    } else {
        std::process::Command::new("opencode")
            .args(["models"])
            .output()
            .context("Failed to run 'opencode models'. Is opencode CLI installed and in PATH?")?
    };

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err(anyhow::anyhow!(
            "'opencode models' exited with code {:?}: {}",
            output.status.code(),
            stderr
        ));
    }

    let stdout = String::from_utf8_lossy(&output.stdout);
    let models: std::collections::HashSet<String> = stdout
        .lines()
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect();

    Ok(models)
}

/// List available models from opencode.
fn run_agent_list_available_models() -> Result<()> {
    let models = fetch_available_models()?;
    let model_list: Vec<String> = models.into_iter().collect();
    let json =
        serde_json::to_string_pretty(&model_list).context("Failed to serialize model list")?;
    println!("{}", json);
    Ok(())
}

/// Set model for an agent.
fn run_agent_set_model(
    agent_id: &str,
    model: &str,
    fallback: Option<&str>,
    layer_str: &str,
    custom_config: Option<&str>,
) -> Result<()> {
    // Validate model exists
    let available = fetch_available_models()?;
    let base_model = model.split(':').next().unwrap_or(model);
    if !available.contains(base_model) && !available.contains(model) {
        return Err(anyhow::anyhow!(
            "Model '{}' not found in available models. Run 'opm agent-config list-available-models' to see available models.",
            model
        ));
    }

    // Parse layer
    let layer = match layer_str.to_lowercase().as_str() {
        "global" => omo_config::ConfigLayer::Global,
        "project" => omo_config::ConfigLayer::Project,
        other => {
            return Err(anyhow::anyhow!(
                "Invalid layer '{}'. Must be one of: global, project",
                other
            ));
        }
    };

    // Load existing config
    let manager = build_agent_manager(custom_config)?;
    let mut config = manager.load_layer(layer)?.unwrap_or_default();

    // Ensure agents exists
    if config.agents.is_none() {
        config.agents = Some(omo_config::AgentsConfig::default());
    }
    let agents = config.agents.as_mut().unwrap();

    // Get or create agent definition
    let agent: &mut omo_config::AgentDefinition = match agent_id {
        "build" => agents
            .build
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "plan" => agents
            .plan
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "sisyphus" => agents
            .sisyphus
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "hephaestus" => agents
            .hephaestus
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "prometheus" => agents
            .prometheus
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "oracle" => agents
            .oracle
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "librarian" => agents
            .librarian
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "explore" => agents
            .explore
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "multimodal-looker" => agents
            .multimodal_looker
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "metis" => agents
            .metis
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "momus" => agents
            .momus
            .get_or_insert_with(omo_config::AgentDefinition::default),
        "atlas" => agents
            .atlas
            .get_or_insert_with(omo_config::AgentDefinition::default),
        custom => agents
            .custom
            .entry(custom.to_string())
            .or_insert_with(omo_config::AgentDefinition::default),
    };

    // Set model
    agent.model = Some(model.to_string());

    // Set fallback if provided
    if let Some(fb_str) = fallback {
        let fb_ids: Vec<String> = fb_str.split(',').map(|s| s.trim().to_string()).collect();
        for fb_id in &fb_ids {
            let base_fb = fb_id.split(':').next().unwrap_or(fb_id);
            if !available.contains(base_fb) && !available.contains(fb_id) {
                return Err(anyhow::anyhow!(
                    "Fallback model '{}' not found in available models.",
                    fb_id
                ));
            }
        }
        agent.fallback_models = Some(omo_config::FallbackModels::StringList(fb_ids));
    }

    // Save
    manager.save(layer, &config)?;
    println!(
        "Set model for agent '{}' to '{}' in {} layer",
        agent_id, model, layer_str
    );
    if let Some(fb) = fallback {
        println!("Fallback models: {}", fb);
    }

    Ok(())
}