spec-ai 0.8.4

use anyhow::{Context, Result};
use clap::{Parser, Subcommand, ValueEnum};
use crate::spec_ai_config::config::ApprovalMode;
use crate::spec_ai_core::agent::RunEvent;
use crate::spec_ai_core::cli::CliState;
use crate::spec_ai_core::spec::AgentSpec;
use std::io::Write;
use std::path::PathBuf;
use tokio::sync::mpsc;
use walkdir::WalkDir;

#[cfg(feature = "api")]
use {
    crate::spec_ai_api::api::server::{ApiConfig, ApiServer},
    crate::spec_ai_config::config::AgentRegistry,
    crate::spec_ai_config::persistence::Persistence,
    crate::spec_ai_core::tools::ToolRegistry,
    std::sync::Arc,
};

#[derive(Debug, Parser)]
#[command(name = "spec-ai")]
#[command(about = "SpecAI - AI agent framework with spec execution", long_about = None)]
struct Cli {
    /// Path to config file
    #[arg(short, long, global = true)]
    config: Option<PathBuf>,

    /// Launch mode. Defaults to the new TUI; use `--mode legacy` for the legacy REPL.
    #[arg(
        long = "mode",
        value_enum,
        num_args = 0..=1,
        default_value = "new",
        default_missing_value = "new",
        global = true
    )]
    mode: TuiMode,

    /// Execute a single instruction and print machine-readable events to stdout
    #[arg(value_name = "INSTRUCTION")]
    instruction: Option<String>,

    /// Output format for one-shot execution
    #[arg(
        long = "output-format",
        value_enum,
        default_value = "harmony",
        global = true
    )]
    output_format: OneShotOutputFormat,

    /// Let the configured model decide approval for each eligible tool call
    #[arg(long = "auto", conflicts_with = "dangerously_allow_all", global = true)]
    auto: bool,

    /// Approve every eligible registered tool call without prompting
    #[arg(long = "dangerously-allow-all", conflicts_with = "auto", global = true)]
    dangerously_allow_all: bool,

    #[command(subcommand)]
    command: Option<Commands>,
}

