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

//! OWLEN CLI - Chat TUI client

mod cloud;

use anyhow::{Result, anyhow};
use async_trait::async_trait;
use clap::{Parser, Subcommand};
use cloud::{CloudCommand, load_runtime_credentials, set_env_var};
use owlen_core::config as core_config;
use owlen_core::{
    ChatStream, Error, Provider,
    config::{Config, McpMode},
    mcp::remote_client::RemoteMcpClient,
    mode::Mode,
    providers::OllamaProvider,
    session::SessionController,
    storage::StorageManager,
    types::{ChatRequest, ChatResponse, Message, ModelInfo},
};
use owlen_tui::tui_controller::{TuiController, TuiRequest};
use owlen_tui::{AppState, ChatApp, Event, EventHandler, SessionEvent, config, ui};
use std::any::Any;
use std::borrow::Cow;
use std::io;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

use crossterm::{
    event::{DisableBracketedPaste, DisableMouseCapture, EnableBracketedPaste, EnableMouseCapture},
    execute,
    terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode},
};
use futures::stream;
use ratatui::{Terminal, prelude::CrosstermBackend};

/// Owlen - Terminal UI for LLM chat
#[derive(Parser, Debug)]
#[command(name = "owlen")]
#[command(about = "Terminal UI for LLM chat via MCP", long_about = None)]
struct Args {
    /// Start in code mode (enables all tools)
    #[arg(long, short = 'c')]
    code: bool,
    #[command(subcommand)]
    command: Option<OwlenCommand>,
}

#[derive(Debug, Subcommand)]
enum OwlenCommand {
    /// Inspect or upgrade configuration files
    #[command(subcommand)]
    Config(ConfigCommand),
    /// Manage Ollama Cloud credentials
    #[command(subcommand)]
    Cloud(CloudCommand),
    /// Show manual steps for updating Owlen to the latest revision
    Upgrade,
}

#[derive(Debug, Subcommand)]
enum ConfigCommand {
    /// Automatically upgrade legacy configuration values and ensure validity
    Doctor,
    /// Print the resolved configuration file path
    Path,
}

fn build_provider(cfg: &Config) -> anyhow::Result<Arc<dyn Provider>> {
    match cfg.mcp.mode {
        McpMode::RemotePreferred => {
            let remote_result = if let Some(mcp_server) = cfg.mcp_servers.first() {
                RemoteMcpClient::new_with_config(mcp_server)
            } else {
                RemoteMcpClient::new()
            };

            match remote_result {
                Ok(client) => {
                    let provider: Arc<dyn Provider> = Arc::new(client);
                    Ok(provider)
                }
                Err(err) if cfg.mcp.allow_fallback => {
                    log::warn!(
                        "Remote MCP client unavailable ({}); falling back to local provider.",
                        err
                    );
                    build_local_provider(cfg)
                }
                Err(err) => Err(anyhow::Error::from(err)),
            }
        }
        McpMode::RemoteOnly => {
            let mcp_server = cfg.mcp_servers.first().ok_or_else(|| {
                anyhow::anyhow!(
                    "[[mcp_servers]] must be configured when [mcp].mode = \"remote_only\""
                )
            })?;
            let client = RemoteMcpClient::new_with_config(mcp_server)?;
            let provider: Arc<dyn Provider> = Arc::new(client);
            Ok(provider)
        }
        McpMode::LocalOnly | McpMode::Legacy => build_local_provider(cfg),
        McpMode::Disabled => Err(anyhow::anyhow!(
            "MCP mode 'disabled' is not supported by the owlen TUI"
        )),
    }
}

fn build_local_provider(cfg: &Config) -> anyhow::Result<Arc<dyn Provider>> {
    let provider_name = cfg.general.default_provider.clone();
    let provider_cfg = cfg.provider(&provider_name).ok_or_else(|| {
        anyhow::anyhow!(format!(
            "No provider configuration found for '{provider_name}' in [providers]"
        ))
    })?;

    match provider_cfg.provider_type.as_str() {
        "ollama" | "ollama-cloud" => {
            let provider = OllamaProvider::from_config(provider_cfg, Some(&cfg.general))?;
            Ok(Arc::new(provider) as Arc<dyn Provider>)
        }
        other => Err(anyhow::anyhow!(format!(
            "Provider type '{other}' is not supported in legacy/local MCP mode"
        ))),
    }
}

