mecha10-models 0.9.0

//! Model artifact download (the `.onnx` weights themselves)

use super::ModelService;
use anyhow::{Context, Result};
use indicatif::{ProgressBar, ProgressStyle};
use mecha10_cli_core::paths;
use std::path::PathBuf;
use tokio::fs;

impl ModelService {
    /// Download a model from the catalog
    #[cfg(feature = "model-download")]
    pub async fn pull(&self, name: &str, progress: Option<&ProgressBar>) -> Result<PathBuf> {
        let entry = self
            .get_catalog_entry(name)
            .await?
            .context(format!("Model '{}' not found in catalog", name))?;

        let model_dir = self.catalog_model_dir(name);
        fs::create_dir_all(&model_dir).await?;

        let model_path = self
            .pull_from_repo(&entry.repo, &entry.filename, name, progress)
            .await?;

        if !entry.classes.is_empty() {
            self.write_inline_labels(name, &entry.classes).await?;
        } else if entry.task == "object-detection" {
            self.pull_labels_from_repo(&entry, name, progress).await?;
        } else if entry.task == "image-classification" {
            self.pull_labels_file(name, "imagenet-labels.txt", progress).await?;
        }

        if entry.task == "object-detection" {
            if let Some(class_names) = self.read_resolved_labels(name).await? {
                self.embed_class_names_metadata(&model_path, &class_names, progress)
                    .await?;
            }
        }

        self.generate_model_config(&entry, &model_path, progress).await?;

        if let Some(quantize_config) = &entry.quantize {
            if quantize_config.enabled {
                self.quantize_model(&model_path, quantize_config, progress).await?;
            }
        }

        if let Some(pb) = progress {
            pb.set_message(format!("✅ Model '{}' ready at {}", name, model_dir.display()));
        }

        Ok(model_path)
    }

    /// Download a model via `mecha10-api` (mirrors HuggingFace repo/filename structure)
    ///
    /// The private Minio bucket is never contacted directly - this calls
    /// `GET /api/downloads/model/{repo}/{filename}`, which streams
    /// `models/{repo}/{filename}` out of Minio from inside `mecha10-api`.
    #[cfg(feature = "model-download")]
    pub async fn pull_from_repo(
        &self,
        repo: &str,
        filename: &str,
        name: &str,
        progress: Option<&ProgressBar>,
    ) -> Result<PathBuf> {
        let model_dir = self.catalog_model_dir(name);
        fs::create_dir_all(&model_dir).await?;

        let output_path = model_dir.join("model.onnx");

        if output_path.exists() {
            if let Some(pb) = progress {
                pb.set_message(format!("Model '{}' already cached", name));
            }
            return Ok(output_path);
        }

        let url = paths::urls::mecha10_api_path(&format!("/downloads/model/{repo}/{filename}"));

        if let Some(pb) = progress {
            pb.set_style(
                ProgressStyle::default_spinner()
                    .template("{spinner:.green} {msg}")
                    .unwrap(),
            );
            pb.set_message(format!("Downloading {}", name));
        }

        tracing::info!("Downloading model from mecha10-api: {}", url);

        let client = reqwest::Client::builder()
            .user_agent("mecha10-cli")
            .build()
            .context("Failed to build HTTP client")?;

        let response = client
            .get(&url)
            .send()
            .await
            .context(format!("Failed to connect to mecha10-api for model '{}'", name))?;

        if !response.status().is_success() {
            return Err(anyhow::anyhow!(
                "Model not found: HTTP {}\nURL: {}",
                response.status(),
                url
            ));
        }

        let mut file = tokio::fs::File::create(&output_path).await?;
        let mut stream = response.bytes_stream();

        use futures_util::StreamExt;
        use tokio::io::AsyncWriteExt;

        while let Some(chunk) = stream.next().await {
            let chunk = chunk.context("Error downloading model chunk")?;
            file.write_all(&chunk).await?;
        }
        file.flush().await?;

        if let Some(pb) = progress {
            pb.set_message(format!("Downloaded {} successfully", name));
        }

        Ok(output_path)
    }
}