optimclaw/src/main.rs

//! IronClaw - Main entry point.

use std::sync::Arc;

use clap::Parser;
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};

use ironclaw::{
    agent::{Agent, AgentDeps},
    channels::{
        ChannelManager, HttpChannel, ReplChannel, WebhookServer, WebhookServerConfig,
        wasm::{
            RegisteredEndpoint, SharedWasmChannel, WasmChannelLoader, WasmChannelRouter,
            WasmChannelRuntime, WasmChannelRuntimeConfig, create_wasm_channel_router,
        },
    },
    cli::{
        Cli, Command, run_mcp_command, run_memory_command, run_status_command, run_tool_command,
    },
    config::Config,
    context::ContextManager,
    extensions::ExtensionManager,
    history::Store,
    llm::{SessionConfig, create_llm_provider, create_session_manager},
    safety::SafetyLayer,
    secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
    settings::Settings,
    setup::{SetupConfig, SetupWizard},
    tools::{
        ToolRegistry,
        mcp::{McpClient, McpSessionManager, config::load_mcp_servers, is_authenticated},
        wasm::{WasmToolLoader, WasmToolRuntime},
    },
    workspace::{EmbeddingProvider, NearAiEmbeddings, OpenAiEmbeddings, Workspace},
};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();

    // Handle non-agent commands first (they don't need full setup)
    match &cli.command {
        Some(Command::Tool(tool_cmd)) => {
            // Simple logging for CLI commands
            tracing_subscriber::fmt()
                .with_env_filter(
                    EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
                )
                .init();

            return run_tool_command(tool_cmd.clone()).await;
        }
        Some(Command::Config(config_cmd)) => {
            // Config commands don't need logging setup
            return ironclaw::cli::run_config_command(config_cmd.clone())
                .map_err(|e| anyhow::anyhow!("{}", e));
        }
        Some(Command::Mcp(mcp_cmd)) => {
            // Simple logging for MCP commands
            tracing_subscriber::fmt()
                .with_env_filter(
                    EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
                )
                .init();

            return run_mcp_command(mcp_cmd.clone()).await;
        }
        Some(Command::Memory(mem_cmd)) => {
            tracing_subscriber::fmt()
                .with_env_filter(
                    EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
                )
                .init();

            // Memory commands need database (and optionally embeddings)
            let _ = dotenvy::dotenv();
            let config = Config::from_env().map_err(|e| anyhow::anyhow!("{}", e))?;
            let store = ironclaw::history::Store::new(&config.database).await?;
            store.run_migrations().await?;

            // Set up embeddings if available
            let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
                auth_base_url: config.llm.nearai.auth_base_url.clone(),
                session_path: config.llm.nearai.session_path.clone(),
                ..Default::default()
            })
            .await;

            let embeddings: Option<Arc<dyn ironclaw::workspace::EmbeddingProvider>> =
                if config.embeddings.enabled {
                    match config.embeddings.provider.as_str() {
                        "nearai" => Some(Arc::new(
                            ironclaw::workspace::NearAiEmbeddings::new(
                                &config.llm.nearai.base_url,
                                session,
                            )
                            .with_model(&config.embeddings.model, 1536),
                        )),
                        _ => {
                            if let Some(api_key) = config.embeddings.openai_api_key() {
                                let dim = match config.embeddings.model.as_str() {
                                    "text-embedding-3-large" => 3072,
                                    _ => 1536,
                                };
                                Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model(
                                    api_key,
                                    &config.embeddings.model,
                                    dim,
                                )))
                            } else {
                                None
                            }
                        }
                    }
                } else {
                    None
                };

            return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await;
        }
        Some(Command::Status) => {
            let _ = dotenvy::dotenv();
            tracing_subscriber::fmt()
                .with_env_filter(
                    EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
                )
                .init();

            return run_status_command().await;
        }
        Some(Command::Setup {
            skip_auth,
            channels_only,
        }) => {
            // Load .env before running setup wizard
            let _ = dotenvy::dotenv();

            // Run setup wizard
            let config = SetupConfig {
                skip_auth: *skip_auth,
                channels_only: *channels_only,
            };
            let mut wizard = SetupWizard::with_config(config);
            wizard.run().await?;
            return Ok(());
        }
        None | Some(Command::Run) => {
            // Continue to run agent
        }
    }

    // Load .env if present
    let _ = dotenvy::dotenv();

    // Enhanced first-run detection
    if !cli.no_setup {
        if let Some(reason) = check_setup_needed() {
            println!("Setup needed: {}", reason);
            println!();
            let mut wizard = SetupWizard::new();
            wizard.run().await?;
        }
    }

    // Load configuration (after potential setup)
    let config = match Config::from_env() {
        Ok(c) => c,
        Err(ironclaw::error::ConfigError::MissingRequired { key, hint }) => {
            eprintln!("Configuration error: Missing required setting '{}'", key);
            eprintln!("  {}", hint);
            eprintln!();
            eprintln!(
                "Run 'ironclaw setup' to configure, or set the required environment variables."
            );
            std::process::exit(1);
        }
        Err(e) => return Err(e.into()),
    };

    // Initialize session manager and authenticate before channel setup
    let session_config = SessionConfig {
        auth_base_url: config.llm.nearai.auth_base_url.clone(),
        session_path: config.llm.nearai.session_path.clone(),
        ..Default::default()
    };
    let session = create_session_manager(session_config).await;

    // Ensure we're authenticated before proceeding (may trigger login flow)
    session.ensure_authenticated().await?;

    // Initialize tracing
    let env_filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=debug"));

    tracing_subscriber::registry()
        .with(env_filter)
        .with(tracing_subscriber::fmt::layer().with_target(false))
        .init();

    // Create CLI channel
    let repl_channel = if let Some(ref msg) = cli.message {
        Some(ReplChannel::with_message(msg.clone()))
    } else if config.channels.cli.enabled {
        Some(ReplChannel::new())
    } else {
        None
    };

    tracing::info!("Starting IronClaw...");
    tracing::info!("Loaded configuration for agent: {}", config.agent.name);
    tracing::info!("NEAR AI session authenticated");

    // Initialize database store (optional for testing)
    let store = if cli.no_db {
        tracing::warn!("Running without database connection");
        None
    } else {
        let store = Store::new(&config.database).await?;
        store.run_migrations().await?;
        tracing::info!("Database connected and migrations applied");
        Some(Arc::new(store))
    };

    // Initialize LLM provider (clone session so we can reuse it for embeddings)
    let llm = create_llm_provider(&config.llm, session.clone())?;
    tracing::info!("LLM provider initialized: {}", llm.model_name());

    // Initialize safety layer
    let safety = Arc::new(SafetyLayer::new(&config.safety));
    tracing::info!("Safety layer initialized");

    // Initialize tool registry
    let tools = Arc::new(ToolRegistry::new());
    tools.register_builtin_tools();
    tracing::info!("Registered {} built-in tools", tools.count());

    // Create embeddings provider if configured
    let embeddings: Option<Arc<dyn EmbeddingProvider>> = if config.embeddings.enabled {
        match config.embeddings.provider.as_str() {
            "nearai" => {
                tracing::info!(
                    "Embeddings enabled via NEAR AI (model: {})",
                    config.embeddings.model
                );
                Some(Arc::new(
                    NearAiEmbeddings::new(&config.llm.nearai.base_url, session.clone())
                        .with_model(&config.embeddings.model, 1536),
                ))
            }
            _ => {
                // Default to OpenAI for unknown providers
                if let Some(api_key) = config.embeddings.openai_api_key() {
                    tracing::info!(
                        "Embeddings enabled via OpenAI (model: {})",
                        config.embeddings.model
                    );
                    Some(Arc::new(OpenAiEmbeddings::with_model(
                        api_key,
                        &config.embeddings.model,
                        match config.embeddings.model.as_str() {
                            "text-embedding-3-large" => 3072,
                            _ => 1536, // text-embedding-3-small and ada-002
                        },
                    )))
                } else {
                    tracing::warn!("Embeddings configured but OPENAI_API_KEY not set");
                    None
                }
            }
        }
    } else {
        tracing::info!("Embeddings disabled (set OPENAI_API_KEY or EMBEDDING_ENABLED=true)");
        None
    };

    // Register memory tools if database is available
    if let Some(ref store) = store {
        let mut workspace = Workspace::new("default", store.pool());
        if let Some(ref emb) = embeddings {
            workspace = workspace.with_embeddings(emb.clone());
        }
        let workspace = Arc::new(workspace);
        tools.register_memory_tools(workspace);
    }

    // Register builder tool if enabled
    if config.builder.enabled {
        tools
            .register_builder_tool(
                llm.clone(),
                safety.clone(),
                Some(config.builder.to_builder_config()),
            )
            .await;
        tracing::info!("Builder mode enabled");
    }

    // Load installed WASM tools (save runtime handle for extension manager)
    let wasm_tool_runtime: Option<Arc<WasmToolRuntime>> = if config.wasm.enabled
        && config.wasm.tools_dir.exists()
    {
        match WasmToolRuntime::new(config.wasm.to_runtime_config()) {
            Ok(runtime) => {
                let runtime = Arc::new(runtime);
                let loader = WasmToolLoader::new(Arc::clone(&runtime), Arc::clone(&tools));

                match loader.load_from_dir(&config.wasm.tools_dir).await {
                    Ok(results) => {
                        if !results.loaded.is_empty() {
                            tracing::info!(
                                "Loaded {} WASM tools from {}",
                                results.loaded.len(),
                                config.wasm.tools_dir.display()
                            );
                        }
                        for (path, err) in &results.errors {
                            tracing::warn!("Failed to load WASM tool {}: {}", path.display(), err);
                        }
                    }
                    Err(e) => {
                        tracing::warn!("Failed to scan WASM tools directory: {}", e);
                    }
                }

                Some(runtime)
            }
            Err(e) => {
                tracing::warn!("Failed to initialize WASM runtime: {}", e);
                None
            }
        }
    } else {
        None
    };

    // Create secrets store if master key is configured (needed for MCP auth and WASM channels)
    let secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>> =
        if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) {
            match SecretsCrypto::new(master_key.clone()) {
                Ok(crypto) => Some(Arc::new(PostgresSecretsStore::new(
                    store.pool(),
                    Arc::new(crypto),
                ))),
                Err(e) => {
                    tracing::warn!("Failed to initialize secrets crypto: {}", e);
                    None
                }
            }
        } else {
            None
        };

    // Load configured MCP servers
    let mcp_session_manager = Arc::new(McpSessionManager::new());
    if let Some(ref secrets) = secrets_store {
        match load_mcp_servers().await {
            Ok(servers) => {
                let enabled_count = servers.servers.iter().filter(|s| s.enabled).count();
                if enabled_count > 0 {
                    tracing::info!("Loading {} configured MCP server(s)...", enabled_count);
                }

                for server in servers.enabled_servers() {
                    tracing::debug!(
                        "Checking authentication for MCP server '{}'...",
                        server.name
                    );
                    // Check for stored tokens (from either pre-configured OAuth or DCR)
                    let has_tokens = is_authenticated(server, secrets, "default").await;
                    tracing::debug!("MCP server '{}' has_tokens={}", server.name, has_tokens);

                    let client = if has_tokens || server.requires_auth() {
                        // Use authenticated client if we have tokens or OAuth is configured
                        McpClient::new_authenticated(
                            server.clone(),
                            Arc::clone(&mcp_session_manager),
                            Arc::clone(secrets),
                            "default",
                        )
                    } else {
                        // No tokens and no OAuth - try unauthenticated
                        McpClient::new_with_name(&server.name, &server.url)
                    };

                    tracing::debug!("Fetching tools from MCP server '{}'...", server.name);
                    match client.list_tools().await {
                        Ok(mcp_tools) => {
                            tracing::debug!(
                                "Got {} tools from MCP server '{}'",
                                mcp_tools.len(),
                                server.name
                            );
                            match client.create_tools().await {
                                Ok(tool_impls) => {
                                    for tool in tool_impls {
                                        tools.register(tool).await;
                                    }
                                    tracing::info!(
                                        "Loaded {} tools from MCP server '{}'",
                                        mcp_tools.len(),
                                        server.name
                                    );
                                }
                                Err(e) => {
                                    tracing::warn!(
                                        "Failed to create tools from MCP server '{}': {}",
                                        server.name,
                                        e
                                    );
                                }
                            }
                        }
                        Err(e) => {
                            // Check if it's an auth error
                            let err_str = e.to_string();
                            if err_str.contains("401") || err_str.contains("authentication") {
                                tracing::warn!(
                                    "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}",
                                    server.name,
                                    server.name
                                );
                            } else {
                                tracing::warn!(
                                    "Failed to connect to MCP server '{}': {}",
                                    server.name,
                                    e
                                );
                            }
                        }
                    }
                }
            }
            Err(e) => {
                tracing::debug!("No MCP servers configured ({})", e);
            }
        }
    }

    // Create extension manager for in-chat discovery/install/auth/activate
    let extension_manager = if let Some(ref secrets) = secrets_store {
        let manager = Arc::new(ExtensionManager::new(
            Arc::clone(&mcp_session_manager),
            Arc::clone(secrets),
            Arc::clone(&tools),
            wasm_tool_runtime.clone(),
            config.wasm.tools_dir.clone(),
            config.channels.wasm_channels_dir.clone(),
            config.tunnel.public_url.clone(),
            "default".to_string(),
        ));
        tools.register_extension_tools(Arc::clone(&manager));
        tracing::info!("Extension manager initialized with in-chat discovery tools");
        Some(manager)
    } else {
        tracing::debug!(
            "Extension manager not available (no secrets store). \
             Extension tools won't be registered."
        );
        None
    };

    tracing::info!(
        "Tool registry initialized with {} total tools",
        tools.count()
    );

    // Initialize channel manager
    let mut channels = ChannelManager::new();

    if let Some(repl) = repl_channel {
        channels.add(Box::new(repl));
        if cli.message.is_some() {
            tracing::info!("Single message mode");
        } else {
            tracing::info!("REPL mode enabled");
        }
    }

    // Collect webhook route fragments; a single WebhookServer hosts them all.
    let mut webhook_routes: Vec<axum::Router> = Vec::new();

    // Load WASM channels and register their webhook routes.
    if config.channels.wasm_channels_enabled && config.channels.wasm_channels_dir.exists() {
        match WasmChannelRuntime::new(WasmChannelRuntimeConfig::default()) {
            Ok(runtime) => {
                let runtime = Arc::new(runtime);
                let loader = WasmChannelLoader::new(Arc::clone(&runtime));

                match loader
                    .load_from_dir(&config.channels.wasm_channels_dir)
                    .await
                {
                    Ok(results) => {
                        let wasm_router = Arc::new(WasmChannelRouter::new());
                        let mut has_webhook_channels = false;

                        for loaded in results.loaded {
                            let channel_name = loaded.name().to_string();
                            tracing::info!("Loaded WASM channel: {}", channel_name);

                            let secret_name = loaded.webhook_secret_name();

                            let webhook_secret = if let Some(ref secrets) = secrets_store {
                                secrets
                                    .get_decrypted("default", &secret_name)
                                    .await
                                    .ok()
                                    .map(|s| s.expose().to_string())
                            } else {
                                None
                            };

                            let secret_header =
                                loaded.webhook_secret_header().map(|s| s.to_string());

                            let webhook_path = format!("/webhook/{}", channel_name);
                            let endpoints = vec![RegisteredEndpoint {
                                channel_name: channel_name.clone(),
                                path: webhook_path.clone(),
                                methods: vec!["POST".to_string()],
                                require_secret: webhook_secret.is_some(),
                            }];

                            let channel_arc = Arc::new(loaded.channel);

                            {
                                let mut config_updates = std::collections::HashMap::new();

                                if let Some(ref tunnel_url) = config.tunnel.public_url {
                                    config_updates.insert(
                                        "tunnel_url".to_string(),
                                        serde_json::Value::String(tunnel_url.clone()),
                                    );
                                }

                                if let Some(ref secret) = webhook_secret {
                                    config_updates.insert(
                                        "webhook_secret".to_string(),
                                        serde_json::Value::String(secret.clone()),
                                    );
                                }

                                if !config_updates.is_empty() {
                                    channel_arc.update_config(config_updates).await;
                                    tracing::info!(
                                        channel = %channel_name,
                                        has_tunnel = config.tunnel.public_url.is_some(),
                                        has_webhook_secret = webhook_secret.is_some(),
                                        "Injected runtime config into channel"
                                    );
                                }
                            }

                            tracing::info!(
                                channel = %channel_name,
                                has_webhook_secret = webhook_secret.is_some(),
                                secret_header = ?secret_header,
                                "Registering channel with router"
                            );

                            wasm_router
                                .register(
                                    Arc::clone(&channel_arc),
                                    endpoints,
                                    webhook_secret.clone(),
                                    secret_header,
                                )
                                .await;
                            has_webhook_channels = true;

                            if let Some(ref secrets) = secrets_store {
                                match inject_channel_credentials(
                                    &channel_arc,
                                    secrets.as_ref(),
                                    &channel_name,
                                )
                                .await
                                {
                                    Ok(count) => {
                                        if count > 0 {
                                            tracing::info!(
                                                channel = %channel_name,
                                                credentials_injected = count,
                                                "Channel credentials injected"
                                            );
                                        }
                                    }
                                    Err(e) => {
                                        tracing::error!(
                                            channel = %channel_name,
                                            error = %e,
                                            "Failed to inject channel credentials"
                                        );
                                    }
                                }
                            }

                            channels.add(Box::new(SharedWasmChannel::new(channel_arc)));
                        }

                        if has_webhook_channels && config.tunnel.public_url.is_some() {
                            webhook_routes.push(create_wasm_channel_router(
                                wasm_router,
                                extension_manager.as_ref().map(Arc::clone),
                            ));
                        }

                        for (path, err) in &results.errors {
                            tracing::warn!(
                                "Failed to load WASM channel {}: {}",
                                path.display(),
                                err
                            );
                        }
                    }
                    Err(e) => {
                        tracing::warn!("Failed to scan WASM channels directory: {}", e);
                    }
                }
            }
            Err(e) => {
                tracing::warn!("Failed to initialize WASM channel runtime: {}", e);
            }
        }
    }

    // Add HTTP channel if configured and not CLI-only mode.
    // Extract its routes for the unified server; the channel itself just
    // provides the mpsc stream.
    let mut webhook_server_addr: Option<std::net::SocketAddr> = None;
    if !cli.cli_only {
        if let Some(ref http_config) = config.channels.http {
            let http_channel = HttpChannel::new(http_config.clone());
            webhook_routes.push(http_channel.routes());
            let (host, port) = http_channel.addr();
            webhook_server_addr = Some(
                format!("{}:{}", host, port)
                    .parse()
                    .expect("HttpConfig host:port must be a valid SocketAddr"),
            );
            channels.add(Box::new(http_channel));
            tracing::info!(
                "HTTP channel enabled on {}:{}",
                http_config.host,
                http_config.port
            );
        }
    }

    // Start the unified webhook server if any routes were registered.
    let mut webhook_server = if !webhook_routes.is_empty() {
        let addr =
            webhook_server_addr.unwrap_or_else(|| std::net::SocketAddr::from(([0, 0, 0, 0], 8080)));
        let mut server = WebhookServer::new(WebhookServerConfig { addr });
        for routes in webhook_routes {
            server.add_routes(routes);
        }
        server.start().await?;
        Some(server)
    } else {
        None
    };

    // Create workspace for agent (shared with memory tools)
    let workspace = store.as_ref().map(|s| {
        let mut ws = Workspace::new("default", s.pool());
        if let Some(ref emb) = embeddings {
            ws = ws.with_embeddings(emb.clone());
        }
        Arc::new(ws)
    });

    // Backfill embeddings if we just enabled the provider
    if let (Some(ws), Some(_)) = (&workspace, &embeddings) {
        match ws.backfill_embeddings().await {
            Ok(count) if count > 0 => {
                tracing::info!("Backfilled embeddings for {} chunks", count);
            }
            Ok(_) => {}
            Err(e) => {
                tracing::warn!("Failed to backfill embeddings: {}", e);
            }
        }
    }

    // Create context manager (shared between job tools and agent)
    let context_manager = Arc::new(ContextManager::new(config.agent.max_parallel_jobs));

    // Register job tools
    tools.register_job_tools(Arc::clone(&context_manager));

    // Create and run the agent
    let deps = AgentDeps {
        store,
        llm,
        safety,
        tools,
        workspace,
    };
    let agent = Agent::new(
        config.agent.clone(),
        deps,
        channels,
        Some(config.heartbeat.clone()),
        Some(context_manager),
    );

    tracing::info!("Agent initialized, starting main loop...");

    // Run the agent (blocks until shutdown)
    agent.run().await?;

    // Shut down the webhook server if one was started
    if let Some(ref mut server) = webhook_server {
        server.shutdown().await;
    }

    tracing::info!("Agent shutdown complete");
    Ok(())
}