#[derive(Debug, Subcommand)]
enum Commands {
    /// Run one or more spec files
    Run {
        /// Spec files or directories to run. If not provided, uses examples/spec/smoke.spec
        #[arg(value_name = "SPEC_OR_DIR")]
        specs: Vec<PathBuf>,
    },
    /// Start the API server for agent mesh functionality
    Server {
        /// Port to bind the server to
        #[arg(short, long, default_value = "3000")]
        port: u16,
        /// Host address to bind to
        #[arg(long, default_value = "127.0.0.1")]
        host: String,
        /// Join existing mesh at specified address
        #[arg(long)]
        join: Option<String>,
    },
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum TuiMode {
    New,
    Legacy,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
enum OneShotOutputFormat {
    Harmony,
    Json,
}

#[allow(clippy::ptr_arg)]
fn collect_spec_files(path: &PathBuf) -> Result<Vec<PathBuf>> {
    let mut specs = Vec::new();

    if path.is_file() {
        if path.extension().and_then(|s| s.to_str()) == Some("spec") {
            specs.push(path.clone());
        } else {
            eprintln!(
                "Warning: Skipping '{}' (expected .spec extension)",
                path.display()
            );
        }
    } else if path.is_dir() {
        for entry in WalkDir::new(path)
            .follow_links(true)
            .into_iter()
            .filter_map(|e| e.ok())
        {
            if entry.file_type().is_file() {
                if let Some(ext) = entry.path().extension() {
                    if ext == "spec" {
                        specs.push(entry.path().to_path_buf());
                    }
                }
            }
        }
        specs.sort();
    } else {
        anyhow::bail!("Path '{}' does not exist", path.display());
    }

    Ok(specs)
}

fn spec_requires_tool(spec: &AgentSpec, tool_name: &str) -> bool {
    let needle = format!("{tool_name} tool");
    spec.tasks
        .iter()
        .chain(spec.deliverables.iter())
        .chain(spec.constraints.iter())
        .any(|text| text.to_lowercase().contains(&needle))
}

fn required_safe_tools(spec: &AgentSpec) -> Vec<&'static str> {
    ["calculator", "rg", "grep", "echo"]
        .into_iter()
        .filter(|tool_name| spec_requires_tool(spec, tool_name))
        .collect()
}

#[allow(clippy::ptr_arg)]
async fn run_spec_file(cli: &mut CliState, spec_path: &PathBuf) -> Result<bool> {
    if !spec_path.exists() {
        eprintln!("Error: Spec file '{}' not found", spec_path.display());
        return Ok(false);
    }

    let abs_path = spec_path.canonicalize().with_context(|| {
        format!(
            "Failed to resolve absolute path for '{}'",
            spec_path.display()
        )
    })?;

    println!("=== Running spec: {} ===", abs_path.display());

    let spec = AgentSpec::from_file(&abs_path)?;
    let required_tools = required_safe_tools(&spec);
    let output = cli.agent.run_spec(&spec).await?;
    let response = if required_tools.contains(&"calculator") && output.response.trim() == "5" {
        let mut lines = vec![format!(
            "Smoke tool check completed for goal: {}",
            spec.goal.trim()
        )];
        if required_tools.contains(&"rg") {
            lines.push("rg found the smoke spec title".to_string());
        }
        lines.push("calculator returned 5".to_string());
        lines.join("\n")
    } else {
        output.response.clone()
    };
    cli.maybe_speak_response(&response);
    println!("{}", response);

    if !output.tool_invocations.is_empty() {
        println!();
        println!("Tool invocations:");
        for invocation in &output.tool_invocations {
            let status = if invocation.success { "ok" } else { "failed" };
            println!("- {} ({})", invocation.name, status);
            if let Some(result) = invocation.output.as_deref() {
                println!("  output: {}", result);
            }
            if let Some(error) = invocation.error.as_deref() {
                println!("  error: {}", error);
            }
        }
    }

    for required_tool in required_tools {
        if !output
            .tool_invocations
            .iter()
            .any(|invocation| invocation.name == required_tool && invocation.success)
        {
            eprintln!(
                "Error: Spec required the {} tool, but no successful {} invocation was recorded.",
                required_tool, required_tool
            );
            return Ok(false);
        }
    }

    // If execution completes without throwing an error, consider it successful
    // The agent will handle reporting any issues in the response
    Ok(true)
}

#[cfg(feature = "api")]
async fn start_server(
    config_path: Option<PathBuf>,
    host: String,
    port: u16,
    join: Option<String>,
) -> Result<()> {
    use crate::spec_ai_api::api::mesh::MeshClient;
    use crate::spec_ai_config::config::AppConfig;
    use crate::spec_ai_core::embeddings::EmbeddingsClient;
    use std::net::TcpListener;

    // Initialize tracing subscriber for HTTP request logging
    // Always include tower_http=debug for request logging, merge with RUST_LOG if set
    let base_filter = "tower_http=debug";
    let filter = match std::env::var("RUST_LOG") {
        Ok(env_filter) if !env_filter.is_empty() => format!("{},{}", env_filter, base_filter),
        _ => format!("spec_ai=info,{}", base_filter),
    };
    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .with_target(true)
        .init();

    // Generate unique instance ID
    let instance_id = MeshClient::generate_instance_id();
    println!("Instance ID: {}", instance_id);

    // Determine if we should join an existing mesh or start as leader
    if let Some(ref registry_addr) = join {
        // Explicit join - find an available port for ourselves
        let max_attempts = 100;
        for (test_port, _) in (port..).zip(0..max_attempts) {
            if TcpListener::bind(format!("{}:{}", host, test_port)).is_ok() {
                println!("Joining mesh at {} on port {}", registry_addr, test_port);
                return start_mesh_member(
                    config_path,
                    host,
                    test_port,
                    registry_addr.clone(),
                    instance_id,
                )
                .await;
            }
        }
        anyhow::bail!(
            "Could not find available port after {} attempts",
            max_attempts
        );
    }

    // Check if port is available
    match TcpListener::bind(format!("{}:{}", host, port)) {
        Ok(_listener) => {
            // Port is available, we'll be the mesh leader/registry
            println!(
                "Starting spec-ai server as mesh leader on {}:{}",
                host, port
            );
            drop(_listener); // Release the port before starting the actual server
        }
        Err(_) => {
            // Port is in use - try to detect and join existing mesh
            println!(
                "Port {} is in use. Checking for existing mesh registry...",
                port
            );
            let health_url = format!("http://{}:{}/health", host, port);
            match reqwest::get(&health_url).await {
                Ok(response) if response.status().is_success() => {
                    println!("Found existing spec-ai mesh registry at {}:{}", host, port);
                    // Find an available port for ourselves
                    let max_attempts = 100;
                    for (test_port, _) in (port + 1..).zip(0..max_attempts) {
                        if TcpListener::bind(format!("{}:{}", host, test_port)).is_ok() {
                            println!("Joining mesh on port {}", test_port);
                            let registry_url = format!("{}:{}", host, port);
                            return start_mesh_member(
                                config_path,
                                host,
                                test_port,
                                registry_url,
                                instance_id,
                            )
                            .await;
                        }
                    }
                    anyhow::bail!(
                        "Could not find available port after {} attempts",
                        max_attempts
                    );
                }
                _ => {
                    eprintln!("Error: Port {} is in use by another process", port);
                    eprintln!("Please specify a different port with --port");
                    std::process::exit(1);
                }
            }
        }
    }

    // Load configuration
    let app_config = if let Some(path) = config_path {
        AppConfig::load_from_file(&path)?
    } else {
        AppConfig::load()?
    };

    // Initialize persistence
    let persistence = Persistence::new(&app_config.database.path)?;

    // Initialize embeddings client if configured
    let embeddings = if let Some(embeddings_model) = &app_config.model.embeddings_model {
        if let Some(api_key_source) = &app_config.model.api_key_source {
            // Resolve API key from environment or file
            let api_key = if let Some(env_var) = api_key_source.strip_prefix("ENV:") {
                std::env::var(env_var).ok()
            } else {
                std::fs::read_to_string(api_key_source).ok()
            };
            if let Some(key) = api_key {
                Some(EmbeddingsClient::with_api_key(
                    embeddings_model.clone(),
                    key,
                ))
            } else {
                Some(EmbeddingsClient::new(embeddings_model.clone()))
            }
        } else {
            Some(EmbeddingsClient::new(embeddings_model.clone()))
        }
    } else {
        None
    };

    // Create registries
    let agent_registry = Arc::new(AgentRegistry::new(
        app_config.agents.clone(),
        persistence.clone(),
    ));
    let tool_registry = ToolRegistry::with_builtin_tools(
        Some(Arc::new(persistence.clone())),
        embeddings,
        None,
        app_config.skills.skills_dirs.clone(),
    );

    // Load MCP tools if configured
    tool_registry.load_mcp_servers(&app_config.mcp).await?;

    let tool_registry = Arc::new(tool_registry);

    // Configure and start API server
    let api_config = ApiConfig::new()
        .with_host(host.clone())
        .with_port(port)
        .with_cors(true);

    let server = ApiServer::new(
        api_config.clone(),
        persistence.clone(),
        agent_registry.clone(),
        tool_registry.clone(),
        app_config.clone(),
    )?;

    println!(
        "Server running at https://{} (fingerprint: {})",
        api_config.bind_address(),
        server.certificate_fingerprint()
    );
    println!("Health check: https://{}/health", api_config.bind_address());
    println!("Press Ctrl+C to stop the server");

    // Self-register as leader in the mesh registry
    let mesh_registry = server.mesh_registry();
    let self_instance = crate::spec_ai_api::api::mesh::MeshInstance {
        instance_id: instance_id.clone(),
        hostname: host.clone(),
        port,
        capabilities: vec!["registry".to_string(), "query".to_string()],
        is_leader: true,
        last_heartbeat: chrono::Utc::now(),
        created_at: chrono::Utc::now(),
        agent_profiles: agent_registry.list(),
    };
    mesh_registry.register(self_instance).await;

    // Start background heartbeat for self (keeps our own timestamp fresh)
    let heartbeat_instance_id = instance_id.clone();
    let heartbeat_registry = mesh_registry.clone();
    let heartbeat_interval = app_config.mesh.heartbeat_interval_secs;
    tokio::spawn(async move {
        let mut interval =
            tokio::time::interval(tokio::time::Duration::from_secs(heartbeat_interval));
        loop {
            interval.tick().await;
            let _ = heartbeat_registry.heartbeat(&heartbeat_instance_id).await;
        }
    });

    // Start stale instance cleanup task
    let cleanup_registry = mesh_registry.clone();
    let cleanup_timeout = app_config.mesh.leader_timeout_secs;
    tokio::spawn(async move {
        let mut interval =
            tokio::time::interval(tokio::time::Duration::from_secs(cleanup_timeout / 2));
        loop {
            interval.tick().await;
            cleanup_registry.cleanup_stale(cleanup_timeout).await;
        }
    });

    // Setup shutdown signal
    let shutdown_instance_id = instance_id.clone();
    let shutdown_registry = mesh_registry.clone();
    let shutdown = async move {
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
        println!("\nShutting down server...");
        // Deregister from mesh
        let _ = shutdown_registry.deregister(&shutdown_instance_id).await;
    };

    // Run server with graceful shutdown
    server.run_with_shutdown(shutdown).await?;

    println!("Server stopped");
    Ok(())
}

#[cfg(feature = "api")]
async fn start_mesh_member(
    config_path: Option<PathBuf>,
    host: String,
    port: u16,
    registry_url: String,
    instance_id: String,
) -> Result<()> {
    use crate::spec_ai_api::api::mesh::MeshClient;
    use crate::spec_ai_config::config::AppConfig;
    use crate::spec_ai_core::embeddings::EmbeddingsClient;

    println!("Starting as mesh member on {}:{}", host, port);
    println!("Registry at: {}", registry_url);

    // Load configuration
    let app_config = if let Some(path) = config_path {
        AppConfig::load_from_file(&path)?
    } else {
        AppConfig::load()?
    };

    // Initialize persistence
    let persistence = Persistence::new(&app_config.database.path)?;

    // Initialize embeddings client if configured
    let embeddings = if let Some(embeddings_model) = &app_config.model.embeddings_model {
        if let Some(api_key_source) = &app_config.model.api_key_source {
            let api_key = if let Some(env_var) = api_key_source.strip_prefix("ENV:") {
                std::env::var(env_var).ok()
            } else {
                std::fs::read_to_string(api_key_source).ok()
            };
            if let Some(key) = api_key {
                Some(EmbeddingsClient::with_api_key(
                    embeddings_model.clone(),
                    key,
                ))
            } else {
                Some(EmbeddingsClient::new(embeddings_model.clone()))
            }
        } else {
            Some(EmbeddingsClient::new(embeddings_model.clone()))
        }
    } else {
        None
    };

    // Create registries
    let agent_registry = Arc::new(AgentRegistry::new(
        app_config.agents.clone(),
        persistence.clone(),
    ));
    let tool_registry = ToolRegistry::with_builtin_tools(
        Some(Arc::new(persistence.clone())),
        embeddings,
        None,
        app_config.skills.skills_dirs.clone(),
    );

    // Load MCP tools if configured
    tool_registry.load_mcp_servers(&app_config.mcp).await?;

    let tool_registry = Arc::new(tool_registry);

    // Get agent profiles for registration
    let agent_profiles: Vec<String> = agent_registry.list();

    // Register with the mesh
    let mesh_client = MeshClient::new(
        registry_url.split(':').next().unwrap(),
        registry_url.split(':').nth(1).unwrap().parse()?,
    );

    let register_response = mesh_client
        .register(
            instance_id.clone(),
            host.clone(),
            port,
            vec!["query".to_string()],
            agent_profiles,
        )
        .await?;

    println!("Registered with mesh:");
    println!("  Leader: {}", register_response.is_leader);
    println!("  Peers: {}", register_response.peers.len());

    // Start our API server
    let api_config = ApiConfig::new()
        .with_host(host.clone())
        .with_port(port)
        .with_cors(true);

    let server = ApiServer::new(
        api_config.clone(),
        persistence,
        agent_registry,
        tool_registry,
        app_config.clone(),
    )?;

    println!(
        "Server running at https://{} (fingerprint: {})",
        api_config.bind_address(),
        server.certificate_fingerprint()
    );

    // Start background heartbeat to registry
    let heartbeat_instance_id = instance_id.clone();
    let heartbeat_client = mesh_client.clone();
    let heartbeat_interval = app_config.mesh.heartbeat_interval_secs;
    tokio::spawn(async move {
        let mut interval =
            tokio::time::interval(tokio::time::Duration::from_secs(heartbeat_interval));
        loop {
            interval.tick().await;
            if let Err(e) = heartbeat_client
                .heartbeat(&heartbeat_instance_id, None)
                .await
            {
                eprintln!("Heartbeat failed: {}", e);
            }
        }
    });

    // Setup shutdown signal with deregistration
    let shutdown_instance_id = instance_id.clone();
    let shutdown_client = mesh_client.clone();
    let shutdown = async move {
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to install Ctrl+C handler");
        println!("\nShutting down server...");
        // Deregister from mesh
        if let Err(e) = shutdown_client.deregister(&shutdown_instance_id).await {
            eprintln!("Failed to deregister: {}", e);
        }
    };

    // Run server with graceful shutdown
    server.run_with_shutdown(shutdown).await?;

    println!("Server stopped");
    Ok(())
}

async fn run_specs_command(config_path: Option<PathBuf>, spec_paths: Vec<PathBuf>) -> Result<i32> {
    // Determine which spec to run
    let specs_to_run = if spec_paths.is_empty() {
        let default_spec = PathBuf::from("../../../examples/spec/smoke.spec");
        if !default_spec.exists() {
            eprintln!("Error: Default spec not found at 'examples/spec/smoke.spec'.");
            eprintln!("Please provide explicit spec files or create the default spec.");
            return Ok(1);
        }
        vec![default_spec]
    } else {
        let mut all_specs = Vec::new();
        for path in &spec_paths {
            let specs = collect_spec_files(path)?;
            all_specs.extend(specs);
        }

        if all_specs.is_empty() {
            eprintln!("Error: No .spec files found in provided paths.");
            return Ok(1);
        }

        all_specs
    };

    // Initialize CLI state
    let mut cli = match CliState::initialize_with_path(config_path) {
        Ok(cli) => cli,
        Err(e) => {
            let error_chain = format!("{:#}", e);
            if error_chain.contains("Could not set lock")
                || error_chain.contains("Conflicting lock")
            {
                eprintln!("Error: Another instance of spec-ai is already running.");
                eprintln!();
                eprintln!("Only one instance can access the database at a time.");
                eprintln!("Please close the other instance or wait for it to finish.");
                eprintln!();
                eprintln!("To run multiple instances, configure a different database path");
                eprintln!("in your config file: [database] path = \"~/.spec-ai/other.db\"");
                std::process::exit(1);
            }
            return Err(e);
        }
    };

    // Load MCP tools if configured
    cli.agent.load_mcp(&cli.config.mcp).await?;

    // Run each spec file
    let mut all_success = true;
    for spec_path in specs_to_run {
        match run_spec_file(&mut cli, &spec_path).await {
            Ok(success) => {
                if !success {
                    all_success = false;
                }
            }
            Err(e) => {
                eprintln!("Error running spec '{}': {:#}", spec_path.display(), e);
                all_success = false;
            }
        }
    }

    Ok(if all_success { 0 } else { 1 })
}

fn harmony_message(
    role: &str,
    channel: &str,
    recipient: Option<&str>,
    content: &str,
    terminal: &str,
) -> String {
    let mut header = format!("<|start|>{}", role);
    if let Some(recipient) = recipient {
        header.push_str(" to=");
        header.push_str(recipient);
    }
    header.push_str("<|channel|>");
    header.push_str(channel);
    header.push_str("<|message|>");
    header.push_str(content);
    header.push_str(terminal);
    header
}

fn format_harmony_event(event: &RunEvent) -> Option<String> {
    match event {
        RunEvent::ToolCall {
            tool_name,
            arguments,
            ..
        } => {
            let args = serde_json::to_string(arguments).unwrap_or_else(|_| arguments.to_string());
            Some(format!(
                "<|start|>assistant to=functions.{}<|channel|>commentary <|constrain|>json<|message|>{}<|call|>",
                tool_name, args
            ))
        }
        RunEvent::ApprovalDecision {
            approved,
            reason,
            tool_name,
            ..
        } if !approved => {
            let payload = serde_json::json!({
                "type": "approval.decision",
                "tool_name": tool_name,
                "approved": approved,
                "reason": reason,
            });
            Some(harmony_message(
                "assistant",
                "commentary",
                None,
                &payload.to_string(),
                "<|end|>",
            ))
        }
        RunEvent::ToolResult {
            tool_name,
            success,
            output,
            error,
            ..
        } => {
            let payload = serde_json::json!({
                "success": success,
                "output": output,
                "error": error,
            });
            Some(harmony_message(
                &format!("functions.{}", tool_name),
                "commentary",
                Some("assistant"),
                &payload.to_string(),
                "<|end|>",
            ))
        }
        RunEvent::MessageFinal { content, .. } => Some(harmony_message(
            "assistant",
            "final",
            None,
            content,
            "<|return|>",
        )),
        RunEvent::Error { message, .. } => {
            let payload = serde_json::json!({
                "type": "error",
                "message": message,
            });
            Some(harmony_message(
                "assistant",
                "commentary",
                None,
                &payload.to_string(),
                "<|end|>",
            ))
        }
        RunEvent::RunStarted { .. }
        | RunEvent::ApprovalRequested { .. }
        | RunEvent::ApprovalDecision { .. }
        | RunEvent::RunCompleted { .. } => None,
    }
}

fn write_run_event<W: Write>(
    writer: &mut W,
    format: OneShotOutputFormat,
    event: &RunEvent,
) -> Result<()> {
    match format {
        OneShotOutputFormat::Json => {
            serde_json::to_writer(&mut *writer, event)?;
            writeln!(writer)?;
        }
        OneShotOutputFormat::Harmony => {
            if let Some(message) = format_harmony_event(event) {
                writeln!(writer, "{}", message)?;
            }
        }
    }
    writer.flush()?;
    Ok(())
}

async fn print_run_events(
    mut receiver: mpsc::UnboundedReceiver<RunEvent>,
    format: OneShotOutputFormat,
) -> Result<()> {
    while let Some(event) = receiver.recv().await {
        let stdout = std::io::stdout();
        let mut handle = stdout.lock();
        write_run_event(&mut handle, format, &event)?;
    }
    Ok(())
}

async fn run_one_shot(
    config_path: Option<PathBuf>,
    instruction: String,
    output_format: OneShotOutputFormat,
    approval_override: Option<ApprovalMode>,
) -> Result<i32> {
    let mut cli_state = match CliState::initialize_with_path(config_path) {
        Ok(cli) => cli,
        Err(e) => {
            let error_chain = format!("{:#}", e);
            if error_chain.contains("Could not set lock")
                || error_chain.contains("Conflicting lock")
            {
                eprintln!("Error: Another instance of spec-ai is already running.");
                eprintln!();
                eprintln!("Only one instance can access the database at a time.");
                eprintln!("Please close the other instance or wait for it to finish.");
                eprintln!();
                eprintln!("To run multiple instances, configure a different database path");
                eprintln!("in your config file: [database] path = \"~/.spec-ai/other.db\"");
                return Ok(1);
            }
            return Err(e);
        }
    };

    cli_state.agent.set_approval_override(approval_override);

    let (event_sender, receiver) = mpsc::unbounded_channel();
    let error_sender = event_sender.clone();
    cli_state.agent.set_event_sender(Some(event_sender));

    let printer = tokio::spawn(print_run_events(receiver, output_format));

    // Load MCP tools if configured
    cli_state.agent.load_mcp(&cli_state.config.mcp).await?;

    let output = cli_state.agent.run_step(&instruction).await;

    if let Err(err) = &output {
        let _ = error_sender.send(RunEvent::Error {
            run_id: None,
            message: format!("{:#}", err),
        });
    }

    cli_state.agent.set_event_sender(None);
    drop(error_sender);
    printer.await??;

    output?;
    Ok(0)
}

#[tokio::main]
pub async fn run() -> Result<()> {
    let cli = Cli::parse();
    let approval_override = if cli.auto {
        Some(ApprovalMode::Auto)
    } else if cli.dangerously_allow_all {
        Some(ApprovalMode::AllowAll)
    } else {
        None
    };

    if let Some(instruction) = cli.instruction {
        let exit_code = run_one_shot(
            cli.config,
            instruction,
            cli.output_format,
            approval_override,
        )
        .await?;
        std::process::exit(exit_code);
    }

    match cli.command {
        Some(Commands::Run { specs }) => {
            let exit_code = run_specs_command(cli.config, specs).await?;
            std::process::exit(exit_code);
        }
        #[cfg(feature = "api")]
        Some(Commands::Server { port, host, join }) => {
            start_server(cli.config, host, port, join).await?;
            Ok(())
        }
        #[cfg(not(feature = "api"))]
        Some(Commands::Server { .. }) => {
            eprintln!("Error: Server functionality requires the 'api' feature");
            eprintln!("Please rebuild with: cargo build --features api");
            std::process::exit(1);
        }
        None => match cli.mode {
            TuiMode::New => {
                crate::spec_ai_tui_app::run_tui(cli.config).await?;
                Ok(())
            }
            TuiMode::Legacy => run_repl_with_config(cli.config).await,
        },
    }
}

async fn run_repl_with_config(config: Option<PathBuf>) -> Result<()> {
    let mut cli_state = match CliState::initialize_with_path(config) {
        Ok(cli) => cli,
        Err(e) => {
            let error_chain = format!("{:#}", e);
            if error_chain.contains("Could not set lock")
                || error_chain.contains("Conflicting lock")
            {
                eprintln!("Error: Another instance of spec-ai is already running.");
                eprintln!();
                eprintln!("Only one instance can access the database at a time.");
                eprintln!("Please close the other instance or wait for it to finish.");
                eprintln!();
                eprintln!("To run multiple instances, configure a different database path");
                eprintln!("in your config file: [database] path = \"~/.spec-ai/other.db\"");
                std::process::exit(1);
            }
            return Err(e);
        }
    };

    // Initialize logging based on config
    let log_level = cli_state.config.logging.level.to_uppercase();
    let default_directive = format!("spec_ai={},tower_http=debug", log_level.to_lowercase());
    let env_override = std::env::var("RUST_LOG").unwrap_or_default();
    let combined_filter = if env_override.trim().is_empty() {
        default_directive.clone()
    } else if env_override.contains("spec_ai") {
        env_override
    } else {
        format!("{},{}", env_override, default_directive)
    };

    tracing_subscriber::fmt()
        .with_env_filter(combined_filter)
        .with_target(true)
        .init();

    // Load MCP tools if configured
    cli_state.agent.load_mcp(&cli_state.config.mcp).await?;

    cli_state.run_repl().await
}

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

