PicoBot/src/gateway/mod.rs

pub mod auth;
pub mod http;
mod router;
pub mod uploads;
pub mod ws;

use axum::{Router, middleware, routing};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;

use crate::bus::{MessageBus, OutboundDispatcher};
use crate::channels::{ChannelManager, CliChatChannel};
use crate::config::{Config, ensure_workspace_dir, expand_path};
use crate::delivery::{ConversationWriteLocks, DeliveryCoordinator, TurnDeliveryService};
use crate::logging;
use crate::mcp;
use crate::memory::MemoryManager;
use crate::scheduler::Scheduler;
use crate::session::{SessionManager, SessionManagerServices};
use crate::task_supervisor::TaskSupervisor;

pub struct GatewayState {
    pub config: Config,
    pub config_path: std::path::PathBuf,
    pub workspace_dir: std::path::PathBuf,
    pub session_manager: Arc<SessionManager>,
    pub channel_manager: ChannelManager,
    pub storage: Arc<crate::storage::Storage>,
    pub task_supervisor: TaskSupervisor,
    pub delivery_coordinator: DeliveryCoordinator,
    pub connection_shutdown: tokio_util::sync::CancellationToken,
    pub auth: auth::AuthManager,
    pub uploads: uploads::UploadRegistry,
}

impl GatewayState {
    pub async fn new() -> Result<Self, Box<dyn std::error::Error>> {
        let config_path = crate::config::resolve_default_config_path();
        let config = Config::load_from(&config_path)?;
        Self::from_config(config, config_path).await
    }

