modelexpress-client 0.7.0

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::args::{DownloadStrategy, ModelCommands, OutputFormat};
use super::output::{print_human_readable, print_output};
use super::payload::read_payload;
use colored::*;
use modelexpress_client::{Client, ClientConfig, ModelDownloadResult, ModelProvider};
use modelexpress_common::{
    cache::{CacheConfig, CacheStats, ModelInfo, resolve_model_path},
    download,
};
use serde_json::Value;
use std::borrow::Cow;
use std::fs;
use std::io::Write;
use std::path::{Path, PathBuf};
use tracing::{debug, error, info, warn};

fn format_model_line(stats: &CacheStats, model: &ModelInfo, detailed: bool) -> String {
    if detailed {
        format!(
            "  [{}] {} ({}) - {:?}",
            model.provider,
            model.name,
            stats.format_model_size(model),
            model.path
        )
    } else {
        format!(
            "  [{}] {} ({})",
            model.provider,
            model.name,
            stats.format_model_size(model)
        )
    }
}

fn model_json(stats: &CacheStats, model: &ModelInfo, detailed: bool) -> serde_json::Value {
    if detailed {
        serde_json::json!({
            "provider": model.provider.to_string(),
            "name": model.name,
            "size": model.size,
            "formatted_size": stats.format_model_size(model),
            "path": model.path
        })
    } else {
        serde_json::json!({
            "provider": model.provider.to_string(),
            "name": model.name,
            "size": model.size,
            "formatted_size": stats.format_model_size(model)
        })
    }
}

fn strip_ascii_prefix_ignore_case<'a>(value: &'a str, prefix: &str) -> Option<&'a str> {
    let candidate = value.get(..prefix.len())?;
    if candidate.eq_ignore_ascii_case(prefix) {
        value.get(prefix.len()..)
    } else {
        None
    }
}

fn normalize_model_name_scheme(model_name: &str) -> Cow<'_, str> {
    let model_name = model_name.trim_start();
    if let Some(rest) = strip_ascii_prefix_ignore_case(model_name, "gs://") {
        Cow::Owned(format!("gs://{rest}"))
    } else if let Some(rest) = strip_ascii_prefix_ignore_case(model_name, "s3://") {
        Cow::Owned(format!("s3://{rest}"))
    } else if let Some(rest) = strip_ascii_prefix_ignore_case(model_name, "ngc://") {
        Cow::Owned(format!("ngc://{rest}"))
    } else {
        Cow::Borrowed(model_name)
    }
}

fn resolve_validation_model_path(cache_root: &Path, model_name: &str) -> (ModelProvider, PathBuf) {
    let normalized_name = normalize_model_name_scheme(model_name);
    let provider = ModelProvider::resolve_provider_for_model_name(
        normalized_name.as_ref(),
        ModelProvider::HuggingFace,
    );
    let model_path = resolve_model_path(cache_root, provider, normalized_name.as_ref(), None)
        .unwrap_or_else(|_| cache_root.join(normalized_name.as_ref()));
    (provider, model_path)
}

fn resolve_download_options(
    model_name: &str,
    default_provider: ModelProvider,
) -> (ModelProvider, bool) {
    let provider = ModelProvider::resolve_provider_for_model_name(model_name, default_provider);
    (provider, provider == ModelProvider::S3)
}

fn has_huggingface_weights(model_path: &Path) -> bool {
    if model_path.join("pytorch_model.bin").exists() {
        return true;
    }

    fs::read_dir(model_path).is_ok_and(|entries| {
        entries.flatten().any(|entry| {
            entry.file_name().to_str().is_some_and(|name| {
                name.ends_with(".safetensors") || name == "model.safetensors.index.json"
            })
        })
    })
}

/// Handle the health check command
pub async fn handle_health_command(
    config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    debug!(
        "Initiating health check to server: {}",
        config.connection.endpoint
    );

    let mut client = Client::new(config).await?;
    let status = client.health_check().await?;

    info!("Health check completed successfully");

    match format {
        OutputFormat::Human => {
            println!("{}", "Server Health Status".green().bold());
            println!(
                "  {}: {}",
                "Status".cyan().bold(),
                if status.status == "healthy" || status.status == "ok" {
                    status.status.green()
                } else {
                    status.status.red()
                }
            );
            println!("  {}: {}", "Version".cyan().bold(), status.version);
            println!("  {}: {} seconds", "Uptime".cyan().bold(), status.uptime);
        }
        _ => print_output(&status, format),
    }

    Ok(())
}

