mecha10-models 0.6.2

//! 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<()> {
        let model_dir = self.models_dir.join(&entry.name);
        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(())
    }
}