mecha10-cli 0.6.3

//! `mecha10 model` command handler.
//!
//! Orchestrates AI model management operations - unifies what used to be two
//! separate CLI verbs (`LAB-1808` follow-up): the plural `mecha10 models`
//! (catalog-facing, `list`/`pull`/`installed`/`remove`/`info`, via `ModelService`)
//! and the singular `mecha10 model` (registry-facing `set`, resolves and downloads
//! a version from the private per-project model registry in
//! `packages/mecha10-api/src/models/registry` - a different system that just
//! happens to share the same on-disk `models/` directory and the same
//! `model.name` config field, disambiguated by `ObjectDetectorConfig::source`).

use crate::commands::ModelCommand;
use crate::context::CliContext;
use crate::paths;
use crate::services::credentials::CredentialsService;
use crate::services::ModelService;
use anyhow::{Context, Result};
use console::style;
use dialoguer::Confirm;
use indicatif::{ProgressBar, ProgressStyle};
use mecha10_models::services::model_service::ModelOrigin;
use serde::Deserialize;
use std::path::PathBuf;
use uuid::Uuid;

#[derive(Debug, Deserialize)]
struct ModelResponse {
    id: Uuid,
    name: String,
    latest_version: i32,
}

#[derive(Debug, Deserialize)]
struct ModelVersionResponse {
    version: i32,
}

/// `GET {base}/api/models/by-name/{name}`
fn model_by_name_url(base: &str, name: &str) -> String {
    format!(
        "{}/api/models/by-name/{}",
        base.trim_end_matches('/'),
        urlencoding_path(name)
    )
}

/// `GET {base}/api/models/{model_id}/versions/number/{version}`
fn model_version_url(base: &str, model_id: Uuid, version: i32) -> String {
    format!(
        "{}/api/models/{model_id}/versions/number/{version}",
        base.trim_end_matches('/')
    )
}

/// Minimal path-segment escaping - model names are expected to be simple
/// identifiers (no existing registry route encodes them either), but this keeps
/// at least whitespace/`/` from producing a malformed URL.
fn urlencoding_path(segment: &str) -> String {
    segment.replace(' ', "%20").replace('/', "%2F")
}

/// Merge the resolved registry selection into an existing node config JSON
/// object, preserving every other field (mirrors
/// `robot::robot_config::push_model_config`'s fetch-then-merge approach
/// server-side). Writes into the nested `model` object (`name`/`path`/
/// `source`/`version`) rather than flat top-level fields (`LAB-2081`) - this
/// must match `ObjectDetectorConfig`'s nested `model: ModelConfig` shape
/// (`packages/nodes/object-detector/src/config.rs`), which `#[serde(deny_unknown_fields)]`
/// otherwise rejects. `model_path` is the project-relative path the artifact
/// was downloaded to.
fn merge_registry_selection(
    mut config: serde_json::Value,
    model_name: &str,
    version: i32,
    model_path: &str,
) -> Result<serde_json::Value> {
    let obj = config
        .as_object_mut()
        .context("existing node config is not a JSON object")?;

    let model_obj = obj
        .entry("model".to_string())
        .or_insert_with(|| serde_json::json!({}))
        .as_object_mut()
        .context("existing node config's \"model\" field is not a JSON object")?;

    model_obj.insert("name".to_string(), serde_json::json!(model_name));
    model_obj.insert("version".to_string(), serde_json::json!(version));
    model_obj.insert("source".to_string(), serde_json::json!("registry"));
    model_obj.insert("path".to_string(), serde_json::json!(model_path));

    Ok(config)
}

/// Local path (relative to the project root) an artifact is downloaded to:
/// `models/registry/{model_name}/{version}/model.onnx` - namespaced separately
/// from the public catalog's `models/catalog/{name}/model.onnx` layout
/// (`ModelService::pull`/`ModelService::catalog_model_dir`) so a registry model
/// name colliding with a catalog name can never silently overwrite the other's
/// artifact (`LAB-1802`). Mirrors `ModelService::registry_model_dir`, just
/// relative to the project root instead of an arbitrary `models_dir`.
fn local_artifact_path(model_name: &str, version: i32) -> PathBuf {
    PathBuf::from(paths::project::MODELS_DIR)
        .join("registry")
        .join(model_name)
        .join(version.to_string())
        .join("model.onnx")
}

async fn control_plane_base_url(ctx: &mut CliContext) -> Result<String> {
    let project_config = ctx.load_project_config().await?;
    Ok(project_config.environments.control_plane_url())
}