/// Handle model commands (unified model management)
pub async fn handle_model_command(
    command: ModelCommands,
    storage_path_override: Option<PathBuf>,
    server_config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    match command {
        ModelCommands::Download {
            model_name,
            provider,
            strategy,
            revision,
        } => {
            download_model(
                storage_path_override,
                model_name,
                provider,
                strategy,
                revision,
                server_config,
                format,
            )
            .await
        }
        ModelCommands::Init {
            storage_path,
            server_endpoint,
        } => init_model_storage(storage_path, server_endpoint, format).await,
        ModelCommands::List { detailed } => {
            list_models(storage_path_override, detailed, format).await
        }
        ModelCommands::Status => show_model_status(storage_path_override, format).await,
        ModelCommands::Clear {
            provider,
            model_name,
        } => {
            clear_model(
                storage_path_override,
                provider,
                &model_name,
                server_config,
                format,
            )
            .await
        }
        ModelCommands::ClearAll { yes } => {
            clear_all_models(storage_path_override, yes, server_config, format).await
        }
        ModelCommands::Validate { model_name } => {
            validate_models(storage_path_override, model_name, format).await
        }
        ModelCommands::Stats { detailed } => {
            show_model_stats(storage_path_override, detailed, format).await
        }
    }
}

async fn download_model(
    storage_path_override: Option<PathBuf>,
    model_name: String,
    provider: ModelProvider,
    strategy: DownloadStrategy,
    revision: Option<String>,
    config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let (provider, ignore_weights) = resolve_download_options(&model_name, provider);
    debug!(
        "Starting model download: {} with provider {:?} and strategy {:?}",
        model_name, provider, strategy
    );

    if let OutputFormat::Human = format {
        println!("{}", "Model Download".green().bold());
        println!("  {}: {}", "Model".cyan().bold(), model_name);
        println!("  {}: {}", "Provider".cyan().bold(), provider);
        println!("  {}: {:?}", "Strategy".cyan().bold(), strategy);
        if let Some(revision) = &revision {
            println!("  {}: {}", "Revision".cyan().bold(), revision);
        }
        println!();
    }

    info!("Downloading model: {}", model_name);

    // Get cache config if available, applying settings from ClientConfig
    let cache_config = if let Some(path) = storage_path_override {
        Some(CacheConfig::from_path(path)?)
    } else {
        CacheConfig::discover().ok()
    };

    // Apply shared_storage and transfer_chunk_size settings from ClientConfig
    let cache_config = cache_config.map(|mut c| {
        c.shared_storage = config.cache.shared_storage;
        c.transfer_chunk_size = config.cache.transfer_chunk_size;
        c.server_endpoint = config.connection.endpoint.clone();
        c
    });

    let revision = revision.as_deref();
    let result = match strategy {
        DownloadStrategy::SmartFallback => {
            debug!("Using smart fallback strategy");
            let mut config = config.clone();
            if let Some(cache_config) = cache_config {
                config.cache = cache_config;
            }
            Client::request_model_with_smart_fallback_revision(
                model_name.clone(),
                provider,
                config,
                ignore_weights,
                revision,
            )
            .await
        }
        DownloadStrategy::ServerOnly => {
            debug!("Using server-only strategy");
            let mut client = if let Some(cache_config) = cache_config {
                Client::new_with_cache(config.clone(), cache_config).await?
            } else {
                Client::new(config.clone()).await?
            };
            client
                .request_model_revision(&model_name, provider, ignore_weights, revision)
                .await
        }
        DownloadStrategy::Direct => {
            debug!("Using direct download strategy");
            download::download_model_revision(
                &model_name,
                provider,
                cache_config.map(|config| config.local_path),
                ignore_weights,
                revision,
            )
            .await
            .map(ModelDownloadResult::from)
            .map_err(|e| {
                modelexpress_common::Error::Generic(format!("Direct download failed: {e}")).into()
            })
        }
    };

    match result {
        Ok(outcome) => {
            info!("Model download completed successfully: {}", model_name);
            let success_msg = format!("Model '{model_name}' downloaded successfully");
            match format {
                OutputFormat::Human => {
                    println!("{}", "✅ SUCCESS".green().bold());
                    println!("  {success_msg}");
                    if let Some(resolved) = &outcome.resolved_revision {
                        println!("  {}: {}", "Resolved revision".cyan().bold(), resolved);
                    }
                    if let Some(path) = &outcome.path {
                        println!("  {}: {}", "Path".cyan().bold(), path.display());
                    }
                }
                _ => {
                    let output = serde_json::json!({
                        "success": true,
                        "message": success_msg,
                        "model_name": model_name,
                        "provider": provider.to_string(),
                        "strategy": format!("{:?}", strategy),
                        "resolved_revision": outcome.resolved_revision,
                        "path": outcome.path.as_ref().map(|path| path.display().to_string())
                    });
                    print_output(&output, format);
                }
            }
        }
        Err(e) => {
            error!("Model download failed for {}: {}", model_name, e);
            let error_msg = format!("Failed to download model '{model_name}': {e}");
            match format {
                OutputFormat::Human => {
                    println!("{}", "❌ FAILED".red().bold());
                    println!("  {error_msg}");
                }
                _ => {
                    let output = serde_json::json!({
                        "success": false,
                        "error": error_msg,
                        "model_name": model_name,
                        "provider": provider.to_string(),
                        "strategy": format!("{:?}", strategy)
                    });
                    print_output(&output, format);
                }
            }
            return Err(e);
        }
    }

    Ok(())
}

