mecha10-models 0.9.0

//! Label file download and generation (`labels.txt`)

use super::catalog::ModelCatalogEntry;
use super::ModelService;
use anyhow::{Context, Result};
use indicatif::ProgressBar;
use std::path::PathBuf;
use tokio::fs;

impl ModelService {
    /// Download a labels file (e.g., imagenet-labels.txt) into model directory
    #[cfg(feature = "model-download")]
    pub(super) async fn pull_labels_file(
        &self,
        model_name: &str,
        filename: &str,
        progress: Option<&ProgressBar>,
    ) -> Result<PathBuf> {
        let model_dir = self.catalog_model_dir(model_name);
        fs::create_dir_all(&model_dir)
            .await
            .context("Failed to create model directory")?;
        let output_path = model_dir.join("labels.txt");

        if output_path.exists() {
            if let Some(pb) = progress {
                pb.set_message("Labels file already cached".to_string());
            }
            return Ok(output_path);
        }

        let url = match filename {
            "imagenet-labels.txt" => "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt",
            _ => {
                if let Some(pb) = progress {
                    pb.set_message(format!("⚠️  Unknown labels file: {}, skipping", filename));
                }
                return Ok(output_path);
            }
        };

        if let Some(pb) = progress {
            pb.set_message(format!("Downloading labels: {}", filename));
        }

        let client = reqwest::Client::new();
        let response = client
            .get(url)
            .send()
            .await
            .context(format!("Failed to download labels from {}", url))?;

        if !response.status().is_success() {
            anyhow::bail!("Failed to download labels: HTTP {}", response.status());
        }

        let content = response.text().await.context("Failed to read labels content")?;

        fs::write(&output_path, content)
            .await
            .context("Failed to write labels file")?;

        if let Some(pb) = progress {
            pb.set_message(format!("Downloaded labels: {}", filename));
        }

        Ok(output_path)
    }

    /// Write inline classes to labels.txt
    #[cfg(feature = "model-download")]
    pub(super) async fn write_inline_labels(&self, model_name: &str, classes: &[String]) -> Result<()> {
        let model_dir = self.catalog_model_dir(model_name);
        fs::create_dir_all(&model_dir)
            .await
            .context("Failed to create model directory")?;
        let labels_path = model_dir.join("labels.txt");

        let content = classes.join("\n");
        fs::write(&labels_path, content)
            .await
            .context("Failed to write inline labels to labels.txt")?;

        Ok(())
    }

    /// Reads back the `labels.txt` a prior resolution step (`write_inline_labels`,
    /// `pull_labels_from_repo`, or `pull_labels_file`) wrote for this model, one
    /// class name per line.
    ///
    /// Returns `None` when the file doesn't exist or resolved to zero lines - i.e.
    /// label resolution didn't actually find anything for this model - so a caller
    /// (`pull`'s ONNX metadata embedding step) can tell "resolved to nothing" apart
    /// from "resolved to an empty list" without embedding a placeholder.
    #[cfg(feature = "model-download")]
    pub(super) async fn read_resolved_labels(&self, model_name: &str) -> Result<Option<Vec<String>>> {
        let labels_path = self.catalog_model_dir(model_name).join("labels.txt");

        if !labels_path.exists() {
            return Ok(None);
        }

        let content = fs::read_to_string(&labels_path)
            .await
            .context("Failed to read resolved labels.txt")?;

        let labels: Vec<String> = content
            .lines()
            .map(str::trim)
            .filter(|line| !line.is_empty())
            .map(str::to_string)
            .collect();

        Ok(if labels.is_empty() { None } else { Some(labels) })
    }