/// Handle the `mecha10 model` command.
///
/// Spans both the public HuggingFace-backed catalog and the private
/// per-project model registry.
pub async fn handle_model(ctx: &mut CliContext, cmd: &ModelCommand) -> Result<()> {
    // `Set` talks to the control plane / registry and doesn't need the catalog
    // banner or a `ModelService` constructed ahead of time - it builds its own
    // once it has the project root.
    if let ModelCommand::Set { node, model, version } = cmd {
        return handle_set(ctx, node, model, *version).await;
    }

    println!();
    println!("{}", style("🤖 AI Model Management").bold());
    println!("{}", style("━".repeat(50)).dim());
    println!();

    // Determine models directory
    // Use ./models in current directory, or project root if in a project
    let models_dir = if let Ok(project) = ctx.project() {
        project.root().join(paths::project::MODELS_DIR)
    } else {
        PathBuf::from(paths::project::MODELS_DIR)
    };

    let service = ModelService::with_models_dir(models_dir).context("Failed to initialize model service")?;

    match cmd {
        ModelCommand::List { verbose, recommended } => {
            handle_list(&service, *verbose, *recommended)?;
        }
        ModelCommand::Pull {
            name,
            all,
            full_catalog,
            repo,
            file,
            custom_name,
        } => {
            handle_pull(
                ctx,
                &service,
                name.as_deref(),
                *all,
                *full_catalog,
                repo,
                file,
                custom_name,
            )
            .await?;
        }
        ModelCommand::Installed { verbose } => {
            handle_installed(&service, *verbose).await?;
        }
        ModelCommand::Remove { name, yes } => {
            handle_remove(&service, name, *yes).await?;
        }
        ModelCommand::Info { name, version } => {
            handle_info(&service, name, *version).await?;
        }
        ModelCommand::Set { .. } => unreachable!("handled above"),
    }

    Ok(())
}

/// Handle list subcommand
fn handle_list(service: &ModelService, verbose: bool, _recommended: bool) -> Result<()> {
    let catalog = service.list_catalog()?;

    if catalog.is_empty() {
        println!("No models found in catalog.");
        return Ok(());
    }

    println!("{} Models available in catalog:", style("📦").bold());
    println!();

    for model in &catalog {
        if verbose {
            println!("  {} {}", style("•").cyan(), style(&model.name).bold());
            println!("    Description: {}", model.description);
            println!("    Task:        {}", model.task);
            println!("    Repository:  {}", model.repo);
            println!("    File:        {}", model.filename);
            if let Some(preset) = &model.preprocessing_preset {
                println!("    Preset:      {}", preset);
            }
            println!();
        } else {
            println!(
                "  {} {} - {}",
                style("•").cyan(),
                style(&model.name).bold(),
                model.description
            );
        }
    }

    if !verbose {
        println!();
        println!("{}", style("Use --verbose for detailed information").dim());
    }

    println!();
    println!("To download a model: {}", style("mecha10 model pull <name>").cyan());
    println!();

    Ok(())
}