async fn run_command(command: OwlenCommand) -> Result<()> {
    match command {
        OwlenCommand::Config(config_cmd) => run_config_command(config_cmd),
        OwlenCommand::Cloud(cloud_cmd) => cloud::run_cloud_command(cloud_cmd).await,
        OwlenCommand::Upgrade => {
            println!(
                "To update Owlen from source:\n  git pull\n  cargo install --path crates/owlen-cli --force"
            );
            println!(
                "If you installed from the AUR, use your package manager (e.g., yay -S owlen-git)."
            );
            Ok(())
        }
    }
}

fn run_config_command(command: ConfigCommand) -> Result<()> {
    match command {
        ConfigCommand::Doctor => run_config_doctor(),
        ConfigCommand::Path => {
            let path = core_config::default_config_path();
            println!("{}", path.display());
            Ok(())
        }
    }
}

fn run_config_doctor() -> Result<()> {
    let config_path = core_config::default_config_path();
    let existed = config_path.exists();
    let mut config = config::try_load_config().unwrap_or_default();
    let mut changes = Vec::new();

    if !existed {
        changes.push("created configuration file from defaults".to_string());
    }

    if !config
        .providers
        .contains_key(&config.general.default_provider)
    {
        config.general.default_provider = "ollama".to_string();
        changes.push("default provider missing; reset to 'ollama'".to_string());
    }

    if let Some(mut legacy) = config.providers.remove("ollama-cloud") {
        legacy.provider_type = "ollama".to_string();
        use std::collections::hash_map::Entry;
        match config.providers.entry("ollama".to_string()) {
            Entry::Occupied(mut existing) => {
                let entry = existing.get_mut();
                if entry.api_key.is_none() {
                    entry.api_key = legacy.api_key.take();
                }
                if entry.base_url.is_none() && legacy.base_url.is_some() {
                    entry.base_url = legacy.base_url.take();
                }
                entry.extra.extend(legacy.extra);
            }
            Entry::Vacant(slot) => {
                slot.insert(legacy);
            }
        }
        changes.push(
            "migrated legacy 'ollama-cloud' provider into unified 'ollama' entry".to_string(),
        );
    }

    if !config.providers.contains_key("ollama") {
        core_config::ensure_provider_config(&mut config, "ollama");
        changes.push("added default ollama provider configuration".to_string());
    }

    match config.mcp.mode {
        McpMode::Legacy => {
            config.mcp.mode = McpMode::LocalOnly;
            config.mcp.warn_on_legacy = true;
            changes.push("converted [mcp].mode = 'legacy' to 'local_only'".to_string());
        }
        McpMode::RemoteOnly if config.mcp_servers.is_empty() => {
            config.mcp.mode = McpMode::RemotePreferred;
            config.mcp.allow_fallback = true;
            changes.push(
                "downgraded remote-only configuration to remote_preferred because no servers are defined"
                    .to_string(),
            );
        }
        McpMode::RemotePreferred if !config.mcp.allow_fallback && config.mcp_servers.is_empty() => {
            config.mcp.allow_fallback = true;
            changes.push(
                "enabled [mcp].allow_fallback because no remote servers are configured".to_string(),
            );
        }
        _ => {}
    }

    config.validate()?;
    config::save_config(&config)?;

    if changes.is_empty() {
        println!(
            "Configuration already up to date: {}",
            config_path.display()
        );
    } else {
        println!("Updated {}:", config_path.display());
        for change in changes {
            println!("  - {change}");
        }
    }

    Ok(())
}

const BASIC_THEME_NAME: &str = "ansi_basic";

#[derive(Debug, Clone)]
enum TerminalColorSupport {
    Full,
    Limited { term: String },
}

fn detect_terminal_color_support() -> TerminalColorSupport {
    let term = std::env::var("TERM").unwrap_or_else(|_| "unknown".to_string());
    let colorterm = std::env::var("COLORTERM").unwrap_or_default();
    let term_lower = term.to_lowercase();
    let color_lower = colorterm.to_lowercase();

    let supports_extended = term_lower.contains("256color")
        || color_lower.contains("truecolor")
        || color_lower.contains("24bit")
        || color_lower.contains("fullcolor");

    if supports_extended {
        TerminalColorSupport::Full
    } else {
        TerminalColorSupport::Limited { term }
    }
}

