torsh-cli 0.1.2

//! Configuration management for ToRSh CLI

// Framework infrastructure - components designed for future use
#![allow(dead_code)]
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use tracing::{debug, info};

/// CLI configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
    /// General settings
    pub general: GeneralConfig,

    /// Model operations settings
    pub model: ModelConfig,

    /// Training settings
    pub training: TrainingConfig,

    /// Hub settings
    pub hub: HubConfig,

    /// Benchmark settings
    pub benchmark: BenchmarkConfig,

    /// Development settings
    pub dev: DevConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GeneralConfig {
    /// Default output directory
    pub output_dir: PathBuf,

    /// Default cache directory
    pub cache_dir: PathBuf,

    /// Default device (cpu, cuda, cuda:0, etc.)
    pub default_device: String,

    /// Number of worker threads
    pub num_workers: usize,

    /// Memory limit in GB
    pub memory_limit_gb: Option<f64>,

    /// Enable progress bars
    pub show_progress: bool,

    /// Default data type (f32, f16, bf16)
    pub default_dtype: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Default model format for conversion
    pub default_format: String,

    /// Model optimization settings
    pub optimization: OptimizationConfig,

    /// Model validation settings
    pub validation: ValidationConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationConfig {
    /// Enable automatic optimization
    pub auto_optimize: bool,

    /// Quantization settings
    pub quantization: QuantizationConfig,

    /// Pruning settings
    pub pruning: PruningConfig,

    /// Fusion settings
    pub fusion: FusionConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QuantizationConfig {
    /// Enable quantization by default
    pub enabled: bool,

    /// Default quantization method (dynamic, static, qat)
    pub method: String,

    /// Default precision (int8, int4, mixed)
    pub precision: String,

    /// Calibration dataset size
    pub calibration_samples: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PruningConfig {
    /// Enable pruning by default
    pub enabled: bool,

    /// Default sparsity target (0.0-1.0)
    pub sparsity: f64,

    /// Pruning method (magnitude, gradient, structured)
    pub method: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FusionConfig {
    /// Enable operator fusion
    pub enabled: bool,

    /// Fusion patterns to apply
    pub patterns: Vec<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidationConfig {
    /// Enable model validation by default
    pub enabled: bool,

    /// Validation dataset path
    pub dataset_path: Option<PathBuf>,

    /// Number of validation samples
    pub num_samples: usize,

    /// Accuracy threshold for validation
    pub accuracy_threshold: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Default training configuration directory
    pub config_dir: PathBuf,

    /// Default checkpoint directory
    pub checkpoint_dir: PathBuf,

    /// Default logging directory
    pub log_dir: PathBuf,

    /// Auto-resume from latest checkpoint
    pub auto_resume: bool,

    /// Save checkpoint every N epochs
    pub checkpoint_frequency: usize,

    /// Early stopping patience
    pub early_stopping_patience: usize,

    /// Mixed precision training
    pub mixed_precision: bool,

    /// Distributed training settings
    pub distributed: DistributedConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DistributedConfig {
    /// Backend for distributed training (nccl, gloo, mpi)
    pub backend: String,

    /// Master address for distributed training
    pub master_addr: String,

    /// Master port for distributed training
    pub master_port: u16,

    /// Auto-detect distributed environment
    pub auto_detect: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HubConfig {
    /// Hub API endpoint
    pub api_endpoint: String,

    /// Authentication token
    pub auth_token: Option<String>,

    /// Default organization
    pub organization: Option<String>,

    /// Model cache directory
    pub cache_dir: PathBuf,

    /// Enable model signature verification
    pub verify_signatures: bool,

    /// Connection timeout in seconds
    pub timeout_seconds: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkConfig {
    /// Default number of warmup iterations
    pub warmup_iterations: usize,

    /// Default number of benchmark iterations
    pub benchmark_iterations: usize,

    /// Default batch sizes to test
    pub batch_sizes: Vec<usize>,

    /// Enable memory tracking
    pub track_memory: bool,

    /// Enable power tracking (if available)
    pub track_power: bool,

    /// Benchmark output directory
    pub output_dir: PathBuf,

    /// Enable detailed profiling
    pub detailed_profiling: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DevConfig {
    /// Enable development mode features
    pub enabled: bool,

    /// Enable debug logging
    pub debug_logging: bool,

    /// Enable experimental features
    pub experimental_features: bool,

    /// Code generation settings
    pub codegen: CodegenConfig,

    /// Testing settings
    pub testing: TestingConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodegenConfig {
    /// Enable automatic code generation
    pub enabled: bool,

    /// Output directory for generated code
    pub output_dir: PathBuf,

    /// Code generation templates directory
    pub templates_dir: PathBuf,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestingConfig {
    /// Enable automatic testing
    pub enabled: bool,

    /// Test data directory
    pub test_data_dir: PathBuf,

    /// Tolerance for numerical tests
    pub numerical_tolerance: f64,
}

impl Default for Config {
    fn default() -> Self {
        let home_dir = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
        let torsh_dir = home_dir.join(".torsh");

        Self {
            general: GeneralConfig {
                output_dir: std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
                cache_dir: torsh_dir.join("cache"),
                default_device: "cpu".to_string(),
                num_workers: std::thread::available_parallelism()
                    .map(|n| n.get())
                    .unwrap_or(4),
                memory_limit_gb: None,
                show_progress: true,
                default_dtype: "f32".to_string(),
            },
            model: ModelConfig {
                default_format: "torsh".to_string(),
                optimization: OptimizationConfig {
                    auto_optimize: false,
                    quantization: QuantizationConfig {
                        enabled: false,
                        method: "dynamic".to_string(),
                        precision: "int8".to_string(),
                        calibration_samples: 1000,
                    },
                    pruning: PruningConfig {
                        enabled: false,
                        sparsity: 0.5,
                        method: "magnitude".to_string(),
                    },
                    fusion: FusionConfig {
                        enabled: true,
                        patterns: vec!["conv_bn_relu".to_string(), "linear_relu".to_string()],
                    },
                },
                validation: ValidationConfig {
                    enabled: true,
                    dataset_path: None,
                    num_samples: 1000,
                    accuracy_threshold: 0.95,
                },
            },
            training: TrainingConfig {
                config_dir: torsh_dir.join("configs"),
                checkpoint_dir: PathBuf::from("./checkpoints"),
                log_dir: PathBuf::from("./logs"),
                auto_resume: false,
                checkpoint_frequency: 1,
                early_stopping_patience: 10,
                mixed_precision: true,
                distributed: DistributedConfig {
                    backend: "nccl".to_string(),
                    master_addr: "localhost".to_string(),
                    master_port: 29500,
                    auto_detect: true,
                },
            },
            hub: HubConfig {
                api_endpoint: "https://hub.torsh.dev".to_string(),
                auth_token: None,
                organization: None,
                cache_dir: torsh_dir.join("hub"),
                verify_signatures: true,
                timeout_seconds: 300,
            },
            benchmark: BenchmarkConfig {
                warmup_iterations: 10,
                benchmark_iterations: 100,
                batch_sizes: vec![1, 4, 8, 16, 32, 64],
                track_memory: true,
                track_power: false,
                output_dir: PathBuf::from("./benchmarks"),
                detailed_profiling: false,
            },
            dev: DevConfig {
                enabled: false,
                debug_logging: false,
                experimental_features: false,
                codegen: CodegenConfig {
                    enabled: false,
                    output_dir: PathBuf::from("./generated"),
                    templates_dir: torsh_dir.join("templates"),
                },
                testing: TestingConfig {
                    enabled: true,
                    test_data_dir: PathBuf::from("./test_data"),
                    numerical_tolerance: 1e-6,
                },
            },
        }
    }
}

/// Load configuration from file or create default
pub async fn load_config(config_path: Option<&Path>) -> Result<Config> {
    let config_path = if let Some(path) = config_path {
        path.to_path_buf()
    } else {
        get_default_config_path()?
    };

    if config_path.exists() {
        debug!("Loading configuration from: {}", config_path.display());
        load_config_from_file(&config_path).await
    } else {
        info!("Configuration file not found, using defaults");
        let config = Config::default();

        // Create config directory if it doesn't exist
        if let Some(parent) = config_path.parent() {
            tokio::fs::create_dir_all(parent).await.with_context(|| {
                format!("Failed to create config directory: {}", parent.display())
            })?;
        }

        // Save default configuration
        save_config(&config, &config_path)
            .await
            .with_context(|| "Failed to save default configuration")?;

        Ok(config)
    }
}

/// Load configuration from a specific file
async fn load_config_from_file(path: &Path) -> Result<Config> {
    let content = tokio::fs::read_to_string(path)
        .await
        .with_context(|| format!("Failed to read config file: {}", path.display()))?;

    let config =
        match path.extension().and_then(|ext| ext.to_str()) {
            Some("yaml") | Some("yml") => serde_yaml::from_str(&content)
                .with_context(|| "Failed to parse YAML configuration")?,
            Some("json") => serde_json::from_str(&content)
                .with_context(|| "Failed to parse JSON configuration")?,
            Some("toml") => {
                toml::from_str(&content).with_context(|| "Failed to parse TOML configuration")?
            }
            _ => {
                // Try to detect format
                if content.trim_start().starts_with('{') {
                    serde_json::from_str(&content)
                        .with_context(|| "Failed to parse JSON configuration")?
                } else {
                    serde_yaml::from_str(&content)
                        .with_context(|| "Failed to parse YAML configuration")?
                }
            }
        };

    Ok(config)
}

/// Save configuration to file
pub async fn save_config(config: &Config, path: &Path) -> Result<()> {
    let content = match path.extension().and_then(|ext| ext.to_str()) {
        Some("json") => serde_json::to_string_pretty(config)
            .with_context(|| "Failed to serialize configuration to JSON")?,
        Some("toml") => toml::to_string_pretty(config)
            .with_context(|| "Failed to serialize configuration to TOML")?,
        _ => {
            // Default to YAML
            serde_yaml::to_string(config)
                .with_context(|| "Failed to serialize configuration to YAML")?
        }
    };

    tokio::fs::write(path, content)
        .await
        .with_context(|| format!("Failed to write config file: {}", path.display()))?;

    info!("Configuration saved to: {}", path.display());
    Ok(())
}

/// Get the default configuration file path
fn get_default_config_path() -> Result<PathBuf> {
    let config_dir = dirs::config_dir()
        .or_else(|| dirs::home_dir().map(|h| h.join(".config")))
        .unwrap_or_else(|| PathBuf::from("."));

    Ok(config_dir.join("torsh").join("config.yaml"))
}

/// Initialize configuration directories
pub async fn init_config_dirs(config: &Config) -> Result<()> {
    let dirs = [
        &config.general.cache_dir,
        &config.training.config_dir,
        &config.training.checkpoint_dir,
        &config.training.log_dir,
        &config.hub.cache_dir,
        &config.benchmark.output_dir,
    ];

    for dir in dirs {
        if !dir.exists() {
            tokio::fs::create_dir_all(dir)
                .await
                .with_context(|| format!("Failed to create directory: {}", dir.display()))?;
            debug!("Created directory: {}", dir.display());
        }
    }

    Ok(())
}

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

    #[tokio::test]
    async fn test_default_config() {
        let config = Config::default();
        assert_eq!(config.general.default_device, "cpu");
        assert_eq!(config.model.default_format, "torsh");
    }

    #[tokio::test]
    async fn test_config_serialization() {
        let config = Config::default();

        // Test YAML serialization
        let yaml = serde_yaml::to_string(&config).unwrap();
        let parsed: Config = serde_yaml::from_str(&yaml).unwrap();
        assert_eq!(config.general.default_device, parsed.general.default_device);

        // Test JSON serialization
        let json = serde_json::to_string_pretty(&config).unwrap();
        let parsed: Config = serde_json::from_str(&json).unwrap();
        assert_eq!(config.general.default_device, parsed.general.default_device);
    }

    #[tokio::test]
    async fn test_config_file_operations() {
        let temp_dir = tempdir().unwrap();
        let config_path = temp_dir.path().join("test_config.yaml");

        let config = Config::default();

        // Save configuration
        save_config(&config, &config_path).await.unwrap();
        assert!(config_path.exists());

        // Load configuration
        let loaded_config = load_config_from_file(&config_path).await.unwrap();
        assert_eq!(
            config.general.default_device,
            loaded_config.general.default_device
        );
    }
}