mecha10-models 0.6.3

//! 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> {
        // Output path: models/catalog/<name>/labels.txt (LAB-1802: namespaced
        // catalog directory - see `ModelService::catalog_model_dir`)
        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");

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

        // Map known label files to download URLs
        let url = match filename {
            "imagenet-labels.txt" => "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt",
            _ => {
                // Try to download from HuggingFace or skip if unknown
                if let Some(pb) = progress {
                    pb.set_message(format!("⚠️  Unknown labels file: {}, skipping", filename));
                }
                return Ok(output_path); // Return path even if not downloaded
            }
        };

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

        // Download using reqwest
        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")?;

        // Write to file
        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(())
    }

    /// 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");

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

        // Extract directory path from filename (e.g., "yolov8n/model.onnx" -> "yolov8n")
        let model_dir_in_repo = entry.filename.rsplit_once('/').map(|(dir, _)| dir).unwrap_or("");

        if model_dir_in_repo.is_empty() {
            // No labels.json in repo root, skip
            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));
        }

        // Download labels.json directly via HTTP
        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,
            _ => {
                // labels.json not found, skip silently (model might use inline classes)
                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")?;

        // Write to labels.txt (one class per line)
        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,
        }
    }

    // LAB-2050: all three label-writing paths must land in the LAB-1802
    // namespaced `models/catalog/<name>/` directory, not the old flat
    // `models/<name>/` path - and must not assume that directory already
    // exists (create_dir_all up front).

    #[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");

        // Must not fall back to the old flat pre-LAB-1802 location.
        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());

        // Pre-seed the "already downloaded" fast path so the test doesn't
        // depend on network access - this still exercises the directory
        // computation that was the actual bug.
        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());

        // Unknown filename short-circuits before any network call, so this
        // exercises directory creation without requiring network access.
        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();

        // Already-cached fast path returns early without touching the file,
        // but only if it was looking in the right (namespaced) directory.
        let content = tokio::fs::read_to_string(&labels_path).await.unwrap();
        assert_eq!(content, "existing");
    }

    #[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());

        // No directory in the repo filename means it skips before any
        // network call, so this exercises directory creation without
        // requiring network access.
        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());
    }
}