axonml-cli 0.5.0

//! Hub CLI Command
//!
//! # File
//! `crates/axonml-cli/src/commands/hub.rs`
//!
//! # Author
//! Andrew Jewell Sr - AutomataNexus
//!
//! # Updated
//! March 8, 2026
//!
//! # Disclaimer
//! Use at own risk. This software is provided "as is", without warranty of any
//! kind, express or implied. The author and AutomataNexus shall not be held
//! liable for any damages arising from the use of this software.

use std::fs;
use std::path::PathBuf;

use colored::Colorize;

// =============================================================================
// Pretrained Model Information
// =============================================================================

/// Information about a pretrained model.
#[derive(Debug, Clone)]
pub struct ModelInfo {
    pub name: String,
    pub description: String,
    pub size_mb: f64,
    pub accuracy: f32,
    pub dataset: String,
    pub input_size: (usize, usize),
    pub num_classes: usize,
}

/// Get available pretrained models.
pub fn available_models() -> Vec<ModelInfo> {
    vec![
        ModelInfo {
            name: "resnet18".to_string(),
            description: "ResNet-18 (18 layers, ~11M params)".to_string(),
            size_mb: 44.7,
            accuracy: 69.76,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
        ModelInfo {
            name: "resnet34".to_string(),
            description: "ResNet-34 (34 layers, ~21M params)".to_string(),
            size_mb: 83.3,
            accuracy: 73.31,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
        ModelInfo {
            name: "resnet50".to_string(),
            description: "ResNet-50 (50 layers, ~23M params)".to_string(),
            size_mb: 97.8,
            accuracy: 76.13,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
        ModelInfo {
            name: "vgg16".to_string(),
            description: "VGG-16 (16 layers, ~138M params)".to_string(),
            size_mb: 528.0,
            accuracy: 71.59,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
        ModelInfo {
            name: "vgg19".to_string(),
            description: "VGG-19 (19 layers, ~144M params)".to_string(),
            size_mb: 548.0,
            accuracy: 72.38,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
        ModelInfo {
            name: "vgg16_bn".to_string(),
            description: "VGG-16 with BatchNorm (~138M params)".to_string(),
            size_mb: 528.0,
            accuracy: 73.36,
            dataset: "ImageNet-1K".to_string(),
            input_size: (224, 224),
            num_classes: 1000,
        },
    ]
}

/// Get the cache directory for pretrained weights.
pub fn cache_dir() -> PathBuf {
    let base = dirs::cache_dir()
        .or_else(dirs::home_dir)
        .unwrap_or_else(|| PathBuf::from("."));
    base.join("axonml").join("hub").join("weights")
}

/// Check if a model is cached.
pub fn is_cached(model_name: &str) -> bool {
    cache_dir()
        .join(format!("{}.safetensors", model_name))
        .exists()
}

/// Get cached model path.
pub fn cached_path(model_name: &str) -> PathBuf {
    cache_dir().join(format!("{}.safetensors", model_name))
}

// =============================================================================
// CLI Commands
// =============================================================================

/// Execute hub list command.
pub fn execute_list() -> Result<(), String> {
    let models = available_models();

    println!("{}", "Available Pretrained Models".bold());
    println!("{}", "═".repeat(80));
    println!(
        "{:<12} {:<40} {:>8} {:>8} {}",
        "NAME".bold(),
        "DESCRIPTION".bold(),
        "SIZE".bold(),
        "ACC".bold(),
        "CACHED".bold()
    );
    println!("{}", "─".repeat(80));

    for model in &models {
        let cached = if is_cached(&model.name) {
            "✓".green().to_string()
        } else {
            String::new()
        };

        println!(
            "{:<12} {:<40} {:>6.1} MB {:>6.1}% {}",
            model.name.cyan(),
            model.description,
            model.size_mb,
            model.accuracy,
            cached
        );
    }

    println!("{}", "─".repeat(80));
    println!(
        "Download with: {} <model-name>",
        "axonml hub download".cyan()
    );
    println!("Cache directory: {:?}", cache_dir());

    Ok(())
}

/// Execute hub info command.
pub fn execute_info(model_name: &str) -> Result<(), String> {
    let models = available_models();
    let model = models
        .iter()
        .find(|m| m.name == model_name)
        .ok_or_else(|| format!("Model '{}' not found", model_name))?;

    println!("{}", format!("Model: {}", model.name).bold());
    println!("{}", "─".repeat(50));
    println!("Description:  {}", model.description);
    println!("Dataset:      {}", model.dataset);
    println!("Accuracy:     {:.2}%", model.accuracy);
    println!("Size:         {:.1} MB", model.size_mb);
    println!(
        "Input Size:   {}x{}",
        model.input_size.0, model.input_size.1
    );
    println!("Classes:      {}", model.num_classes);
    println!(
        "Cached:       {}",
        if is_cached(&model.name) { "Yes" } else { "No" }
    );

    if is_cached(&model.name) {
        println!("Path:         {:?}", cached_path(&model.name));
    }

    println!("{}", "─".repeat(50));

    // Usage example
    println!("\n{}", "Usage Example:".bold());
    println!(
        r#"
use axonml_vision::{{hub, models::ResNet}};

// Download weights (if not cached)
let weights_path = hub::download_weights("{}", false)?;

// Load model with pretrained weights
let state_dict = hub::load_state_dict(&weights_path)?;
let model = ResNet::resnet18(1000);
// model.load_state_dict(&state_dict);  // Apply weights
"#,
        model.name
    );

    Ok(())
}

/// Execute hub download command.
pub fn execute_download(model_name: &str, force: bool) -> Result<(), String> {
    let models = available_models();
    let model = models
        .iter()
        .find(|m| m.name == model_name)
        .ok_or_else(|| {
            format!(
                "Model '{}' not found. Run 'axonml hub list' to see available models.",
                model_name
            )
        })?;

    if is_cached(model_name) && !force {
        println!(
            "{} Model '{}' is already cached at {:?}",
            "✓".green(),
            model_name,
            cached_path(model_name)
        );
        println!("Use --force to re-download.");
        return Ok(());
    }

    println!(
        "{} Downloading {} ({:.1} MB)...",
        "⬇".cyan(),
        model_name,
        model.size_mb
    );

    // Ensure cache directory exists
    let cache = cache_dir();
    fs::create_dir_all(&cache).map_err(|e| e.to_string())?;

    // For now, create synthetic weights (real download would require the download feature)
    let weights_path = cached_path(model_name);

    #[cfg(feature = "hub-download")]
    {
        // Real download implementation would go here
        // use reqwest to download from model.url
        println!("Downloading from remote...");
    }

    #[cfg(not(feature = "hub-download"))]
    {
        // Create synthetic weights for demonstration
        println!("Note: Creating synthetic weights (enable hub-download feature for real weights)");
        create_synthetic_weights(model_name, &weights_path)?;
    }

    println!("{} Downloaded to {:?}", "✓".green(), weights_path);

    Ok(())
}

/// Create synthetic weights for testing.
#[cfg(not(feature = "hub-download"))]
fn create_synthetic_weights(model_name: &str, path: &PathBuf) -> Result<(), String> {
    use rand::Rng;
    use std::io::Write;

    let mut rng = rand::thread_rng();

    // Create a simple binary file with random data
    let mut file = fs::File::create(path).map_err(|e| e.to_string())?;

    // Write header
    let num_tensors: u32 = match model_name {
        "resnet18" => 62,
        "resnet34" => 110,
        "resnet50" => 159,
        "vgg16" | "vgg16_bn" => 32,
        "vgg19" => 38,
        _ => 20,
    };

    file.write_all(&num_tensors.to_le_bytes())
        .map_err(|e| e.to_string())?;

    // Write some dummy tensor data
    for i in 0..num_tensors {
        let name = format!("layer_{}", i);
        let name_bytes = name.as_bytes();
        let name_len = name_bytes.len() as u32;

        file.write_all(&name_len.to_le_bytes())
            .map_err(|e| e.to_string())?;
        file.write_all(name_bytes).map_err(|e| e.to_string())?;

        // Shape: [64] for simplicity
        let ndim: u32 = 1;
        file.write_all(&ndim.to_le_bytes())
            .map_err(|e| e.to_string())?;
        let dim: u64 = 64;
        file.write_all(&dim.to_le_bytes())
            .map_err(|e| e.to_string())?;

        // Data
        for _ in 0..64 {
            let val: f32 = rng.r#gen::<f32>() * 0.1;
            file.write_all(&val.to_le_bytes())
                .map_err(|e| e.to_string())?;
        }
    }

    Ok(())
}

/// Execute hub clear command.
pub fn execute_clear(model_name: Option<&str>) -> Result<(), String> {
    let cache = cache_dir();

    if let Some(name) = model_name {
        let path = cached_path(name);
        if path.exists() {
            fs::remove_file(&path).map_err(|e| e.to_string())?;
            println!("{} Removed cached weights for '{}'", "✓".green(), name);
        } else {
            println!("Model '{}' is not cached.", name);
        }
    } else if cache.exists() {
        let count = fs::read_dir(&cache)
            .map_err(|e| e.to_string())?
            .flatten()
            .count();

        fs::remove_dir_all(&cache).map_err(|e| e.to_string())?;
        println!("{} Cleared {} cached model(s)", "✓".green(), count);
    } else {
        println!("Cache is empty.");
    }

    Ok(())
}

/// Execute hub cached command.
pub fn execute_cached() -> Result<(), String> {
    let cache = cache_dir();

    if !cache.exists() {
        println!("No cached models.");
        return Ok(());
    }

    println!("{}", "Cached Models".bold());
    println!("{}", "─".repeat(60));

    let mut total_size: u64 = 0;
    let mut count = 0;

    for entry in fs::read_dir(&cache).map_err(|e| e.to_string())? {
        let entry = entry.map_err(|e| e.to_string())?;
        let path = entry.path();

        if path.is_file() {
            let filename = path.file_name().unwrap().to_string_lossy();
            let metadata = fs::metadata(&path).map_err(|e| e.to_string())?;
            let size = metadata.len();
            total_size += size;

            let model_name = filename.trim_end_matches(".safetensors");
            println!(
                "  {} {} ({:.1} MB)",
                "✓".green(),
                model_name.cyan(),
                size as f64 / 1_000_000.0
            );
            count += 1;
        }
    }

    if count == 0 {
        println!("  No models cached.");
    } else {
        println!("{}", "─".repeat(60));
        println!(
            "Total: {} model(s), {:.1} MB",
            count,
            total_size as f64 / 1_000_000.0
        );
    }

    println!("\nCache directory: {:?}", cache);

    Ok(())
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_available_models() {
        let models = available_models();
        assert!(!models.is_empty());

        let resnet = models.iter().find(|m| m.name == "resnet18");
        assert!(resnet.is_some());
    }

    #[test]
    fn test_cache_dir() {
        let dir = cache_dir();
        assert!(dir.to_string_lossy().contains("axonml"));
    }

    #[test]
    fn test_cached_path() {
        let path = cached_path("resnet18");
        assert!(path.to_string_lossy().contains("resnet18"));
        assert!(path.to_string_lossy().contains("safetensors"));
    }
}