/// Handle pull subcommand
#[allow(clippy::too_many_arguments)]
async fn handle_pull(
    ctx: &mut CliContext,
    service: &ModelService,
    name: Option<&str>,
    all: bool,
    full_catalog: bool,
    repo: &Option<String>,
    file: &Option<String>,
    custom_name: &Option<String>,
) -> Result<()> {
    let pb = ProgressBar::new_spinner();
    pb.set_style(
        ProgressStyle::default_spinner()
            .template("{spinner:.green} {msg}")
            .unwrap(),
    );
    pb.enable_steady_tick(std::time::Duration::from_millis(100));

    if full_catalog {
        // Download every model in the curated catalog, regardless of what the current
        // project references (LAB-2043: this is the old, pre-scoping `--all` behavior,
        // now its own explicit opt-in).
        println!("Downloading the entire model catalog...");
        println!();

        let catalog = service.list_catalog()?;
        let mut paths = Vec::new();

        for entry in catalog {
            pb.set_message(format!("Downloading {}", entry.name));
            match service.pull(&entry.name, Some(&pb)).await {
                Ok(path) => paths.push(path),
                Err(e) => {
                    eprintln!("⚠️  Failed to download {}: {}", entry.name, e);
                }
            }
        }

        pb.finish_and_clear();

        println!("{} Downloaded {} models:", style("✓").green().bold(), paths.len());
        for path in &paths {
            println!("  • {}", path.display());
        }
        println!();
    } else if all {
        // LAB-2043: scope `--all` to the catalog models this project's own node configs
        // reference, rather than the entire curated catalog - `mecha10 setup`'s guidance to
        // run `mecha10 model pull --all` was pulling every model in `model_catalog.toml`
        // (dozens of entries) for a fresh project that only references 1-2 of them.
        let model_names = project_referenced_catalog_models(ctx)
            .await
            .context("failed to scan this project's node configs for referenced models")?;

        if model_names.is_empty() {
            pb.finish_and_clear();
            println!("No catalog models are referenced by this project's node configs.");
            println!();
            println!(
                "  {} to download a specific model",
                style("mecha10 model pull <name>").cyan()
            );
            println!(
                "  {} to download the entire curated catalog",
                style("mecha10 model pull --full-catalog").cyan()
            );
            println!();
            return Ok(());
        }

        println!(
            "Downloading {} model(s) referenced by this project's node configs...",
            model_names.len()
        );
        println!();

        let mut paths = Vec::new();
        for model_name in &model_names {
            pb.set_message(format!("Downloading {}", model_name));
            match service.pull(model_name, Some(&pb)).await {
                Ok(path) => paths.push(path),
                Err(e) => {
                    eprintln!("⚠️  Failed to download {}: {}", model_name, e);
                }
            }
        }

        pb.finish_and_clear();

        println!("{} Downloaded {} models:", style("✓").green().bold(), paths.len());
        for path in &paths {
            println!("  • {}", path.display());
        }
        println!();
    } else if let (Some(repo), Some(file), Some(custom_name)) = (repo, file, custom_name) {
        // Download custom model from HF repo
        pb.set_message(format!("Downloading {} from {}", custom_name, repo));

        let path = service.pull_from_repo(repo, file, custom_name, Some(&pb)).await?;

        pb.finish_and_clear();

        println!(
            "{} Downloaded custom model '{}'",
            style("✓").green().bold(),
            custom_name
        );
        println!("  Location: {}", path.display());
        println!();
    } else if let Some(name) = name {
        // Download specific model from catalog
        pb.set_message(format!("Downloading {}", name));

        let path = service.pull(name, Some(&pb)).await?;

        pb.finish_and_clear();

        println!("{} Downloaded '{}'", style("✓").green().bold(), name);
        println!("  Location: {}", path.display());
        println!();
    } else {
        pb.finish_and_clear();
        return Err(anyhow::anyhow!(
            "Please specify a model name, --all, --full-catalog, or --repo with --file and --custom-name"
        ));
    }

    Ok(())
}

/// Discover the catalog model names actually referenced by the current project's node
/// configs (`LAB-2043`).
///
/// `mecha10 model pull --all` used to iterate the entire curated catalog
/// (`model_catalog.toml`) regardless of project size - a fresh `mecha10 init` project only
/// references 1-2 models (e.g. `object-detector`'s `yolov8n`, `image-classifier`'s
/// `mobilenet-v2`) via their node configs, so pulling "all" pulled dozens of unrelated models
/// too. This scans every node configured in the project's `mecha10.json`
/// (`ProjectService::list_enabled_nodes`), resolves each node's `config.json` through the same
/// 2-tier fallback `Context::resolve_node_config` uses (framework `@mecha10` tier, then
/// project tier - first one that exists on disk wins), and collects the `model.name` of every
/// `model` block that isn't `source: "registry"` (a registry model is resolved by
/// `mecha10 model set`, not `model pull`).
///
/// Mirrors node-runner's own generic `NodeConfig.model` field (`packages/node-runner/src/
/// main.rs`) and `handlers/dev/ops/models.rs`'s `ensure_node_models_available`, but scans every
/// node in the project rather than a caller-supplied subset of already-spawned ones.
///
/// Returns an error if `ctx` isn't rooted in a mecha10 project - `--all` has nothing to scope
/// to outside one (see the caller's `--full-catalog` escape hatch for that case).
async fn project_referenced_catalog_models(ctx: &mut CliContext) -> Result<Vec<String>> {
    let node_names = ctx
        .project()
        .context(
            "`mecha10 model pull --all` must be run inside a mecha10 project (no mecha10.json found) - \
                   use `mecha10 model pull --full-catalog` to download the entire catalog instead",
        )?
        .list_enabled_nodes()
        .await?;
    let project_root = ctx.project()?.root().to_path_buf();

    let mut names = std::collections::BTreeSet::new();
    for node_name in &node_names {
        for candidate in paths::config::node_config_candidates(node_name) {
            let path = project_root.join(&candidate);
            if !path.exists() {
                continue;
            }

            let content = tokio::fs::read_to_string(&path)
                .await
                .with_context(|| format!("failed to read node config: {}", path.display()))?;

            if let Some(model_name) = extract_catalog_model_reference(&content) {
                names.insert(model_name);
            }

            // First existing tier wins - mirrors `Context::resolve_node_config`'s file-level
            // fallback (no key-by-key merging across tiers).
            break;
        }
    }

    Ok(names.into_iter().collect())
}

#[derive(Debug, Deserialize, Default)]
struct NodeModelSection {
    model: Option<ModelReference>,
}