fn apply_terminal_theme(cfg: &mut Config, support: &TerminalColorSupport) -> Option<String> {
    match support {
        TerminalColorSupport::Full => None,
        TerminalColorSupport::Limited { .. } => {
            if cfg.ui.theme != BASIC_THEME_NAME {
                let previous = std::mem::replace(&mut cfg.ui.theme, BASIC_THEME_NAME.to_string());
                Some(previous)
            } else {
                None
            }
        }
    }
}

struct OfflineProvider {
    reason: String,
    placeholder_model: String,
}

impl OfflineProvider {
    fn new(reason: String, placeholder_model: String) -> Self {
        Self {
            reason,
            placeholder_model,
        }
    }

    fn friendly_response(&self, requested_model: &str) -> ChatResponse {
        let mut message = String::new();
        message.push_str("⚠️  Owlen is running in offline mode.\n\n");
        message.push_str(&self.reason);
        if !requested_model.is_empty() && requested_model != self.placeholder_model {
            message.push_str(&format!(
                "\n\nYou requested model '{}', but no providers are reachable.",
                requested_model
            ));
        }
        message.push_str(
            "\n\nStart your preferred provider (e.g. `ollama serve`) or switch providers with `:provider` once connectivity is restored.",
        );

        ChatResponse {
            message: Message::assistant(message),
            usage: None,
            is_streaming: false,
            is_final: true,
        }
    }
}

#[async_trait]
impl Provider for OfflineProvider {
    fn name(&self) -> &str {
        "offline"
    }

    async fn list_models(&self) -> Result<Vec<ModelInfo>, Error> {
        Ok(vec![ModelInfo {
            id: self.placeholder_model.clone(),
            provider: "offline".to_string(),
            name: format!("Offline (fallback: {})", self.placeholder_model),
            description: Some("Placeholder model used while no providers are reachable".into()),
            context_window: None,
            capabilities: vec![],
            supports_tools: false,
        }])
    }

    async fn send_prompt(&self, request: ChatRequest) -> Result<ChatResponse, Error> {
        Ok(self.friendly_response(&request.model))
    }

    async fn stream_prompt(&self, request: ChatRequest) -> Result<ChatStream, Error> {
        let response = self.friendly_response(&request.model);
        Ok(Box::pin(stream::iter(vec![Ok(response)])))
    }

    async fn health_check(&self) -> Result<(), Error> {
        Err(Error::Provider(anyhow!(
            "offline provider cannot reach any backing models"
        )))
    }

    fn as_any(&self) -> &(dyn Any + Send + Sync) {
        self
    }
}

