mecha10-cli 0.8.0

//! AI model availability checks before spawning nodes.
//!
//! Used in standalone mode to ensure a node's declared models are downloaded before it spawns.

use crate::services::ModelService;
use anyhow::Result;
use mecha10_core::context::split_base_and_custom;
use serde::Deserialize;
use std::path::PathBuf;

/// Minimal config structure to extract model requirements from node configs
///
/// Deserialized from the merged output of `split_base_and_custom`, so this parser is agnostic
/// to whether a field lives on the shared base or under `custom`.
#[derive(Deserialize, Default)]
struct NodeModelConfig {
    model: Option<ModelInfo>,
}

#[derive(Deserialize)]
struct ModelInfo {
    name: Option<String>,
    /// LLM provider (`"local"` means Ollama - see `ensure_ollama_models_available`).
    #[serde(default)]
    provider: Option<String>,
    /// `"registry"` means the model is resolved via `mecha10 model set`, not `model pull` -
    /// see `ModelReference::is_catalog()` in `handlers/model.rs`.
    #[serde(default)]
    source: Option<String>,
}

/// Parse a node's `config.json` contents, resolving the active `dev`/`production` environment
/// section and then running it through the real [`split_base_and_custom`] split before reading
/// `model`/`model.provider` off the merged result.
fn parse_node_config(content: &str) -> Option<NodeModelConfig> {
    let profile = std::env::var("MECHA10_ENVIRONMENT").unwrap_or_else(|_| "dev".to_string());

    let value: serde_json::Value = serde_json::from_str(content).ok()?;
    let obj = value.as_object()?;

    let env_key = match profile.as_str() {
        "production" | "prod" => "production",
        _ => "dev",
    };
    let section = obj.get(env_key)?;

    let merged = split_base_and_custom(section).ok()?;
    serde_json::from_value(merged).ok()
}

/// Ensure all required AI models are available for nodes
///
/// Reads node configs and downloads missing models before spawning nodes.
pub(super) async fn ensure_node_models_available(node_names: &[String]) -> Result<()> {
    let model_service = match ModelService::with_models_dir("models".into()) {
        Ok(service) => service,
        Err(e) => {
            eprintln!("⚠️  Warning: Could not initialize model service: {}", e);
            return Ok(());
        }
    };

    for node_name in node_names {
        let (node_name, _instance_id) = node_name.split_once("::").unwrap_or((node_name, ""));
        let short_name = node_name.strip_prefix("@mecha10/").unwrap_or(node_name);

        let config_paths = crate::paths::config::node_config_candidates(short_name);

        for config_path in &config_paths {
            let path = PathBuf::from(config_path);
            if !path.exists() {
                continue;
            }

            let content = match std::fs::read_to_string(&path) {
                Ok(c) => c,
                Err(_) => continue,
            };

            let config = match parse_node_config(&content) {
                Some(c) => c,
                None => continue,
            };

            if let Some(model_info) = config.model {
                if model_info.source.as_deref() == Some("registry") {
                    break;
                }
                if model_info.provider.as_deref() == Some("local") {
                    break;
                }
                if let Some(model_name) = model_info.name {
                    let model_path = PathBuf::from(format!("models/catalog/{}/model.onnx", model_name));
                    if !model_path.exists() {
                        println!("📦 Downloading model '{}' for node '{}'...", model_name, node_name);
                        match model_service.pull(&model_name, None).await {
                            Ok(_) => println!("✅ Model '{}' ready", model_name),
                            Err(e) => {
                                eprintln!("⚠️  Warning: Failed to download model '{}': {}", model_name, e);
                                eprintln!("   You can manually download with: mecha10 model pull {}", model_name);
                            }
                        }
                    }
                }
            }
            break;
        }
    }

    Ok(())
}

/// Ollama's default host-mapped endpoint.
fn ollama_base_url() -> String {
    std::env::var("OLLAMA_ENDPOINT")
        .map(|endpoint| endpoint.trim_end_matches("/api/chat").to_string())
        .unwrap_or_else(|_| "http://localhost:11434".to_string())
}