/// A node config's `"model": {"name": ..., "source": ...}` reference. Deliberately lenient
/// about the *shape* of the raw `model` value at the call site (see
/// [`extract_catalog_model_reference`]) since some nodes (e.g. `simulator`) have an unrelated
/// top-level `model` field of their own (a plain string, not an AI-model reference).
#[derive(Debug, Deserialize)]
struct ModelReference {
    name: Option<String>,
    /// `"catalog"` (default) or `"registry"`. Mirrors node-runner's own `ModelConfig::source`
    /// and `ObjectDetectorConfig::ModelSource` (`LAB-1802`).
    #[serde(default)]
    source: Option<String>,
}

impl ModelReference {
    fn is_catalog(&self) -> bool {
        !matches!(self.source.as_deref(), Some("registry"))
    }
}

/// Extract a catalog-sourced model name from a node's raw `config.json` contents, resolving the
/// active `MECHA10_ENV` (`dev`/`production`, default `dev`) section first - mirrors
/// `Context::resolve_node_config`'s environment selection and
/// `handlers/dev/ops/models.rs::parse_node_config`'s fallback-to-flat-parsing behavior.
///
/// Returns `None` for a node with no `model` block, a non-object `model` value (e.g. the
/// `simulator` node's own unrelated `model: "rover"` string field - deserializing the whole
/// section fails on that field's type mismatch, so this falls through to `None` rather than
/// erroring, exactly like node-runner's own tolerant `deserialize_optional_model_config`), or a
/// `source: "registry"` reference (resolved by `mecha10 model set`, not `model pull`).
fn extract_catalog_model_reference(content: &str) -> Option<String> {
    let profile = std::env::var("MECHA10_ENV").unwrap_or_else(|_| "dev".to_string());

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

    // Only treat this as a `{"dev": {...}, "production": {...}}` envelope if it actually has
    // one of those keys - `dev`/`production` being `Option` fields on `NodeModelSection` means
    // *any* JSON object would otherwise deserialize into that shape trivially (both `None`),
    // silently discarding a flat (non-environment-aware) config's `model` block instead of
    // falling through to the flat-parsing branch below.
    let section_value = if obj.contains_key("dev") || obj.contains_key("production") {
        let env_key = match profile.as_str() {
            "production" | "prod" => "production",
            _ => "dev",
        };
        let fallback_key = if env_key == "dev" { "production" } else { "dev" };
        obj.get(env_key).or_else(|| obj.get(fallback_key))?.clone()
    } else {
        // Flat (non-environment-aware) config - use the whole object directly.
        value
    };

    let section: NodeModelSection = serde_json::from_value(section_value).ok()?;

    let model = section.model?;
    if !model.is_catalog() {
        return None;
    }
    model.name
}

/// Handle installed subcommand
async fn handle_installed(service: &ModelService, verbose: bool) -> Result<()> {
    let installed = service.list_installed().await?;

    if installed.is_empty() {
        println!("No models installed.");
        println!();
        println!("To download models: {}", style("mecha10 model pull <name>").cyan());
        println!();
        return Ok(());
    }

    let count_str = if installed.len() == 1 {
        "1 model".to_string()
    } else {
        format!("{} models", installed.len())
    };
    println!("{} {} installed:", style("📦").bold(), count_str);
    println!();

    for model in &installed {
        // LAB-1802: `list_installed` now also surfaces registry models (e.g. an
        // auto-registered training-job artifact locally fetched via
        // `mecha10 model set`), not just catalog ones - so every listing labels
        // which source (and, for registry, which version) each entry came from.
        let source_label = match (model.source, model.version) {
            (ModelOrigin::Registry, Some(version)) => format!("registry v{version}"),
            (ModelOrigin::Registry, None) => "registry".to_string(),
            (ModelOrigin::Catalog, _) => "catalog".to_string(),
        };

        if verbose {
            println!("  {} {}", style("•").cyan(), style(&model.name).bold());
            println!("    Path:   {}", model.path.display());
            println!("    Size:   {} bytes", model.size);
            println!("    Source: {}", source_label);

            if let Some(entry) = &model.catalog_entry {
                println!("    From catalog: {}", entry.description);
                println!("    Task: {}", entry.task);
            } else if model.source == ModelOrigin::Catalog {
                println!("    {} Custom model", style("ℹ").dim());
            }
            println!();
        } else {
            let size_mb = model.size as f64 / 1_048_576.0;
            let extra_info = if let Some(entry) = &model.catalog_entry {
                format!(" ({})", entry.task)
            } else if model.source == ModelOrigin::Catalog {
                " (custom)".to_string()
            } else {
                format!(" ({source_label})")
            };

            println!(
                "  {} {} - {:.2} MB{}",
                style("•").cyan(),
                style(&model.name).bold(),
                size_mb,
                style(extra_info).dim()
            );
        }
    }

    if !verbose {
        println!();
        println!("{}", style("Use --verbose for detailed information").dim());
    }

    println!();

    Ok(())
}