/// Handle API send command
pub async fn handle_api_send(
    action: String,
    payload: Option<String>,
    payload_file: Option<String>,
    config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    debug!("Preparing API request for action: {}", action);

    let mut client = Client::new(config).await?;
    let payload_data = read_payload(payload, payload_file)?;

    if payload_data.is_some() {
        debug!("API request includes payload data");
    }

    if let OutputFormat::Human = format {
        println!("{}", "API Request".green().bold());
        println!("  {}: {}", "Action".cyan().bold(), action);
        if payload_data.is_some() {
            println!("  {}: Yes", "Payload".cyan().bold());
        }
        println!();
    }

    info!("Sending API request: {}", action);

    let response: Value = client.send_request(&action, payload_data).await?;

    info!("API request completed successfully");

    match format {
        OutputFormat::Human => {
            println!("{}", "Response:".green().bold());
            print_human_readable(&response);
        }
        _ => print_output(&response, format),
    }

    Ok(())
}

async fn init_model_storage(
    storage_path: Option<PathBuf>,
    server_endpoint: Option<String>,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let config = if let Some(path) = storage_path {
        CacheConfig::from_path(path)?
    } else {
        // Use default configuration instead of prompting
        CacheConfig::default()
    };

    // Override with command line options if provided
    let mut config = config;
    if let Some(endpoint) = server_endpoint {
        config.server_endpoint = endpoint;
    }

    // Save configuration
    config.save_to_config_file()?;

    match format {
        OutputFormat::Human => {
            println!("{}", "ModelExpress Storage Configuration".green().bold());
            println!("{}", "===================================".green().bold());
            println!("Configuration saved successfully!");
            println!("Storage path: {:?}", config.local_path);
            println!("Server endpoint: {}", config.server_endpoint);
        }
        _ => {
            let output = serde_json::json!({
                "success": true,
                "storage_path": config.local_path,
                "server_endpoint": config.server_endpoint,
            });
            print_output(&output, format);
        }
    }

    Ok(())
}

