mecha10-models 0.6.3

//! Preprocessing auto-detection and `config.json` generation
//!
//! Auto-detects normalization (mean/std) and input size for a downloaded
//! model, falling back through: HuggingFace `preprocessor_config.json` ->
//! catalog preprocessing preset -> ONNX metadata / sensible defaults. The
//! result is written out as the model's `config.json`.

use super::ModelService;
use std::path::Path;
#[cfg(feature = "model-download")]
use {
    super::catalog::{HFPreprocessorConfig, ModelCatalogEntry, PreprocessingPreset},
    anyhow::{Context, Result},
    indicatif::ProgressBar,
    mecha10_core::model::{CustomLabelsConfig, ModelConfig, PreprocessingConfig},
    tokio::fs,
};

impl ModelService {
    /// Fetch preprocessor_config.json from HuggingFace repo
    #[cfg(feature = "model-download")]
    pub(super) async fn fetch_hf_preprocessor_config(&self, repo: &str) -> Result<HFPreprocessorConfig> {
        let url = format!("https://huggingface.co/{}/raw/main/preprocessor_config.json", repo);

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

        if !response.status().is_success() {
            anyhow::bail!(
                "HuggingFace preprocessor_config.json not found for {} (HTTP {})",
                repo,
                response.status()
            );
        }

        let config: HFPreprocessorConfig = response
            .json()
            .await
            .context("Failed to parse preprocessor_config.json")?;

        Ok(config)
    }

    /// Extract input size from ONNX model metadata
    ///
    /// Note: This is best-effort and may fail depending on the ONNX model format.
    /// Falls back to None if extraction fails.
    #[cfg(feature = "vision")]
    pub(super) fn extract_input_size_from_onnx(&self, model_path: &Path) -> Option<[u32; 2]> {
        use ort::session::Session;

        // Best-effort extraction - don't fail if ONNX model is not parseable
        let session = Session::builder().ok()?.commit_from_file(model_path).ok()?;

        // Get first input tensor
        let _input = session.inputs.first()?;

        // Try to extract dimensions from the input type
        // Note: The ort API has changed between versions, so this is best-effort
        // For now, we'll return None and rely on HF config or presets
        // TODO: Update when ort API is stabilized

        None
    }

    /// Stub when vision feature is disabled
    #[cfg(not(feature = "vision"))]
    pub(super) fn extract_input_size_from_onnx(&self, _model_path: &Path) -> Option<[u32; 2]> {
        None
    }

    /// Auto-detect preprocessing parameters with fallback priority:
    /// 1. HuggingFace preprocessor_config.json
    /// 2. Preprocessing preset from catalog
    /// 3. ONNX metadata (input size only, use sensible defaults for mean/std)
    #[cfg(feature = "model-download")]
    pub(super) async fn auto_detect_preprocessing(
        &self,
        entry: &ModelCatalogEntry,
        model_path: &Path,
        progress: Option<&ProgressBar>,
    ) -> Result<(PreprocessingConfig, [u32; 2])> {
        // Try HuggingFace preprocessor_config.json first
        if let Some(pb) = progress {
            pb.set_message(format!("🔍 Auto-detecting preprocessing for {}", entry.name));
        }

        if let Ok(hf_config) = self.fetch_hf_preprocessor_config(&entry.repo).await {
            tracing::debug!(
                "HF config: size={:?}, crop_size={:?}",
                hf_config.size,
                hf_config.crop_size
            );

            let preprocessing = hf_config.to_preprocessing();
            let input_size = hf_config.input_size().unwrap_or([224, 224]);

            tracing::debug!(
                "Detected preprocessing: mean={:?}, std={:?}, input_size={:?}",
                preprocessing.mean,
                preprocessing.std,
                input_size
            );

            if let Some(pb) = progress {
                pb.set_message(format!(
                    "✅ Auto-detected from HuggingFace (input_size={:?})",
                    input_size
                ));
            }

            return Ok((preprocessing, input_size));
        } else {
            tracing::debug!("Failed to fetch HuggingFace preprocessor config, falling back to preset");
        }

        // Try preprocessing preset
        if let Some(preset_name) = &entry.preprocessing_preset {
            if let Ok(preset) = PreprocessingPreset::from_name(preset_name) {
                let preprocessing = preset.to_config();

                // Try to get input_size from ONNX (best-effort)
                let input_size = self.extract_input_size_from_onnx(model_path).unwrap_or([224, 224]);

                if let Some(pb) = progress {
                    pb.set_message(format!(
                        "✅ Using preset '{}' (input_size={:?})",
                        preset_name, input_size
                    ));
                }

                return Ok((preprocessing, input_size));
            }
        }

        // Fallback: Try ONNX metadata for input_size, use sensible defaults
        let input_size = self.extract_input_size_from_onnx(model_path).unwrap_or([224, 224]);

        let preprocessing = PreprocessingConfig {
            mean: [0.0, 0.0, 0.0],
            std: [1.0, 1.0, 1.0],
            channel_order: "RGB".to_string(),
        };

        if let Some(pb) = progress {
            pb.set_message(format!(
                "⚠️  Using fallback preprocessing (input_size={:?}). Consider editing config.json",
                input_size
            ));
        }

        Ok((preprocessing, input_size))
    }

