owlen/crates/owlen-cli/src/commands/cloud.rs

use std::ffi::OsStr;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::{Context, Result, anyhow, bail};
use clap::Subcommand;
use owlen_core::LlmProvider;
use owlen_core::ProviderConfig;
use owlen_core::config::{
    self as core_config, Config, OLLAMA_CLOUD_API_KEY_ENV, OLLAMA_CLOUD_BASE_URL,
    OLLAMA_CLOUD_ENDPOINT_KEY, OLLAMA_MODE_KEY,
};
use owlen_core::credentials::{ApiCredentials, CredentialManager, OLLAMA_CLOUD_CREDENTIAL_ID};
use owlen_core::encryption;
use owlen_core::providers::OllamaProvider;
use owlen_core::storage::StorageManager;
use serde_json::Value;

const DEFAULT_CLOUD_ENDPOINT: &str = OLLAMA_CLOUD_BASE_URL;
const CLOUD_ENDPOINT_KEY: &str = OLLAMA_CLOUD_ENDPOINT_KEY;
const CLOUD_PROVIDER_KEY: &str = "ollama_cloud";

#[derive(Debug, Subcommand)]
pub enum CloudCommand {
    /// Configure Ollama Cloud credentials
    Setup {
        /// API key passed directly on the command line (prompted when omitted)
        #[arg(long)]
        api_key: Option<String>,
        /// Override the cloud endpoint (default: https://ollama.com)
        #[arg(long)]
        endpoint: Option<String>,
        /// Provider name to configure (default: ollama_cloud)
        #[arg(long, default_value = "ollama_cloud")]
        provider: String,
        /// Overwrite the provider base URL with the cloud endpoint
        #[arg(long)]
        force_cloud_base_url: bool,
    },
    /// Check connectivity to Ollama Cloud
    Status {
        /// Provider name to check (default: ollama_cloud)
        #[arg(long, default_value = "ollama_cloud")]
        provider: String,
    },
    /// List available cloud-hosted models
    Models {
        /// Provider name to query (default: ollama_cloud)
        #[arg(long, default_value = "ollama_cloud")]
        provider: String,
    },
    /// Remove stored Ollama Cloud credentials
    Logout {
        /// Provider name to clear (default: ollama_cloud)
        #[arg(long, default_value = "ollama_cloud")]
        provider: String,
    },
}

pub async fn run_cloud_command(command: CloudCommand) -> Result<()> {
    match command {
        CloudCommand::Setup {
            api_key,
            endpoint,
            provider,
            force_cloud_base_url,
        } => setup(provider, api_key, endpoint, force_cloud_base_url).await,
        CloudCommand::Status { provider } => status(provider).await,
        CloudCommand::Models { provider } => models(provider).await,
        CloudCommand::Logout { provider } => logout(provider).await,
    }
}

async fn setup(
    provider: String,
    api_key: Option<String>,
    endpoint: Option<String>,
    force_cloud_base_url: bool,
) -> Result<()> {
    let provider = canonical_provider_name(&provider);
    let mut config = crate::config::try_load_config().unwrap_or_default();
    let endpoint =
        normalize_endpoint(&endpoint.unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string()));

    let base_changed = {
        let entry = ensure_provider_entry(&mut config, &provider);
        entry.enabled = true;
        configure_cloud_endpoint(entry, &endpoint, force_cloud_base_url)
    };

    let key = match api_key {
        Some(value) if !value.trim().is_empty() => value,
        _ => {
            let prompt = format!("Enter API key for {provider}: ");
            encryption::prompt_password(&prompt)?
        }
    };

    if config.privacy.encrypt_local_data {
        let storage = Arc::new(StorageManager::new().await?);
        let manager = unlock_credential_manager(&config, storage.clone())?;
        let credentials = ApiCredentials {
            api_key: key.clone(),
            endpoint: endpoint.clone(),
        };
        manager
            .store_credentials(OLLAMA_CLOUD_CREDENTIAL_ID, &credentials)
            .await?;
        // Ensure plaintext key is not persisted to disk.
        if let Some(entry) = config.providers.get_mut(&provider) {
            entry.api_key = None;
        }
    } else if let Some(entry) = config.providers.get_mut(&provider) {
        entry.api_key = Some(key.clone());
    }

    crate::config::save_config(&config)?;
    println!("Saved Ollama configuration for provider '{provider}'.");
    if config.privacy.encrypt_local_data {
        println!("API key stored securely in the encrypted credential vault.");
    } else {
        println!("API key stored in plaintext configuration (encryption disabled).");
    }
    if !force_cloud_base_url && !base_changed {
        println!(
            "Local base URL preserved; cloud endpoint stored as {}.",
            CLOUD_ENDPOINT_KEY
        );
    }
    Ok(())
}