    async fn from_config(
        config: Config,
        config_path: std::path::PathBuf,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let task_supervisor = TaskSupervisor::new();
        let delivery_coordinator = DeliveryCoordinator::new(ConversationWriteLocks::default());
        let connection_shutdown = tokio_util::sync::CancellationToken::new();
        let auth = auth::AuthManager::load(
            config.gateway.require_pairing,
            crate::config::get_user_config_dir().join("web_auth.json"),
        )
        .await?;
        let uploads = uploads::UploadRegistry::new(config.gateway.file_transfer.clone());

        // Initialize workspace directory: expand path and ensure it exists
        let workspace_path = expand_path(&config.workspace_dir);
        let workspace_path = ensure_workspace_dir(&workspace_path)?;

        // Switch current working directory to workspace
        std::env::set_current_dir(&workspace_path).map_err(|e| {
            format!(
                "Failed to switch to workspace directory {}: {}",
                workspace_path.display(),
                e
            )
        })?;

        tracing::info!("Using workspace directory: {}", workspace_path.display());

        // Release default AGENTS.md and USER.md to ~/.picobot/ if not exist
        ensure_default_config_files();

        // Get provider config for SessionManager
        let mut provider_config = config.get_provider_config("default")?;
        // Override workspace_dir with the ensured path
        provider_config.workspace_dir = workspace_path.clone();

        // Initialize Storage
        let db_path = if let Some(ref path) = config.gateway.session_db_path {
            std::path::PathBuf::from(path)
        } else {
            workspace_path.join("picobot.db")
        };
        let storage = Arc::new(
            crate::storage::Storage::new(&db_path)
                .await
                .map_err(|e| format!("failed to initialize session storage: {}", e))?,
        );
        tracing::info!("Session storage: {}", db_path.display());

        // Resolve consolidation provider/model with fallback to main agent config
        let consolidation_provider = config
            .memory
            .resolve_consolidation_provider(&provider_config.name);
        let consolidation_model = config
            .memory
            .resolve_consolidation_model(&provider_config.model_id);
        let memory_manager = Arc::new(MemoryManager::new(
            storage.clone(),
            consolidation_provider,
            consolidation_model,
        ));
        tracing::info!(
            consolidation_provider = %memory_manager.consolidation_provider,
            consolidation_model = %memory_manager.consolidation_model,
            "Memory system initialized"
        );

        // Create MessageBus first (shared by SessionManager and ChannelManager)
        let bus = MessageBus::new(100);

        // Channels are resolved by TurnDeliveryService, while Session workers
        // depend only on that protocol-neutral delivery facade.
        let cli_chat_channel = Arc::new(CliChatChannel::with_upload_registry(uploads.clone()));
        let channel_manager = ChannelManager::with_bus(cli_chat_channel, bus.clone());
        channel_manager
            .init(&config, workspace_path.clone())
            .await
            .map_err(|e| format!("Failed to init channels: {}", e))?;
        let turn_delivery = TurnDeliveryService::new(
            delivery_coordinator.clone(),
            channel_manager.clone(),
            task_supervisor.clone(),
        );

        let browser_config = if config.browser.enabled {
            Some(config.browser.clone())
        } else {
            None
        };

        // Create SessionManager with bus injection
        let session_manager = SessionManager::new(
            provider_config.clone(),
            storage.clone(),
            SessionManagerServices::new(
                bus.clone(),
                memory_manager,
                task_supervisor.clone(),
                turn_delivery,
            ),
            browser_config,
            config.gateway.max_concurrent_background_tasks,
        )?;
        let session_manager = Arc::new(session_manager);

        // Register send_message tool with available channel names
        let available_channels = channel_manager.list_channel_names().await;
        let valid_channels = available_channels.clone();
        session_manager.register_outbound_tool(available_channels);

        // Register chat_manager tool
        session_manager
            .tools()
            .register(crate::tools::ChatManagerTool::new(
                storage.clone(),
                valid_channels.clone(),
            ));

        // Initialize MCP servers — connect and register discovered tools
        if !config.mcp.servers.is_empty() {
            let mcp_tools = mcp::connect_all(&config.mcp).await;
            for tool_info in mcp_tools {
                let wrapper = mcp::McpToolWrapper::new(
                    &tool_info.server_name,
                    tool_info.tool_name,
                    tool_info.description,
                    tool_info.schema,
                    tool_info.connection,
                );
                session_manager.tools().register(wrapper);
            }
        }

        // Initialize scheduler if enabled in config
        let scheduler_config = config.gateway.scheduler.clone().unwrap_or_default();
        if scheduler_config.enabled {
            session_manager
                .tools()
                .register(crate::tools::RoutineMaintenanceTool::new(
                    storage.clone(),
                    config.memory.timeline_retention_days,
                ));
            storage
                .ensure_default_maintenance_job()
                .await
                .map_err(|e| format!("failed to seed default maintenance job: {e}"))?;
            // Register cron tools
            session_manager
                .tools()
                .register(crate::tools::cron::CronAddTool::new(
                    storage.clone(),
                    valid_channels,
                ));
            session_manager
                .tools()
                .register(crate::tools::cron::CronListTool::new(storage.clone()));
            session_manager
                .tools()
                .register(crate::tools::cron::CronRemoveTool::new(storage.clone()));
            session_manager
                .tools()
                .register(crate::tools::cron::CronEnableTool::new(storage.clone()));
            session_manager
                .tools()
                .register(crate::tools::cron::CronDisableTool::new(storage.clone()));
            session_manager
                .tools()
                .register(crate::tools::cron::CronUpdateTool::new(storage.clone()));
            tracing::info!("Cron tools registered");
        }

        Ok(Self {
            config,
            config_path,
            workspace_dir: workspace_path,
            session_manager: session_manager.clone(),
            channel_manager,
            storage,
            task_supervisor,
            delivery_coordinator,
            connection_shutdown,
            auth,
            uploads,
        })
    }

    /// Get a reference to the MessageBus
    pub fn bus(&self) -> Arc<crate::bus::MessageBus> {
        self.channel_manager.bus()
    }

    /// Get CLI chat channel for WebSocket handling
    pub fn cli_chat_channel(&self) -> Arc<CliChatChannel> {
        self.channel_manager.cli_chat_channel()
    }