async fn list_models(
    storage_path_override: Option<PathBuf>,
    detailed: bool,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;
    let stats = storage_config.get_cache_stats()?;

    match format {
        OutputFormat::Human => {
            println!("{}", "Downloaded Models".green().bold());
            println!("{}", "=================".green().bold());
            println!("Total models: {}", stats.total_models);
            println!("Total size: {}", stats.format_total_size());

            if stats.models.is_empty() {
                println!("No models found in storage.");
                return Ok(());
            }

            println!("Models:");
            for model in &stats.models {
                println!("{}", format_model_line(&stats, model, detailed));
            }
        }
        _ => {
            let models_json: Vec<serde_json::Value> = stats
                .models
                .iter()
                .map(|model| model_json(&stats, model, detailed))
                .collect();

            let output = serde_json::json!({
                "total_models": stats.total_models,
                "total_size": stats.format_total_size(),
                "models": models_json
            });
            print_output(&output, format);
        }
    }

    Ok(())
}

async fn show_model_status(
    storage_path_override: Option<PathBuf>,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;
    let stats = storage_config.get_cache_stats()?;

    let storage_accessible = storage_config.local_path.exists();
    let server_available = Client::new(ClientConfig::for_testing(&storage_config.server_endpoint))
        .await
        .is_ok();

    match format {
        OutputFormat::Human => {
            println!("{}", "Model Storage Status".green().bold());
            println!("{}", "====================".green().bold());
            println!("Storage path: {:?}", storage_config.local_path);
            println!("Server endpoint: {}", storage_config.server_endpoint);
            println!("Total models: {}", stats.total_models);
            println!("Total size: {}", stats.format_total_size());

            // Check if storage directory exists and is accessible
            if storage_accessible {
                println!("Storage directory: ✅ Accessible");
            } else {
                println!("Storage directory: ❌ Not found");
            }

            // Try to connect to server
            if server_available {
                println!("Server connection: ✅ Available");
            } else {
                println!("Server connection: ❌ Unavailable");
            }
        }
        _ => {
            let output = serde_json::json!({
                "storage_path": storage_config.local_path,
                "server_endpoint": storage_config.server_endpoint,
                "total_models": stats.total_models,
                "total_size": stats.format_total_size(),
                "storage_accessible": storage_accessible,
                "server_available": server_available
            });
            print_output(&output, format);
        }
    }

    Ok(())
}

/// Best-effort deletion of a model's server-side registry record. Removing the local
/// files is the primary intent of `model clear`, so a server that is unreachable or
/// returns an error is surfaced as a warning rather than failing the command. Returns
/// `Ok(())` when the record was deleted, `Err(reason)` otherwise.
async fn delete_model_registry_record(
    server_config: &ClientConfig,
    model_name: &str,
    provider: ModelProvider,
) -> Result<(), String> {
    let mut client = Client::new(server_config.clone())
        .await
        .map_err(|e| format!("could not connect to server: {e}"))?;
    client
        .delete_model_on_server(model_name, provider)
        .await
        .map_err(|e| format!("server registry delete failed: {e}"))
}

async fn clear_model(
    storage_path_override: Option<PathBuf>,
    provider: ModelProvider,
    model_name: &str,
    server_config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;

    storage_config.clear_model(model_name, provider)?;

    // Also drop the server-side registry record so a later `model download` re-fetches
    // the files instead of hitting a stale DOWNLOADED record and returning a false success.
    let registry_warning =
        match delete_model_registry_record(&server_config, model_name, provider).await {
            Ok(()) => None,
            Err(reason) => {
                warn!("Cleared local files for '{model_name}' but {reason}");
                Some(reason)
            }
        };

    match format {
        OutputFormat::Human => {
            println!(
                "✅ Model '{model_name}' cleared from storage for provider {}",
                provider
            );
            match &registry_warning {
                None => println!("   Server registry record removed"),
                Some(reason) => println!(
                    "{}",
                    format!("⚠️  Server registry record NOT removed: {reason}").yellow()
                ),
            }
        }
        _ => {
            let output = serde_json::json!({
                "success": true,
                "message": format!("Model '{}' cleared from storage", model_name),
                "model_name": model_name,
                "provider": provider.to_string(),
                "registry_cleared": registry_warning.is_none(),
                "registry_warning": registry_warning,
            });
            print_output(&output, format);
        }
    }

    Ok(())
}