async fn status(provider: String) -> Result<()> {
    let provider = canonical_provider_name(&provider);
    let mut config = crate::config::try_load_config().unwrap_or_default();
    let storage = Arc::new(StorageManager::new().await?);
    let manager = if config.privacy.encrypt_local_data {
        Some(unlock_credential_manager(&config, storage.clone())?)
    } else {
        None
    };

    let api_key = hydrate_api_key(&mut config, manager.as_ref()).await?;
    {
        let entry = ensure_provider_entry(&mut config, &provider);
        entry.enabled = true;
        configure_cloud_endpoint(entry, DEFAULT_CLOUD_ENDPOINT, false);
    }

    let provider_cfg = config
        .provider(&provider)
        .cloned()
        .ok_or_else(|| anyhow!("Provider '{provider}' is not configured"))?;

    let endpoint =
        resolve_cloud_endpoint(&provider_cfg).unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string());
    let mut runtime_cfg = provider_cfg.clone();
    runtime_cfg.base_url = Some(endpoint.clone());
    runtime_cfg.extra.insert(
        OLLAMA_MODE_KEY.to_string(),
        Value::String("cloud".to_string()),
    );

    let ollama = OllamaProvider::from_config(&runtime_cfg, Some(&config.general))
        .with_context(|| "Failed to construct Ollama provider. Run `owlen cloud setup` first.")?;

    match ollama.health_check().await {
        Ok(_) => {
            println!("✓ Connected to {provider} ({})", endpoint);
            if api_key.is_none() && config.privacy.encrypt_local_data {
                println!(
                    "Warning: No API key stored; connection succeeded via environment variables."
                );
            }
        }
        Err(err) => {
            println!("✗ Failed to reach {provider}: {err}");
        }
    }

    Ok(())
}

async fn models(provider: String) -> Result<()> {
    let provider = canonical_provider_name(&provider);
    let mut config = crate::config::try_load_config().unwrap_or_default();
    let storage = Arc::new(StorageManager::new().await?);
    let manager = if config.privacy.encrypt_local_data {
        Some(unlock_credential_manager(&config, storage.clone())?)
    } else {
        None
    };
    hydrate_api_key(&mut config, manager.as_ref()).await?;

    {
        let entry = ensure_provider_entry(&mut config, &provider);
        entry.enabled = true;
        configure_cloud_endpoint(entry, DEFAULT_CLOUD_ENDPOINT, false);
    }

    let provider_cfg = config
        .provider(&provider)
        .cloned()
        .ok_or_else(|| anyhow!("Provider '{provider}' is not configured"))?;

    let endpoint =
        resolve_cloud_endpoint(&provider_cfg).unwrap_or_else(|| DEFAULT_CLOUD_ENDPOINT.to_string());
    let mut runtime_cfg = provider_cfg.clone();
    runtime_cfg.base_url = Some(endpoint);
    runtime_cfg.extra.insert(
        OLLAMA_MODE_KEY.to_string(),
        Value::String("cloud".to_string()),
    );

    let ollama = OllamaProvider::from_config(&runtime_cfg, Some(&config.general))
        .with_context(|| "Failed to construct Ollama provider. Run `owlen cloud setup` first.")?;

    match ollama.list_models().await {
        Ok(models) => {
            if models.is_empty() {
                println!("No cloud models reported by '{}'.", provider);
            } else {
                println!("Models available via '{}':", provider);
                for model in models {
                    if let Some(description) = &model.description {
                        println!("  - {} ({})", model.id, description);
                    } else {
                        println!("  - {}", model.id);
                    }
                }
            }
        }
        Err(err) => {
            bail!("Failed to list models: {err}");
        }
    }

    Ok(())
}