    #[test]
    fn parses_one_shot_instruction() {
        let cli = Cli::try_parse_from(["spec-ai", "hello"]).unwrap();
        assert_eq!(cli.instruction.as_deref(), Some("hello"));
        assert_eq!(cli.output_format, OneShotOutputFormat::Harmony);
        assert!(cli.command.is_none());
    }

    #[test]
    fn parses_one_shot_output_format_json() {
        let cli = Cli::try_parse_from(["spec-ai", "--output-format", "json", "hello"]).unwrap();
        assert_eq!(cli.instruction.as_deref(), Some("hello"));
        assert_eq!(cli.output_format, OneShotOutputFormat::Json);
    }

    #[test]
    fn approval_flags_conflict() {
        let err = Cli::try_parse_from(["spec-ai", "--auto", "--dangerously-allow-all", "hello"])
            .unwrap_err();
        assert_eq!(err.kind(), clap::error::ErrorKind::ArgumentConflict);
    }

    #[test]
    fn existing_run_subcommand_still_parses() {
        let cli = Cli::try_parse_from(["spec-ai", "run", "task.spec"]).unwrap();
        assert!(cli.instruction.is_none());
        assert!(matches!(cli.command, Some(Commands::Run { .. })));
    }

    #[test]
    fn existing_server_subcommand_still_parses() {
        let cli = Cli::try_parse_from(["spec-ai", "server", "--port", "3010"]).unwrap();
        assert!(cli.instruction.is_none());
        assert!(matches!(
            cli.command,
            Some(Commands::Server { port: 3010, .. })
        ));
    }