    /// Start the message processing loops
    pub async fn start_message_processing(&self) {
        let bus = self.bus();
        let bus_for_outbound = bus.clone();
        let session_manager = self.session_manager.clone();

        if self.uploads.enabled() {
            let uploads = self.uploads.clone();
            self.task_supervisor
                .spawn("pending-upload-cleanup", async move {
                    let mut interval = tokio::time::interval(std::time::Duration::from_secs(300));
                    interval.tick().await;
                    loop {
                        interval.tick().await;
                        let removed = uploads.cleanup_expired().await;
                        if removed > 0 {
                            tracing::debug!(removed, "Expired pending uploads removed");
                        }
                    }
                });
        }

        // Relay structured plan changes to WebSocket clients. This remains
        // separate from chat messages, so task UI updates never pollute history.
        let mut plan_events = self.session_manager.work_manager().subscribe();
        let cli_chat = self.cli_chat_channel();
        self.task_supervisor.spawn("task-plan-events", async move {
            loop {
                match plan_events.recv().await {
                    Ok(event) => cli_chat.publish_plan_changed(event).await,
                    Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
                        tracing::warn!(skipped, "Task plan event relay lagged");
                    }
                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
                }
            }
        });

        router::spawn_message_routers(bus.clone(), session_manager, self.task_supervisor.clone());

        // Spawn outbound dispatcher
        let dispatcher = OutboundDispatcher::new(
            bus_for_outbound,
            self.channel_manager.clone(),
            self.task_supervisor.clone(),
            self.delivery_coordinator.write_locks(),
        );

        self.task_supervisor
            .spawn("outbound-dispatcher", async move {
                tracing::info!("Outbound dispatcher started");
                dispatcher.run().await;
            });

        // Spawn scheduler background task if enabled
        let scheduler_config = self.config.gateway.scheduler.clone().unwrap_or_default();
        if scheduler_config.enabled {
            let sched = Arc::new(Scheduler::new(
                self.storage.clone(),
                self.session_manager.clone(),
                scheduler_config,
            ));
            self.task_supervisor.spawn("scheduler", async move {
                sched.run().await;
            });
            tracing::info!("Scheduler background task spawned");
        }
    }
}

pub async fn run(
    host: Option<String>,
    port: Option<u16>,
) -> Result<(), Box<dyn std::error::Error>> {
    let config_path = crate::config::resolve_default_config_path();
    let config = Config::load_from(&config_path)?;

    // Initialize logging
    logging::init_logging();
    tracing::info!(config_path = %config_path.display(), "Starting PicoBot Gateway");

    let state = Arc::new(GatewayState::from_config(config, config_path).await?);

    // Start all channels (init already done while constructing GatewayState)
    state.channel_manager.start_all().await?;

    // Start message processing (inbound processor + control processor + outbound dispatcher)
    state.start_message_processing().await;

    // CLI args override config file values
    let bind_host = host.unwrap_or_else(|| state.config.gateway.host.clone());
    let bind_port = port.unwrap_or(state.config.gateway.port);

    let protected = Router::new()
        .route("/api/health", routing::get(http::health))
        .route(
            "/api/config",
            routing::get(http::get_config).put(http::put_config),
        )
        .route(
            "/api/profiles/{name}",
            routing::get(http::get_profile).put(http::put_profile),
        )
        .route("/api/logs", routing::get(http::get_logs))
        .route("/api/tasks", routing::get(http::get_tasks))
        .route("/api/jobs", routing::get(http::get_jobs))
        .route("/api/jobs/{id}/runs", routing::get(http::get_job_runs))
        .route("/api/memories", routing::get(http::get_memories))
        .route(
            "/api/chat/{client_id}/uploads",
            routing::post(http::upload_file).layer(axum::extract::DefaultBodyLimit::disable()),
        )
        .route(
            "/api/chat/{client_id}/sessions/{session_id}/messages/{message_id}/attachments/{index}",
            routing::get(http::download_attachment),
        )
        .route("/ws", routing::get(ws::ws_handler))
        .route_layer(middleware::from_fn_with_state(
            state.auth.clone(),
            auth::require_auth,
        ));

    let app = Router::new()
        .route("/", routing::get(http::webui_index))
        .route("/app.js", routing::get(http::webui_script))
        .route("/styles.css", routing::get(http::webui_styles))
        .route("/health", routing::get(http::health))
        .route("/api/auth/status", routing::get(auth::status))
        .route("/api/auth/pair", routing::post(auth::pair))
        .route("/api/auth/code", routing::post(auth::issue_code))
        .merge(protected)
        .with_state(state.clone());

    let addr = format!("{}:{}", bind_host, bind_port);
    let listener = TcpListener::bind(&addr).await?;
    tracing::info!(address = %addr, "Gateway listening");

    let connection_shutdown = state.connection_shutdown.clone();
    let serve_result = axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(async move {
        wait_for_shutdown_signal().await;
        tracing::info!("Shutdown signal received");
        connection_shutdown.cancel();
    })
    .await;

    // Stop external intake before waiting for internal work to finish.
    if let Err(error) = state.channel_manager.stop_all().await {
        tracing::error!(error = %error, "Failed to stop channels cleanly");
    }
    state.task_supervisor.cancel();
    state
        .task_supervisor
        .shutdown(std::time::Duration::from_secs(10))
        .await;
    serve_result?;

    Ok(())
}