async fn logout(provider: String) -> Result<()> {
    let provider = canonical_provider_name(&provider);
    let mut config = crate::config::try_load_config().unwrap_or_default();
    let storage = Arc::new(StorageManager::new().await?);

    if config.privacy.encrypt_local_data {
        let manager = unlock_credential_manager(&config, storage.clone())?;
        manager
            .delete_credentials(OLLAMA_CLOUD_CREDENTIAL_ID)
            .await?;
    }

    if let Some(entry) = config.providers.get_mut(&provider) {
        entry.api_key = None;
        entry.enabled = false;
    }

    crate::config::save_config(&config)?;
    println!("Cleared credentials for provider '{provider}'.");
    Ok(())
}

fn ensure_provider_entry<'a>(config: &'a mut Config, provider: &str) -> &'a mut ProviderConfig {
    core_config::ensure_provider_config_mut(config, provider)
}

fn configure_cloud_endpoint(entry: &mut ProviderConfig, endpoint: &str, force: bool) -> bool {
    let normalized = normalize_endpoint(endpoint);
    let previous_base = entry.base_url.clone();
    entry.extra.insert(
        CLOUD_ENDPOINT_KEY.to_string(),
        Value::String(normalized.clone()),
    );

    if entry.api_key_env.is_none() {
        entry.api_key_env = Some(OLLAMA_CLOUD_API_KEY_ENV.to_string());
    }

    if force
        || entry
            .base_url
            .as_ref()
            .map(|value| value.trim().is_empty())
            .unwrap_or(true)
    {
        entry.base_url = Some(normalized.clone());
    }

    if force {
        entry.enabled = true;
    }

    entry.base_url != previous_base
}

fn resolve_cloud_endpoint(cfg: &ProviderConfig) -> Option<String> {
    if let Some(value) = cfg
        .extra
        .get(CLOUD_ENDPOINT_KEY)
        .and_then(|value| value.as_str())
        .map(normalize_endpoint)
    {
        return Some(value);
    }

    cfg.base_url
        .as_ref()
        .map(|value| value.trim_end_matches('/').to_string())
        .filter(|value| !value.is_empty())
}

fn normalize_endpoint(endpoint: &str) -> String {
    let trimmed = endpoint.trim().trim_end_matches('/');
    if trimmed.is_empty() {
        DEFAULT_CLOUD_ENDPOINT.to_string()
    } else {
        trimmed.to_string()
    }
}

fn canonical_provider_name(provider: &str) -> String {
    let normalized = provider.trim().to_ascii_lowercase().replace('-', "_");
    match normalized.as_str() {
        "" => CLOUD_PROVIDER_KEY.to_string(),
        "ollama" => CLOUD_PROVIDER_KEY.to_string(),
        "ollama_cloud" => CLOUD_PROVIDER_KEY.to_string(),
        value => value.to_string(),
    }
}

pub(crate) fn set_env_var<K, V>(key: K, value: V)
where
    K: AsRef<OsStr>,
    V: AsRef<OsStr>,
{
    // Safety: the CLI updates process-wide environment variables during startup while no
    // other threads are mutating the environment.
    unsafe {
        std::env::set_var(key, value);
    }
}

fn set_env_if_missing(var: &str, value: &str) {
    if std::env::var(var)
        .map(|v| v.trim().is_empty())
        .unwrap_or(true)
    {
        set_env_var(var, value);
    }
}

fn unlock_credential_manager(
    config: &Config,
    storage: Arc<StorageManager>,
) -> Result<Arc<CredentialManager>> {
    if !config.privacy.encrypt_local_data {
        bail!("Credential manager requested but encryption is disabled");
    }

    let secure_path = vault_path(&storage)?;
    let handle = unlock_vault(&secure_path)?;
    let master_key = Arc::new(handle.data.master_key.clone());
    Ok(Arc::new(CredentialManager::new(
        storage,
        master_key.clone(),
    )))
}

