mecha10-models 0.6.2

//! 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/<name>/labels.txt
        let model_dir = self.models_dir.join(model_name);
        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.models_dir.join(model_name);
        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.models_dir.join(model_name);
        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(())
    }
}