async fn wait_for_shutdown_signal() {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{SignalKind, signal};

        match signal(SignalKind::terminate()) {
            Ok(mut terminate) => {
                tokio::select! {
                    result = tokio::signal::ctrl_c() => {
                        if let Err(error) = result {
                            tracing::error!(error = %error, "Failed to listen for Ctrl-C");
                        }
                    }
                    _ = terminate.recv() => {}
                }
            }
            Err(error) => {
                tracing::error!(error = %error, "Failed to listen for SIGTERM");
                if let Err(error) = tokio::signal::ctrl_c().await {
                    tracing::error!(error = %error, "Failed to listen for Ctrl-C");
                }
            }
        }
    }

    #[cfg(not(unix))]
    if let Err(error) = tokio::signal::ctrl_c().await {
        tracing::error!(error = %error, "Failed to listen for Ctrl-C");
    }
}

/// Release default AGENTS.md and USER.md templates to ~/.picobot/ if not already present.
fn ensure_default_config_files() {
    let picobot_dir = dirs::home_dir().unwrap_or_default().join(".picobot");
    if let Err(e) = std::fs::create_dir_all(&picobot_dir) {
        tracing::warn!(dir = %picobot_dir.display(), error = %e, "Failed to create ~/.picobot directory");
        return;
    }

    let agents_path = picobot_dir.join("AGENTS.md");
    if !agents_path.exists() {
        let content = include_str!("../../resources/templates/AGENTS.md");
        if let Err(e) = std::fs::write(&agents_path, content) {
            tracing::warn!(path = %agents_path.display(), error = %e, "Failed to write AGENTS.md template");
        } else {
            tracing::info!(path = %agents_path.display(), "Released default AGENTS.md template");
        }
    }

    let user_path = picobot_dir.join("USER.md");
    if !user_path.exists() {
        let content = include_str!("../../resources/templates/USER.md");
        if let Err(e) = std::fs::write(&user_path, content) {
            tracing::warn!(path = %user_path.display(), error = %e, "Failed to write USER.md template");
        } else {
            tracing::info!(path = %user_path.display(), "Released default USER.md template");
        }
    }

    let config_example_path = picobot_dir.join("config.example.json");
    if !config_example_path.exists() {
        let content = include_str!("../../resources/templates/config.example.json");
        if let Err(e) = std::fs::write(&config_example_path, content) {
            tracing::warn!(path = %config_example_path.display(), error = %e, "Failed to write config.example.json template");
        } else {
            tracing::info!(path = %config_example_path.display(), "Released config.example.json template");
        }
    }
}