fn vault_path(storage: &StorageManager) -> Result<PathBuf> {
    let base_dir = storage
        .database_path()
        .parent()
        .map(|p| p.to_path_buf())
        .or_else(dirs::data_local_dir)
        .unwrap_or_else(|| PathBuf::from("."));
    Ok(base_dir.join("encrypted_data.json"))
}

fn unlock_vault(path: &Path) -> Result<encryption::VaultHandle> {
    use std::env;

    if path.exists() {
        if let Some(password) = env::var("OWLEN_MASTER_PASSWORD")
            .ok()
            .map(|value| value.trim().to_string())
            .filter(|password| !password.is_empty())
        {
            return encryption::unlock_with_password(path.to_path_buf(), &password)
                .context("Failed to unlock vault with OWLEN_MASTER_PASSWORD");
        }

        for attempt in 0..3 {
            let password = encryption::prompt_password("Enter master password: ")?;
            match encryption::unlock_with_password(path.to_path_buf(), &password) {
                Ok(handle) => {
                    set_env_var("OWLEN_MASTER_PASSWORD", password);
                    return Ok(handle);
                }
                Err(err) => {
                    eprintln!("Failed to unlock vault: {err}");
                    if attempt == 2 {
                        return Err(err);
                    }
                }
            }
        }

        bail!("Unable to unlock encrypted credential vault");
    }

    let handle = encryption::unlock_interactive(path.to_path_buf())?;
    if env::var("OWLEN_MASTER_PASSWORD")
        .map(|v| v.trim().is_empty())
        .unwrap_or(true)
    {
        let password = encryption::prompt_password("Cache master password for this session: ")?;
        set_env_var("OWLEN_MASTER_PASSWORD", password);
    }
    Ok(handle)
}

async fn hydrate_api_key(
    config: &mut Config,
    manager: Option<&Arc<CredentialManager>>,
) -> Result<Option<String>> {
    let credentials = match manager {
        Some(manager) => manager.get_credentials(OLLAMA_CLOUD_CREDENTIAL_ID).await?,
        None => None,
    };

    if let Some(credentials) = credentials {
        let key = credentials.api_key.trim().to_string();
        if !key.is_empty() {
            set_env_if_missing("OLLAMA_API_KEY", &key);
            set_env_if_missing("OLLAMA_CLOUD_API_KEY", &key);
        }

        let cfg = core_config::ensure_provider_config_mut(config, CLOUD_PROVIDER_KEY);
        configure_cloud_endpoint(cfg, &credentials.endpoint, false);
        return Ok(Some(key));
    }

    if let Some(key) = config
        .provider(CLOUD_PROVIDER_KEY)
        .and_then(|cfg| cfg.api_key.as_ref())
        .map(|value| value.trim())
        .filter(|value| !value.is_empty())
    {
        set_env_if_missing("OLLAMA_API_KEY", key);
        set_env_if_missing("OLLAMA_CLOUD_API_KEY", key);
        return Ok(Some(key.to_string()));
    }
    Ok(None)
}

pub async fn load_runtime_credentials(
    config: &mut Config,
    storage: Arc<StorageManager>,
) -> Result<()> {
    if config.privacy.encrypt_local_data {
        let manager = unlock_credential_manager(config, storage.clone())?;
        hydrate_api_key(config, Some(&manager)).await?;
    } else {
        hydrate_api_key(config, None).await?;
    }
    Ok(())
}

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

    #[test]
    fn canonicalises_provider_names() {
        assert_eq!(canonical_provider_name("OLLAMA_CLOUD"), CLOUD_PROVIDER_KEY);
        assert_eq!(canonical_provider_name(" ollama-cloud"), CLOUD_PROVIDER_KEY);
        assert_eq!(canonical_provider_name(""), CLOUD_PROVIDER_KEY);
    }
}