    /// Download labels.json from HuggingFace repo and convert to labels.txt
    #[cfg(feature = "model-download")]
    pub(super) async fn pull_labels_from_repo(
        &self,
        entry: &ModelCatalogEntry,
        model_name: &str,
        progress: Option<&ProgressBar>,
    ) -> Result<()> {
        let model_dir = self.catalog_model_dir(model_name);
        fs::create_dir_all(&model_dir)
            .await
            .context("Failed to create model directory")?;
        let labels_path = model_dir.join("labels.txt");

        if labels_path.exists() {
            if let Some(pb) = progress {
                pb.set_message("Labels file already cached".to_string());
            }
            return Ok(());
        }

        let model_dir_in_repo = entry.filename.rsplit_once('/').map(|(dir, _)| dir).unwrap_or("");

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

        let labels_filename = format!("{}/labels.json", model_dir_in_repo);

        if let Some(pb) = progress {
            pb.set_message(format!("Downloading labels from {}", entry.repo));
        }

        let url = format!("https://huggingface.co/{}/raw/main/{}", entry.repo, labels_filename);

        let client = reqwest::Client::new();
        let response = match client.get(&url).send().await {
            Ok(resp) if resp.status().is_success() => resp,
            _ => {
                return Ok(());
            }
        };

        let json_content = response.text().await.context("Failed to read labels.json response")?;

        let labels: Vec<String> = serde_json::from_str(&json_content).context("Failed to parse labels.json")?;

        let content = labels.join("\n");
        fs::write(&labels_path, content)
            .await
            .context("Failed to write labels.txt")?;

        if let Some(pb) = progress {
            pb.set_message(format!("Downloaded {} class labels", labels.len()));
        }

        Ok(())
    }
}

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

    fn service(models_dir: PathBuf) -> ModelService {
        ModelService::with_models_dir(models_dir).unwrap()
    }

    fn entry(name: &str, task: &str, filename: &str, classes: Vec<String>) -> ModelCatalogEntry {
        ModelCatalogEntry {
            name: name.to_string(),
            description: "test entry".to_string(),
            task: task.to_string(),
            repo: "test/repo".to_string(),
            filename: filename.to_string(),
            preprocessing_preset: None,
            classes,
            quantize: None,
            input_size: None,
        }
    }

    #[tokio::test]
    async fn write_inline_labels_writes_to_the_namespaced_catalog_dir() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        service
            .write_inline_labels("mobilenet-v2", &["cat".to_string(), "dog".to_string()])
            .await
            .unwrap();

        let expected_path = models_dir.join("catalog").join("mobilenet-v2").join("labels.txt");
        assert!(expected_path.exists(), "expected labels.txt at {:?}", expected_path);
        let content = tokio::fs::read_to_string(&expected_path).await.unwrap();
        assert_eq!(content, "cat\ndog");

        assert!(!models_dir.join("mobilenet-v2").join("labels.txt").exists());
    }

    #[tokio::test]
    async fn pull_labels_file_writes_to_the_namespaced_catalog_dir_when_already_cached() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let catalog_dir = models_dir.join("catalog").join("mobilenet-v2");
        tokio::fs::create_dir_all(&catalog_dir).await.unwrap();
        let labels_path = catalog_dir.join("labels.txt");
        tokio::fs::write(&labels_path, "existing").await.unwrap();

        let result_path = service
            .pull_labels_file("mobilenet-v2", "imagenet-labels.txt", None)
            .await
            .unwrap();

        assert_eq!(result_path, labels_path);
    }

    #[tokio::test]
    async fn pull_labels_file_creates_the_catalog_dir_if_missing() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let result_path = service
            .pull_labels_file("mobilenet-v2", "unknown-labels.txt", None)
            .await
            .unwrap();

        assert_eq!(
            result_path,
            models_dir.join("catalog").join("mobilenet-v2").join("labels.txt")
        );
        assert!(models_dir.join("catalog").join("mobilenet-v2").exists());
    }

    #[tokio::test]
    async fn pull_labels_from_repo_writes_to_the_namespaced_catalog_dir_when_already_cached() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let catalog_dir = models_dir.join("catalog").join("yolov8n");
        tokio::fs::create_dir_all(&catalog_dir).await.unwrap();
        let labels_path = catalog_dir.join("labels.txt");
        tokio::fs::write(&labels_path, "existing").await.unwrap();

        let entry = entry("yolov8n", "object-detection", "yolov8n/model.onnx", vec![]);
        service.pull_labels_from_repo(&entry, "yolov8n", None).await.unwrap();

        let content = tokio::fs::read_to_string(&labels_path).await.unwrap();
        assert_eq!(content, "existing");
    }

    #[tokio::test]
    async fn read_resolved_labels_returns_none_when_labels_txt_does_not_exist() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir);

        let result = service.read_resolved_labels("mobilenet-v2").await.unwrap();

        assert!(
            result.is_none(),
            "no labels.txt at all must resolve to None, not an empty Vec"
        );
    }

    #[tokio::test]
    async fn read_resolved_labels_returns_none_for_an_empty_labels_txt() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let catalog_dir = models_dir.join("catalog").join("mobilenet-v2");
        tokio::fs::create_dir_all(&catalog_dir).await.unwrap();
        tokio::fs::write(catalog_dir.join("labels.txt"), "").await.unwrap();

        let result = service.read_resolved_labels("mobilenet-v2").await.unwrap();

        assert!(
            result.is_none(),
            "an empty labels.txt must resolve to None, not an empty Vec"
        );
    }

    #[tokio::test]
    async fn read_resolved_labels_parses_one_class_per_line() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        service
            .write_inline_labels("mobilenet-v2", &["cat".to_string(), "dog".to_string()])
            .await
            .unwrap();

        let result = service.read_resolved_labels("mobilenet-v2").await.unwrap();

        assert_eq!(result, Some(vec!["cat".to_string(), "dog".to_string()]));
    }

    #[tokio::test]
    async fn pull_labels_from_repo_creates_the_catalog_dir_if_missing() {
        let tmp = tempfile::tempdir().unwrap();
        let models_dir = tmp.path().to_path_buf();
        let service = service(models_dir.clone());

        let entry = entry("flat-repo-model", "object-detection", "model.onnx", vec![]);
        service
            .pull_labels_from_repo(&entry, "flat-repo-model", None)
            .await
            .unwrap();

        assert!(models_dir.join("catalog").join("flat-repo-model").exists());
        assert!(!models_dir
            .join("catalog")
            .join("flat-repo-model")
            .join("labels.txt")
            .exists());
    }
}