    /// Generate and write config.json for a model with auto-detected values
    #[cfg(feature = "model-download")]
    pub(super) async fn generate_model_config(
        &self,
        entry: &ModelCatalogEntry,
        model_path: &Path,
        progress: Option<&ProgressBar>,
    ) -> Result<()> {
        // Output path: models/catalog/<name>/config.json (LAB-1802: namespaced
        // catalog directory - see `ModelService::catalog_model_dir`)
        let model_dir = self.catalog_model_dir(&entry.name);
        fs::create_dir_all(&model_dir)
            .await
            .context("Failed to create model directory")?;
        let config_path = model_dir.join("config.json");

        // Auto-detect preprocessing
        let (preprocessing, input_size) = self.auto_detect_preprocessing(entry, model_path, progress).await?;

        // Determine num_classes
        let num_classes = if entry.task == "object-detection" {
            entry.classes.len().max(1)
        } else {
            1000 // Default for image classification (ImageNet)
        };

        // Convert catalog entry to ModelConfig
        let config = ModelConfig {
            name: entry.name.clone(),
            task: entry.task.clone(),
            repo: entry.repo.clone(),
            filename: entry.filename.clone(),
            input_size,
            preprocessing,
            num_classes,
            labels_file: "labels.txt".to_string(),
            custom_labels: CustomLabelsConfig::default(),
        };

        // Write config as pretty JSON
        let json = serde_json::to_string_pretty(&config).context("Failed to serialize model config")?;

        fs::write(&config_path, json)
            .await
            .context("Failed to write model config.json")?;

        if let Some(pb) = progress {
            pb.set_message(format!("📝 Wrote config to {}", config_path.display()));
        }

        Ok(())
    }
}

#[cfg(all(test, feature = "model-download"))]
mod tests {
    use super::super::ModelService;
    use super::*;
    use std::path::PathBuf;

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

    fn entry(name: &str, task: &str) -> ModelCatalogEntry {
        ModelCatalogEntry {
            name: name.to_string(),
            description: "test entry".to_string(),
            task: task.to_string(),
            // A repo that does not exist on HuggingFace, so
            // `fetch_hf_preprocessor_config` fails fast (404) and
            // `auto_detect_preprocessing` falls through to defaults instead
            // of depending on real preprocessor_config.json content.
            repo: "mecha10-test-fixtures/does-not-exist".to_string(),
            filename: "model.onnx".to_string(),
            preprocessing_preset: None,
            classes: vec![],
            quantize: None,
        }
    }

    // LAB-2057: generate_model_config must write to the LAB-1802 namespaced
    // `models/catalog/<name>/` directory, not the old flat `models/<name>/`
    // path (same bug class as LAB-2050's labels.rs fix).
    #[tokio::test]
    async fn generate_model_config_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());

        let entry = entry("mobilenet-v2", "image-classification");
        let model_path = models_dir.join("catalog").join("mobilenet-v2").join("model.onnx");

        service.generate_model_config(&entry, &model_path, None).await.unwrap();

        let expected_path = models_dir.join("catalog").join("mobilenet-v2").join("config.json");
        assert!(expected_path.exists(), "expected config.json at {:?}", expected_path);

        let content = tokio::fs::read_to_string(&expected_path).await.unwrap();
        let config: ModelConfig = serde_json::from_str(&content).unwrap();
        assert_eq!(config.name, "mobilenet-v2");

        // Must not fall back to the old flat pre-LAB-1802 location.
        assert!(!models_dir.join("mobilenet-v2").join("config.json").exists());
    }

    #[tokio::test]
    async fn generate_model_config_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 pre-existing directory at all - exercises the create_dir_all
        // call rather than relying on a prior download step having made it.
        let entry = entry("yolov8n", "object-detection");
        let model_path = models_dir.join("catalog").join("yolov8n").join("model.onnx");

        service.generate_model_config(&entry, &model_path, None).await.unwrap();

        assert!(models_dir.join("catalog").join("yolov8n").join("config.json").exists());
    }
}