async fn clear_all_models(
    storage_path_override: Option<PathBuf>,
    yes: bool,
    server_config: ClientConfig,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;

    if !yes && matches!(format, OutputFormat::Human) {
        print!("Are you sure you want to clear all models from storage? [y/N]: ");
        std::io::stdout().flush()?;

        let mut input = String::new();
        std::io::stdin().read_line(&mut input)?;

        if input.trim().to_lowercase() != "y" {
            println!("Operation cancelled.");
            return Ok(());
        }
    }

    // Capture the cached models before clearing so we can drop their registry records too.
    let cached = storage_config
        .get_cache_stats()
        .map(|stats| {
            stats
                .models
                .into_iter()
                .map(|model| (model.name, model.provider))
                .collect::<Vec<_>>()
        })
        .unwrap_or_default();

    storage_config.clear_all()?;

    let mut registry_failures = 0usize;
    for (name, provider) in &cached {
        if let Err(reason) = delete_model_registry_record(&server_config, name, *provider).await {
            warn!("Cleared local files for '{name}' but {reason}");
            registry_failures = registry_failures.saturating_add(1);
        }
    }

    match format {
        OutputFormat::Human => {
            println!("✅ All models cleared from storage");
            if registry_failures == 0 {
                println!("   Server registry records removed");
            } else {
                println!(
                    "{}",
                    format!(
                        "⚠️  {registry_failures} server registry record(s) NOT removed (see warnings)"
                    )
                    .yellow()
                );
            }
        }
        _ => {
            let output = serde_json::json!({
                "success": true,
                "message": "All models cleared from storage",
                "registry_records_cleared": cached.len().saturating_sub(registry_failures),
                "registry_records_failed": registry_failures,
            });
            print_output(&output, format);
        }
    }

    Ok(())
}

async fn validate_models(
    storage_path_override: Option<PathBuf>,
    model_name: Option<String>,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;

    if let Some(name) = model_name {
        let (provider, model_path) =
            resolve_validation_model_path(&storage_config.local_path, &name);
        let exists = model_path.exists();

        match format {
            OutputFormat::Human => {
                println!("{}", "Model Validation".green().bold());
                if exists {
                    println!("✅ Model '{name}' found in storage");

                    if provider == ModelProvider::HuggingFace {
                        // Check for common HuggingFace model files.
                        let required_files = ["config.json", "tokenizer.json"];
                        for file in &required_files {
                            let file_path = model_path.join(file);
                            if file_path.exists() {
                                debug!("  ✅ {} found", file);
                            } else {
                                println!("  ⚠️  {file} missing");
                            }
                        }

                        if !has_huggingface_weights(&model_path) {
                            println!(
                                "  ⚠️  missing model weights (expected pytorch_model.bin or safetensors)"
                            );
                        }
                    }
                } else {
                    println!("❌ Model '{name}' not found in storage");
                }
            }
            _ => {
                let output = serde_json::json!({
                    "model_name": name,
                    "exists": exists,
                    "path": model_path
                });
                print_output(&output, format);
            }
        }
    } else {
        // Validate entire storage
        let stats = storage_config.get_cache_stats()?;

        match format {
            OutputFormat::Human => {
                println!("{}", "Model Validation".green().bold());
                println!("Found {} models in storage", stats.total_models);

                for model in &stats.models {
                    println!("{}", format_model_line(&stats, model, false));
                }
            }
            _ => {
                let output = serde_json::json!({
                    "total_models": stats.total_models,
                    "models": stats.models.iter().map(|model| {
                        model_json(&stats, model, false)
                    }).collect::<Vec<_>>()
                });
                print_output(&output, format);
            }
        }
    }

    Ok(())
}

async fn show_model_stats(
    storage_path_override: Option<PathBuf>,
    detailed: bool,
    format: &OutputFormat,
) -> Result<(), Box<dyn std::error::Error>> {
    let storage_config = get_storage_config(storage_path_override)?;
    let stats = storage_config.get_cache_stats()?;

    match format {
        OutputFormat::Human => {
            println!("{}", "Model Storage Statistics".green().bold());
            println!("{}", "========================".green().bold());
            println!("Total models: {}", stats.total_models);
            println!("Total size: {}", stats.format_total_size());

            if detailed && !stats.models.is_empty() {
                println!("Detailed Statistics:");
                for model in &stats.models {
                    println!(
                        "  [{}] {}: {} bytes ({})",
                        model.provider,
                        model.name,
                        model.size,
                        stats.format_model_size(model)
                    );
                }
            }
        }
        _ => {
            let models_data = if detailed {
                Some(
                    stats
                        .models
                        .iter()
                        .map(|model| model_json(&stats, model, false))
                        .collect::<Vec<_>>(),
                )
            } else {
                None
            };

            let mut output = serde_json::json!({
                "total_models": stats.total_models,
                "total_size": stats.format_total_size()
            });

            if let Some(models) = models_data {
                output["detailed_models"] = serde_json::Value::Array(models);
            }

            print_output(&output, format);
        }
    }

    Ok(())
}