/// Handle remove subcommand
async fn handle_remove(service: &ModelService, name: &str, skip_confirm: bool) -> Result<()> {
    // Check if model exists
    let installed = service.list_installed().await?;
    let model = installed
        .iter()
        .find(|m| m.name == name)
        .ok_or_else(|| anyhow::anyhow!("Model '{}' is not installed", name))?;

    // Confirm removal
    if !skip_confirm {
        let confirm = Confirm::new()
            .with_prompt(format!("Remove model '{}'? ({} bytes)", name, model.size))
            .default(false)
            .interact()?;

        if !confirm {
            println!("Cancelled.");
            return Ok(());
        }
    }

    // Remove model
    service.remove(name).await?;

    println!("{} Removed model '{}'", style("✓").green().bold(), name);
    println!();

    Ok(())
}

/// Handle info subcommand
///
/// Works for a model from either origin (`LAB-1808`) - a Registry-origin
/// model (e.g. one locally fetched via `mecha10 model set`) is just as valid a
/// target as a Catalog one, since `ModelService::info` already looks across
/// both namespaced subtrees. `version` disambiguates which installed Registry
/// version to report when more than one is cached locally; ignored for a
/// Catalog model (versionless/always-latest).
async fn handle_info(service: &ModelService, name: &str, version: Option<i32>) -> Result<()> {
    let info = service.info(name, version).await?;

    println!("{} Model: {}", style("ℹ").blue().bold(), style(name).bold());
    println!();

    // Catalog information
    if let Some(entry) = &info.catalog_entry {
        println!("{}", style("Catalog Information:").underlined());
        println!("  Description:  {}", entry.description);
        println!("  Task:         {}", entry.task);
        println!("  Repository:   {}", entry.repo);
        println!("  File:         {}", entry.filename);
        if let Some(preset) = &entry.preprocessing_preset {
            println!("  Preset:       {}", preset);
        }
        println!();
    } else {
        println!("{}", style("Not found in catalog (custom model)").dim());
        println!();
    }

    // Installation information
    if let Some(installed) = &info.installed_info {
        // LAB-1808: every listing shows a `version` field regardless of origin
        // - "latest" for a Catalog model (versionless by design), an actual
        // integer for a Registry model - so the two origins read as one
        // consistent model concept in CLI output.
        let source_label = match installed.source {
            ModelOrigin::Registry => "registry",
            ModelOrigin::Catalog => "catalog",
        };
        let version_label = match installed.version {
            Some(v) => v.to_string(),
            None => "latest".to_string(),
        };

        println!("{}", style("Installation:").underlined());
        println!("  Status:   {}", style("Installed").green());
        println!("  Source:   {}", source_label);
        println!("  Version:  {}", version_label);
        println!("  Path:     {}", installed.path.display());
        println!(
            "  Size:     {} bytes ({:.2} MB)",
            installed.size,
            installed.size as f64 / 1_048_576.0
        );
        println!();
    } else {
        println!("{}", style("Installation:").underlined());
        println!("  Status:   {}", style("Not installed").yellow());
        println!();
        println!("  To install: {}", style(format!("mecha10 model pull {}", name)).cyan());
        println!(
            "  {} A registry model can be fetched with: {}",
            style("ℹ").dim(),
            style(format!("mecha10 model set <node> --model {}", name)).cyan()
        );
        println!();
    }

    Ok(())
}