/// Ensure every node declaring `"provider": "local"` (Ollama) has its `model.name` pulled before
/// nodes are spawned.
///
/// A no-op (with a warning, not an error) if Ollama isn't reachable at all.
pub(super) async fn ensure_ollama_models_available(node_names: &[String]) -> Result<()> {
    let client = reqwest::Client::new();
    let base_url = ollama_base_url();

    let mut wanted_models = Vec::new();
    for node_name in node_names {
        let (node_name, _instance_id) = node_name.split_once("::").unwrap_or((node_name, ""));
        let short_name = node_name.strip_prefix("@mecha10/").unwrap_or(node_name);

        for config_path in crate::paths::config::node_config_candidates(short_name) {
            let path = PathBuf::from(&config_path);
            if !path.exists() {
                continue;
            }

            let Ok(content) = std::fs::read_to_string(&path) else {
                continue;
            };

            let Some(config) = parse_node_config(&content) else {
                continue;
            };

            if let Some(model_info) = config.model {
                if model_info.provider.as_deref() == Some("local") {
                    if let Some(model) = model_info.name {
                        wanted_models.push((node_name.to_string(), model));
                    }
                }
            }
            break;
        }
    }

    if wanted_models.is_empty() {
        return Ok(());
    }

    let tags_url = format!("{base_url}/api/tags");
    let installed: Vec<String> = match client.get(&tags_url).send().await {
        Ok(response) => response
            .json::<serde_json::Value>()
            .await
            .ok()
            .map(|body| mecha10_core::ollama::parse_installed_model_names(&body))
            .unwrap_or_default(),
        Err(e) => {
            eprintln!("⚠️  Warning: Could not reach Ollama at {base_url}: {e}");
            eprintln!("   Local LLM nodes will fail until Ollama is running and the model is pulled.");
            return Ok(());
        }
    };

    for (node_name, model) in wanted_models {
        let already_installed = installed
            .iter()
            .any(|name| name == &model || name.starts_with(&format!("{model}:")));
        if already_installed {
            continue;
        }

        println!("📦 Pulling Ollama model '{}' for node '{}'...", model, node_name);
        let pull_result = client
            .post(format!("{base_url}/api/pull"))
            .json(&serde_json::json!({ "model": model, "stream": false }))
            .send()
            .await;

        match pull_result {
            Ok(response) if response.status().is_success() => println!("✅ Ollama model '{}' ready", model),
            Ok(response) => {
                let status = response.status();
                let body = response.text().await.unwrap_or_default();
                eprintln!("⚠️  Warning: Ollama rejected pulling '{}': {} {}", model, status, body);
                eprintln!(
                    "   You can manually pull with: docker compose exec ollama ollama pull {}",
                    model
                );
            }
            Err(e) => {
                eprintln!("⚠️  Warning: Failed to pull Ollama model '{}': {}", model, e);
                eprintln!(
                    "   You can manually pull with: docker compose exec ollama ollama pull {}",
                    model
                );
            }
        }
    }

    Ok(())
}

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

    /// Proves `parse_node_config` delegates to the real [`split_base_and_custom`].
    #[test]
    #[serial]
    fn parse_node_config_delegates_to_the_real_base_and_custom_split() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let dev_section = serde_json::json!({
            "custom": {
                "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" },
            },
            "relay_topics": ["/x"],
        });
        let content = serde_json::json!({ "dev": dev_section.clone(), "production": dev_section.clone() }).to_string();

        let config = parse_node_config(&content).expect("config should parse");
        let merged = split_base_and_custom(&dev_section).expect("split should succeed");

        assert_eq!(
            config.model.as_ref().and_then(|m| m.name.clone()),
            merged["model"]["name"].as_str().map(|s| s.to_string())
        );
        assert_eq!(
            config.model.and_then(|m| m.provider),
            merged["model"]["provider"].as_str().map(|s| s.to_string())
        );
    }

    /// `parse_node_config` must correctly read `model` nested under `custom`, like
    /// `object-detector`'s real `config.json`.
    #[test]
    #[serial]
    fn parse_node_config_reads_model_nested_under_custom_like_object_detector() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let content = r#"{
            "dev": {
                "custom": {
                    "model": { "name": "yolov8n" },
                    "input_size": 640,
                    "confidence_threshold": 0.5,
                    "topics": { "publishes": [], "subscribes": [] }
                },
                "relay_topics": []
            },
            "production": {
                "custom": {
                    "model": { "name": "yolov8n" },
                    "input_size": 640,
                    "confidence_threshold": 0.5,
                    "topics": { "publishes": [], "subscribes": [] }
                },
                "relay_topics": []
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        assert_eq!(config.model.and_then(|m| m.name), Some("yolov8n".to_string()));
    }

    /// Same regression as above, for `image-classifier`'s real shape (`packages/nodes/
    /// image-classifier/configs/config.json`), covering the `production` section too.
    #[test]
    #[serial]
    fn parse_node_config_reads_model_nested_under_custom_like_image_classifier_in_production() {
        std::env::set_var("MECHA10_ENVIRONMENT", "production");

        let content = r#"{
            "dev": {
                "custom": {
                    "processing_rate_hz": 10,
                    "top_k": 5,
                    "model": { "name": "mobilenet-v2" },
                    "topics": { "publishes": [], "subscribes": [] }
                },
                "relay_topics": []
            },
            "production": {
                "custom": {
                    "processing_rate_hz": 30,
                    "top_k": 3,
                    "model": { "name": "mobilenet-v2" },
                    "topics": { "publishes": [], "subscribes": [] }
                },
                "relay_topics": []
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        std::env::remove_var("MECHA10_ENVIRONMENT");

        assert_eq!(config.model.and_then(|m| m.name), Some("mobilenet-v2".to_string()));
    }

    /// Same as above, for `llm-command`'s Ollama-backed fields (`model.provider`/`model.name`)
    /// instead of the vision-model-only shape.
    #[test]
    #[serial]
    fn parse_node_config_reads_ollama_provider_and_model_name_nested_under_custom() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let content = r#"{
            "dev": {
                "custom": {
                    "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" }
                }
            },
            "production": {
                "custom": {
                    "model": { "name": "hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M", "provider": "local" }
                }
            }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        let model = config.model.expect("model should be present");
        assert_eq!(model.provider.as_deref(), Some("local"));
        assert_eq!(
            model.name,
            Some("hf.co/bartowski/Qwen2.5-1.5B-Instruct-GGUF:Q4_K_M".to_string())
        );
    }

    /// A node with no `model` field at all (e.g. `speaker`) must still parse cleanly with no
    /// model declared - `custom` defaulting via `#[serde(default)]` must not turn a missing
    /// `custom` block, or a `custom` block with no `model` key, into a parse failure.
    #[test]
    #[serial]
    fn parse_node_config_returns_no_model_for_a_node_with_no_model_field() {
        std::env::remove_var("MECHA10_ENVIRONMENT");

        let content = r#"{
            "dev": { "custom": { "interval_secs": 5 } },
            "production": { "custom": { "interval_secs": 1 } }
        }"#;

        let config = parse_node_config(content).expect("config should parse");
        assert!(config.model.is_none());
    }
}