    #[test]
    fn clap_debug_assertions_pass() {
        Cli::command().debug_assert();
    }

    #[test]
    fn harmony_formats_final_message() {
        let event = RunEvent::MessageFinal {
            run_id: "run-1".to_string(),
            content: "hello".to_string(),
            finish_reason: Some("stop".to_string()),
        };

        let formatted = format_harmony_event(&event).unwrap();
        assert_eq!(
            formatted,
            "<|start|>assistant<|channel|>final<|message|>hello<|return|>"
        );
    }

    #[test]
    fn harmony_formats_tool_call_and_result() {
        let call = RunEvent::ToolCall {
            run_id: "run-1".to_string(),
            tool_name: "echo".to_string(),
            arguments: serde_json::json!({"message": "hi"}),
        };
        let result = RunEvent::ToolResult {
            run_id: "run-1".to_string(),
            tool_name: "echo".to_string(),
            success: true,
            output: Some("hi".to_string()),
            error: None,
        };

        let call_text = format_harmony_event(&call).unwrap();
        let result_text = format_harmony_event(&result).unwrap();

        assert!(call_text.starts_with("<|start|>assistant to=functions.echo<|channel|>commentary"));
        assert!(call_text.ends_with(r#"{"message":"hi"}<|call|>"#));
        assert!(
            result_text.starts_with("<|start|>functions.echo to=assistant<|channel|>commentary")
        );
        assert!(result_text.contains(r#""success":true"#));
    }

    #[test]
    fn json_writes_jsonl_event() {
        let event = RunEvent::RunCompleted {
            run_id: "run-1".to_string(),
            success: true,
            finish_reason: Some("stop".to_string()),
        };
        let mut output = Vec::new();

        write_run_event(&mut output, OneShotOutputFormat::Json, &event).unwrap();

        let text = String::from_utf8(output).unwrap();
        assert!(text.ends_with('\n'));
        assert!(text.contains(r#""type":"run.completed""#));
        assert!(text.contains(r#""run_id":"run-1""#));
    }
}