async fn handle_set(ctx: &mut CliContext, node: &str, model_name: &str, version: Option<i32>) -> Result<()> {
    println!();
    println!("{}", style("Resolving model from registry...").bold());

    let credentials_service = CredentialsService::new();
    let credentials = credentials_service
        .load()?
        .ok_or_else(|| anyhow::anyhow!("Not logged in. Run `mecha10 auth login` first."))?;

    let base_url = control_plane_base_url(ctx).await?;
    let client = reqwest::Client::new();

    let model_resp = client
        .get(model_by_name_url(&base_url, model_name))
        .bearer_auth(&credentials.api_key)
        .send()
        .await
        .context("failed to reach the model registry")?
        .error_for_status()
        .context(format!("model '{model_name}' was not found in the registry"))?
        .json::<ModelResponse>()
        .await
        .context("model registry returned an unexpected response")?;

    let resolved_version = version.unwrap_or(model_resp.latest_version);

    let version_resp = client
        .get(model_version_url(&base_url, model_resp.id, resolved_version))
        .bearer_auth(&credentials.api_key)
        .send()
        .await
        .context("failed to reach the model registry")?
        .error_for_status()
        .context(format!("version {resolved_version} of '{model_name}' was not found"))?
        .json::<ModelVersionResponse>()
        .await
        .context("model registry returned an unexpected response")?;

    println!(
        "  {} {} v{}",
        style("Found:").dim(),
        style(&model_resp.name).bold(),
        version_resp.version
    );

    println!("{}", style("Downloading artifact...").bold());

    let project_root = ctx.project()?.root().to_path_buf();
    let models_dir = project_root.join(paths::project::MODELS_DIR);
    let model_service = ModelService::with_models_dir(models_dir).context("Failed to initialize model service")?;

    // Reuses ModelService::pull_registry_version (LAB-1802) - the same
    // by-name-lookup-then-download primitive node-runner's source-aware
    // preflight uses - instead of hand-rolling the download here, so the two
    // callers can't drift on namespacing (models/registry/<name>/<version>/).
    let absolute_path = model_service
        .pull_registry_version(
            &base_url,
            &credentials.api_key,
            &model_resp.name,
            version_resp.version,
            None,
        )
        .await
        .context("failed to download the model artifact")?;
    let relative_path = local_artifact_path(&model_resp.name, version_resp.version);

    let artifact_size = tokio::fs::metadata(&absolute_path)
        .await
        .context("failed to read downloaded artifact metadata")?
        .len();
    println!(
        "  {} {} ({} bytes)",
        style("Saved:").dim(),
        absolute_path.display(),
        artifact_size
    );

    println!("{}", style("Updating node config...").bold());
    let config_path = project_root
        .join("configs")
        .join("nodes")
        .join(node)
        .join("config.json");
    let existing = if config_path.exists() {
        let content = tokio::fs::read_to_string(&config_path)
            .await
            .context("failed to read existing node config")?;
        serde_json::from_str(&content).context("existing node config is not valid JSON")?
    } else {
        serde_json::json!({})
    };

    let relative_path_str = relative_path.to_string_lossy().replace('\\', "/");
    let updated = merge_registry_selection(existing, &model_resp.name, version_resp.version, &relative_path_str)?;

    if let Some(parent) = config_path.parent() {
        tokio::fs::create_dir_all(parent)
            .await
            .context("failed to create node config directory")?;
    }
    tokio::fs::write(&config_path, serde_json::to_string_pretty(&updated)?)
        .await
        .context("failed to write node config")?;

    println!("  {} {}", style("Updated:").dim(), config_path.display());
    println!();
    println!(
        "{} {} is now set to {} v{} (registry)",
        style("✔").green().bold(),
        node,
        model_resp.name,
        version_resp.version
    );
    println!();

    Ok(())
}

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

    #[test]
    fn model_by_name_url_appends_the_expected_path() {
        assert_eq!(
            model_by_name_url("https://control-plane.example.com", "rover-detector"),
            "https://control-plane.example.com/api/models/by-name/rover-detector"
        );
    }

    #[test]
    fn model_by_name_url_strips_a_trailing_slash_from_the_base() {
        assert_eq!(
            model_by_name_url("https://control-plane.example.com/", "rover-detector"),
            "https://control-plane.example.com/api/models/by-name/rover-detector"
        );
    }

    #[test]
    fn model_version_url_is_derived_consistently() {
        let id = Uuid::nil();
        assert_eq!(
            model_version_url("https://cp.example.com", id, 3),
            format!("https://cp.example.com/api/models/{id}/versions/number/3")
        );
    }

    #[test]
    fn local_artifact_path_is_namespaced_by_source_and_version() {
        assert_eq!(
            local_artifact_path("rover-detector", 4),
            PathBuf::from("models/registry/rover-detector/4/model.onnx")
        );
    }

    #[test]
    fn local_artifact_path_never_collides_across_versions() {
        assert_ne!(
            local_artifact_path("rover-detector", 1),
            local_artifact_path("rover-detector", 2)
        );
    }

    #[test]
    fn merge_registry_selection_preserves_unrelated_fields() {
        let existing = serde_json::json!({
            "confidence_threshold": 0.6,
            "model": {
                "name": "yolov8n",
                "source": "catalog",
            },
        });

        let merged = merge_registry_selection(
            existing,
            "rover-detector",
            4,
            "models/registry/rover-detector/4/model.onnx",
        )
        .unwrap();

        assert_eq!(merged["confidence_threshold"], serde_json::json!(0.6));
        assert_eq!(merged["model"]["name"], serde_json::json!("rover-detector"));
        assert_eq!(merged["model"]["version"], serde_json::json!(4));
        assert_eq!(merged["model"]["source"], serde_json::json!("registry"));
        assert_eq!(
            merged["model"]["path"],
            serde_json::json!("models/registry/rover-detector/4/model.onnx")
        );
    }

    #[test]
    fn merge_registry_selection_works_on_an_empty_config() {
        let merged = merge_registry_selection(serde_json::json!({}), "m", 1, "models/registry/m/1/model.onnx").unwrap();
        assert_eq!(merged["model"]["source"], serde_json::json!("registry"));
    }

    #[test]
    fn merge_registry_selection_rejects_a_non_object_config() {
        let err = merge_registry_selection(serde_json::json!([1, 2, 3]), "m", 1, "models/registry/m/1/model.onnx")
            .unwrap_err();
        assert!(err.to_string().contains("not a JSON object"));
    }

    // ===== extract_catalog_model_reference (LAB-2043) =====
    //
    // These mutate the process-wide `MECHA10_ENV` var, so they're `#[serial]` (same pattern as
    // `mecha10_core::context::node_config::tests`) rather than a manually-locked `Mutex` guard -
    // the latter would need to stay held across the `.await` points in the
    // `project_referenced_catalog_models_*` tests further down, which `clippy::await_holding_lock`
    // rightly flags.

    #[test]
    #[serial]
    fn extract_catalog_model_reference_reads_the_dev_section_by_default() {
        std::env::remove_var("MECHA10_ENV");

        let content = r#"{"dev": {"model": {"name": "yolov8n"}}, "production": {"model": {"name": "yolov8n-prod"}}}"#;
        assert_eq!(extract_catalog_model_reference(content), Some("yolov8n".to_string()));
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_honors_mecha10_env_production() {
        std::env::set_var("MECHA10_ENV", "production");

        let content = r#"{"dev": {"model": {"name": "yolov8n"}}, "production": {"model": {"name": "yolov8n-prod"}}}"#;
        let result = extract_catalog_model_reference(content);

        std::env::remove_var("MECHA10_ENV");
        assert_eq!(result, Some("yolov8n-prod".to_string()));
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_ignores_a_registry_sourced_model() {
        std::env::remove_var("MECHA10_ENV");

        let content = r#"{"dev": {"model": {"name": "my-detector", "source": "registry"}}}"#;
        assert_eq!(extract_catalog_model_reference(content), None);
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_treats_missing_source_as_catalog() {
        std::env::remove_var("MECHA10_ENV");

        let content = r#"{"dev": {"model": {"name": "mobilenet-v2"}}}"#;
        assert_eq!(
            extract_catalog_model_reference(content),
            Some("mobilenet-v2".to_string())
        );
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_returns_none_for_a_node_with_no_model_field() {
        std::env::remove_var("MECHA10_ENV");

        let content = r#"{"dev": {"camera_id": "camera0"}, "production": {"camera_id": "camera0"}}"#;
        assert_eq!(extract_catalog_model_reference(content), None);
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_returns_none_for_a_non_object_model_field() {
        std::env::remove_var("MECHA10_ENV");

        // Mirrors the `simulator` node's own unrelated `model: "rover"` string field (a
        // physics/robot model, not an AI-model-catalog reference) - see node-runner's
        // `deserialize_optional_model_config` for the equivalent tolerance there.
        let content = r#"{"dev": {"model": "rover"}, "production": {"model": "rover"}}"#;
        assert_eq!(extract_catalog_model_reference(content), None);
    }

    #[test]
    #[serial]
    fn extract_catalog_model_reference_falls_back_to_flat_non_environment_aware_parsing() {
        std::env::remove_var("MECHA10_ENV");

        let content = r#"{"model": {"name": "yolov8n"}}"#;
        assert_eq!(extract_catalog_model_reference(content), Some("yolov8n".to_string()));
    }

    // ===== project_referenced_catalog_models (LAB-2043) =====

    /// Scaffolds a temp project with `node_configs` (each `(node identifier as it appears in
    /// `mecha10.json`'s `nodes` map, config-dir-relative node name, config.json body)`) and
    /// returns a `CliContext` rooted there plus the `TempDir` guard (kept alive by the caller
    /// for the duration of the test).
    fn scaffold_project(node_configs: &[(&str, &str, serde_json::Value)]) -> (tempfile::TempDir, CliContext) {
        let dir = tempfile::tempdir().unwrap();

        let nodes: std::collections::HashMap<&str, &str> = node_configs
            .iter()
            .map(|(identifier, _, _)| (*identifier, "0.1.0"))
            .collect();
        let mecha10_json = serde_json::json!({
            "name": "proj",
            "version": "1.0.0",
            "robot": {"id": "robot-1"},
            "nodes": nodes,
        });
        let config_path = dir.path().join("mecha10.json");
        std::fs::write(&config_path, mecha10_json.to_string()).unwrap();

        for (identifier, config_dir_name, body) in node_configs {
            let is_framework = identifier.starts_with("@mecha10/");
            let node_dir = if is_framework {
                dir.path().join("configs/nodes/@mecha10").join(config_dir_name)
            } else {
                dir.path().join("configs/nodes").join(config_dir_name)
            };
            std::fs::create_dir_all(&node_dir).unwrap();
            std::fs::write(node_dir.join("config.json"), body.to_string()).unwrap();
        }

        let ctx = CliContextBuilder::new()
            .working_dir(dir.path().to_path_buf())
            .config_path(Some(config_path))
            .build()
            .unwrap();

        (dir, ctx)
    }

    #[tokio::test]
    #[serial]
    async fn project_referenced_catalog_models_scopes_to_a_projects_own_node_configs() {
        std::env::remove_var("MECHA10_ENV");

        // A project with exactly 2 node configs referencing 2 distinct catalog models must
        // only surface those 2 - not the entire curated catalog (`model_catalog.toml`, which
        // has dozens of entries).
        let (_dir, mut ctx) = scaffold_project(&[
            (
                "@mecha10/object-detector",
                "object-detector",
                serde_json::json!({
                    "dev": {"model": {"name": "yolov8n"}},
                    "production": {"model": {"name": "yolov8n"}},
                }),
            ),
            (
                "@mecha10/image-classifier",
                "image-classifier",
                serde_json::json!({
                    "dev": {"model": {"name": "mobilenet-v2"}},
                    "production": {"model": {"name": "mobilenet-v2"}},
                }),
            ),
        ]);

        let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
        assert_eq!(models, vec!["mobilenet-v2".to_string(), "yolov8n".to_string()]);
    }

    #[tokio::test]
    #[serial]
    async fn project_referenced_catalog_models_excludes_registry_sourced_models() {
        std::env::remove_var("MECHA10_ENV");

        let (_dir, mut ctx) = scaffold_project(&[
            (
                "@mecha10/object-detector",
                "object-detector",
                serde_json::json!({
                    "dev": {"model": {"name": "yolov8n"}},
                    "production": {"model": {"name": "yolov8n"}},
                }),
            ),
            (
                "custom-registry-node",
                "custom-registry-node",
                serde_json::json!({
                    "dev": {"model": {"name": "my-registry-model", "source": "registry"}},
                    "production": {"model": {"name": "my-registry-model", "source": "registry"}},
                }),
            ),
        ]);

        // `mecha10 model pull --all` resolves the public catalog - a `source: "registry"`
        // reference is resolved by `mecha10 model set` instead, so it must not appear here.
        let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
        assert_eq!(models, vec!["yolov8n".to_string()]);
    }

    #[tokio::test]
    #[serial]
    async fn project_referenced_catalog_models_deduplicates_a_model_shared_by_two_nodes() {
        std::env::remove_var("MECHA10_ENV");

        let (_dir, mut ctx) = scaffold_project(&[
            (
                "@mecha10/object-detector",
                "object-detector",
                serde_json::json!({
                    "dev": {"model": {"name": "yolov8n"}},
                    "production": {"model": {"name": "yolov8n"}},
                }),
            ),
            (
                "second-detector",
                "second-detector",
                serde_json::json!({
                    "dev": {"model": {"name": "yolov8n"}},
                    "production": {"model": {"name": "yolov8n"}},
                }),
            ),
        ]);

        let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
        assert_eq!(models, vec!["yolov8n".to_string()]);
    }

    #[tokio::test]
    #[serial]
    async fn project_referenced_catalog_models_returns_empty_when_no_node_declares_a_model() {
        std::env::remove_var("MECHA10_ENV");

        let (_dir, mut ctx) = scaffold_project(&[(
            "@mecha10/camera",
            "camera",
            serde_json::json!({
                "dev": {"camera_id": "camera0"},
                "production": {"camera_id": "camera0"},
            }),
        )]);

        let models = project_referenced_catalog_models(&mut ctx).await.unwrap();
        assert!(models.is_empty());
    }

    #[tokio::test]
    async fn project_referenced_catalog_models_errors_outside_a_project() {
        let dir = tempfile::tempdir().unwrap();
        let mut ctx = CliContextBuilder::new()
            .working_dir(dir.path().to_path_buf())
            .build()
            .unwrap();

        let err = project_referenced_catalog_models(&mut ctx).await.unwrap_err();
        assert!(
            err.to_string().contains("--full-catalog"),
            "error should point users at the --full-catalog escape hatch: {err}"
        );
    }
}