#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<()> {
    // Parse command-line arguments
    let Args { code, command } = Args::parse();
    if let Some(command) = command {
        return run_command(command).await;
    }
    let initial_mode = if code { Mode::Code } else { Mode::Chat };

    // Set auto-consent for TUI mode to prevent blocking stdin reads
    set_env_var("OWLEN_AUTO_CONSENT", "1");

    let color_support = detect_terminal_color_support();
    // Load configuration (or fall back to defaults) for the session controller.
    let mut cfg = config::try_load_config().unwrap_or_default();
    if let Some(previous_theme) = apply_terminal_theme(&mut cfg, &color_support) {
        let term_label = match &color_support {
            TerminalColorSupport::Limited { term } => Cow::from(term.as_str()),
            TerminalColorSupport::Full => Cow::from("current terminal"),
        };
        eprintln!(
            "Terminal '{}' lacks full 256-color support. Using '{}' theme instead of '{}'.",
            term_label, BASIC_THEME_NAME, previous_theme
        );
    } else if let TerminalColorSupport::Limited { term } = &color_support {
        eprintln!(
            "Warning: terminal '{}' may not fully support 256-color themes.",
            term
        );
    }
    cfg.validate()?;
    let storage = Arc::new(StorageManager::new().await?);
    load_runtime_credentials(&mut cfg, storage.clone()).await?;

    let (tui_tx, _tui_rx) = mpsc::unbounded_channel::<TuiRequest>();
    let tui_controller = Arc::new(TuiController::new(tui_tx));

    // Create provider according to MCP configuration (supports legacy/local fallback)
    let provider = build_provider(&cfg)?;
    let mut offline_notice: Option<String> = None;
    let provider = match provider.health_check().await {
        Ok(_) => provider,
        Err(err) => {
            let hint = if matches!(cfg.mcp.mode, McpMode::RemotePreferred | McpMode::RemoteOnly)
                && !cfg.mcp_servers.is_empty()
            {
                "Ensure the configured MCP server is running and reachable."
            } else {
                "Ensure Ollama is running (`ollama serve`) and reachable at the configured base_url."
            };
            let notice =
                format!("Provider health check failed: {err}. {hint} Continuing in offline mode.");
            eprintln!("{notice}");
            offline_notice = Some(notice.clone());
            let fallback_model = cfg
                .general
                .default_model
                .clone()
                .unwrap_or_else(|| "offline".to_string());
            Arc::new(OfflineProvider::new(notice, fallback_model)) as Arc<dyn Provider>
        }
    };

    let controller =
        SessionController::new(provider, cfg, storage.clone(), tui_controller, false).await?;
    let (mut app, mut session_rx) = ChatApp::new(controller).await?;
    app.initialize_models().await?;
    if let Some(notice) = offline_notice {
        app.set_status_message(&notice);
        app.set_system_status(notice);
    }

    // Set the initial mode
    app.set_mode(initial_mode).await;

    // Event infrastructure
    let cancellation_token = CancellationToken::new();
    let (event_tx, event_rx) = mpsc::unbounded_channel();
    let event_handler = EventHandler::new(event_tx, cancellation_token.clone());
    let event_handle = tokio::spawn(async move { event_handler.run().await });

    // Terminal setup
    enable_raw_mode()?;
    let mut stdout = io::stdout();
    execute!(
        stdout,
        EnterAlternateScreen,
        EnableMouseCapture,
        EnableBracketedPaste
    )?;
    let backend = CrosstermBackend::new(stdout);
    let mut terminal = Terminal::new(backend)?;

    let result = run_app(&mut terminal, &mut app, event_rx, &mut session_rx).await;

    // Shutdown
    cancellation_token.cancel();
    event_handle.await?;

    // Persist configuration updates (e.g., selected model)
    config::save_config(&app.config())?;

    disable_raw_mode()?;
    execute!(
        terminal.backend_mut(),
        LeaveAlternateScreen,
        DisableMouseCapture,
        DisableBracketedPaste
    )?;
    terminal.show_cursor()?;

    if let Err(err) = result {
        println!("{err:?}");
    }

    Ok(())
}

async fn run_app(
    terminal: &mut Terminal<CrosstermBackend<io::Stdout>>,
    app: &mut ChatApp,
    mut event_rx: mpsc::UnboundedReceiver<Event>,
    session_rx: &mut mpsc::UnboundedReceiver<SessionEvent>,
) -> Result<()> {
    let stream_draw_interval = tokio::time::Duration::from_millis(50);
    let idle_tick = tokio::time::Duration::from_millis(100);
    let mut last_draw = tokio::time::Instant::now() - stream_draw_interval;

    loop {
        // Advance loading animation frame
        app.advance_loading_animation();

        let streaming_active = app.streaming_count() > 0;
        let draw_due = if streaming_active {
            last_draw.elapsed() >= stream_draw_interval
        } else {
            true
        };

        if draw_due {
            terminal.draw(|f| ui::render_chat(f, app))?;
            last_draw = tokio::time::Instant::now();
        }

        // Process any pending LLM requests AFTER UI has been drawn
        if let Err(e) = app.process_pending_llm_request().await {
            eprintln!("Error processing LLM request: {}", e);
        }

        // Process any pending tool executions AFTER UI has been drawn
        if let Err(e) = app.process_pending_tool_execution().await {
            eprintln!("Error processing tool execution: {}", e);
        }

        let sleep_duration = if streaming_active {
            stream_draw_interval
                .checked_sub(last_draw.elapsed())
                .unwrap_or_else(|| tokio::time::Duration::from_millis(0))
        } else {
            idle_tick
        };

        tokio::select! {
            Some(event) = event_rx.recv() => {
                if let AppState::Quit = app.handle_event(event).await? {
                    return Ok(());
                }
            }
            Some(session_event) = session_rx.recv() => {
                app.handle_session_event(session_event)?;
            }
            _ = tokio::time::sleep(sleep_duration) => {}
        }
    }
}