fn get_storage_config(
    storage_path_override: Option<PathBuf>,
) -> Result<CacheConfig, Box<dyn std::error::Error>> {
    // If storage path is provided via CLI, use it
    if let Some(path) = storage_path_override {
        return Ok(CacheConfig::from_path(path)?);
    }

    // Otherwise, try to discover configuration
    CacheConfig::discover()
        .map_err(|e| format!("Failed to discover storage configuration: {e}").into())
}

#[cfg(test)]
#[allow(clippy::expect_used)]
mod tests {
    use super::*;
    use std::fs;
    use tempfile::TempDir;

    #[test]
    fn test_validate_resolves_gcs_model_name_to_gcs_cache_path() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let model_name = "GS://bucket/foo/bar";
        let expected_path = temp_dir
            .path()
            .join("gcs")
            .join("bucket")
            .join("foo")
            .join("bar");
        fs::create_dir_all(&expected_path).expect("Failed to create GCS cache path");

        let (provider, model_path) = resolve_validation_model_path(temp_dir.path(), model_name);

        assert_eq!(provider, ModelProvider::Gcs);
        assert_eq!(model_path, expected_path);
        assert!(model_path.exists());
        assert!(!temp_dir.path().join("GS://bucket/foo/bar").exists());
    }

    #[test]
    fn test_validate_resolves_s3_model_name_to_s3_cache_path() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        let model_name = "S3://bucket/foo/bar";
        let expected_path = temp_dir
            .path()
            .join("s3")
            .join("bucket")
            .join("foo")
            .join("bar");
        fs::create_dir_all(&expected_path).expect("Failed to create S3 cache path");

        let (provider, model_path) = resolve_validation_model_path(temp_dir.path(), model_name);

        assert_eq!(provider, ModelProvider::S3);
        assert_eq!(model_path, expected_path);
        assert!(model_path.exists());
        assert!(!temp_dir.path().join("S3://bucket/foo/bar").exists());
    }

    #[test]
    fn test_validate_provider_resolution_keeps_ambiguous_names_hugging_face() {
        let (provider, _) =
            resolve_validation_model_path(Path::new("/tmp/mx-cache"), "nvidia/model/version");
        assert_eq!(provider, ModelProvider::HuggingFace);

        let (provider, _) =
            resolve_validation_model_path(Path::new("/tmp/mx-cache"), "ngc://nvidia/model/version");
        assert_eq!(provider, ModelProvider::Ngc);

        let (provider, _) =
            resolve_validation_model_path(Path::new("/tmp/mx-cache"), "NgC://nvidia/model/version");
        assert_eq!(provider, ModelProvider::Ngc);
    }

    #[test]
    fn test_s3_download_options_enable_metadata_only_mode() {
        let (provider, ignore_weights) =
            resolve_download_options("s3://bucket/org/model", ModelProvider::HuggingFace);
        assert_eq!(provider, ModelProvider::S3);
        assert!(ignore_weights);

        let (provider, ignore_weights) =
            resolve_download_options("org/model", ModelProvider::HuggingFace);
        assert_eq!(provider, ModelProvider::HuggingFace);
        assert!(!ignore_weights);
    }

    #[test]
    fn test_huggingface_weights_accepts_safetensors() {
        let temp_dir = TempDir::new().expect("Failed to create temp dir");
        fs::write(
            temp_dir.path().join("model-00001-of-00002.safetensors"),
            b"",
        )
        .expect("Failed to write safetensors shard");

        assert!(has_huggingface_weights(temp_dir.path()));
    }
}