/// Check if setup is needed and return the reason.
///
/// Returns `Some(reason)` if setup should be triggered, `None` otherwise.
fn check_setup_needed() -> Option<&'static str> {
    let settings = Settings::load();

    // Database not configured (and not in env)
    if settings.database_url.is_none() && std::env::var("DATABASE_URL").is_err() {
        return Some("Database not configured");
    }

    // Secrets not configured (and not in env)
    if settings.secrets_master_key_source == ironclaw::settings::KeySource::None
        && std::env::var("SECRETS_MASTER_KEY").is_err()
        && !ironclaw::secrets::keychain::has_master_key()
    {
        // Only require secrets setup if user hasn't explicitly disabled it
        // For now, we don't require it for first run
    }

    // First run (setup never completed and no session)
    let session_path = ironclaw::llm::session::default_session_path();
    if !settings.setup_completed && !session_path.exists() {
        return Some("First run");
    }

    None
}

/// Inject credentials for a channel based on naming convention.
///
/// Looks for secrets matching the pattern `{channel_name}_*` and injects them
/// as credential placeholders (e.g., `telegram_bot_token` -> `{TELEGRAM_BOT_TOKEN}`).
///
/// Returns the number of credentials injected.
async fn inject_channel_credentials(
    channel: &Arc<ironclaw::channels::wasm::WasmChannel>,
    secrets: &dyn SecretsStore,
    channel_name: &str,
) -> anyhow::Result<usize> {
    // List all secrets for this user and filter by channel prefix
    let all_secrets = secrets
        .list("default")
        .await
        .map_err(|e| anyhow::anyhow!("Failed to list secrets: {}", e))?;

    let prefix = format!("{}_", channel_name);
    let mut count = 0;

    for secret_meta in all_secrets {
        // Only process secrets matching the channel prefix
        if !secret_meta.name.starts_with(&prefix) {
            continue;
        }

        // Get the decrypted value
        let decrypted = match secrets.get_decrypted("default", &secret_meta.name).await {
            Ok(d) => d,
            Err(e) => {
                tracing::warn!(
                    secret = %secret_meta.name,
                    error = %e,
                    "Failed to decrypt secret for channel credential injection"
                );
                continue;
            }
        };

        // Convert secret name to placeholder format (SCREAMING_SNAKE_CASE)
        let placeholder = secret_meta.name.to_uppercase();

        tracing::debug!(
            channel = %channel_name,
            secret = %secret_meta.name,
            placeholder = %placeholder,
            "Injecting credential"
        );

        channel
            .set_credential(&placeholder, decrypted.expose().to_string())
            .await;
        count += 1;
    }

    Ok(count)
}