trustformers-mobile 0.1.1

// Size optimization utilities for mobile deployment
// Focuses on meeting the <50MB framework size target

use crate::Result;
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::collections::HashMap;
use tracing::{debug, info, warn};

/// Framework size measurement results
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FrameworkSizeResults {
    /// Measured framework size in MB
    pub measured_size_mb: f64,
    /// Target achieved
    pub target_achieved: bool,
    /// Size breakdown by component
    pub size_breakdown: HashMap<String, f64>,
}

/// Size optimization strategies for mobile frameworks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SizeOptimizationStrategy {
    /// Minimal feature set for basic inference
    Minimal,
    /// Balanced features with size constraints
    Balanced,
    /// Maximum features with aggressive compression
    MaximumCompressed,
    /// Custom optimization strategy
    Custom {
        enabled_features: Vec<String>,
        compression_level: u8,
    },
}

/// Framework size targets and measurements
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SizeMetrics {
    /// Current framework size in MB
    pub current_size_mb: f64,
    /// Target size in MB
    pub target_size_mb: f64,
    /// Size breakdown by component
    pub component_sizes: HashMap<String, f64>,
    /// Estimated achievable size with optimizations
    pub optimized_size_mb: f64,
    /// Size reduction percentage
    pub reduction_percentage: f64,
}

/// Size optimization configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SizeOptimizerConfig {
    /// Optimization strategy
    pub strategy: SizeOptimizationStrategy,
    /// Target framework size in MB
    pub target_size_mb: f64,
    /// Enable dead code elimination
    pub eliminate_dead_code: bool,
    /// Enable symbol stripping
    pub strip_symbols: bool,
    /// Enable compression
    pub enable_compression: bool,
    /// Compression level (1-9)
    pub compression_level: u8,
    /// Enable link-time optimization
    pub enable_lto: bool,
    /// Disable debugging symbols in release
    pub strip_debug_symbols: bool,
}

impl Default for SizeOptimizerConfig {
    fn default() -> Self {
        Self {
            strategy: SizeOptimizationStrategy::Balanced,
            target_size_mb: 50.0, // Target from TODO.md
            eliminate_dead_code: true,
            strip_symbols: true,
            enable_compression: true,
            compression_level: 6,
            enable_lto: true,
            strip_debug_symbols: true,
        }
    }
}

/// Mobile framework size optimizer
pub struct SizeOptimizer {
    config: SizeOptimizerConfig,
    metrics: SizeMetrics,
}

impl SizeOptimizer {
    /// Create a new size optimizer
    pub fn new(config: SizeOptimizerConfig) -> Result<Self> {
        let metrics = SizeMetrics {
            current_size_mb: 0.0,
            target_size_mb: config.target_size_mb,
            component_sizes: HashMap::new(),
            optimized_size_mb: 0.0,
            reduction_percentage: 0.0,
        };

        info!(
            "Initialized size optimizer with target: {}MB",
            config.target_size_mb
        );

        Ok(Self { config, metrics })
    }

    /// Analyze current framework size
    pub fn analyze_current_size(&mut self) -> Result<&SizeMetrics> {
        // Simulate current size analysis
        self.metrics.current_size_mb = 96.0; // From release build static lib

        // Breakdown by components
        self.metrics.component_sizes.insert("core".to_string(), 25.0);
        self.metrics.component_sizes.insert("models".to_string(), 35.0);
        self.metrics.component_sizes.insert("tokenizers".to_string(), 15.0);
        self.metrics.component_sizes.insert("mobile_optimizations".to_string(), 12.0);
        self.metrics.component_sizes.insert("dependencies".to_string(), 9.0);

        // Estimate optimization potential
        self.estimate_optimization_potential()?;

        debug!(
            "Analyzed current framework size: {}MB",
            self.metrics.current_size_mb
        );
        Ok(&self.metrics)
    }

    /// Estimate optimization potential
    fn estimate_optimization_potential(&mut self) -> Result<()> {
        let mut potential_savings = 0.0;

        // Dead code elimination savings
        if self.config.eliminate_dead_code {
            potential_savings += self.metrics.current_size_mb * 0.15; // 15% savings
        }

        // Symbol stripping savings
        if self.config.strip_symbols {
            potential_savings += self.metrics.current_size_mb * 0.10; // 10% savings
        }

        // Compression savings
        if self.config.enable_compression {
            let compression_factor = match self.config.compression_level {
                1..=3 => 0.20,
                4..=6 => 0.30,
                7..=9 => 0.40,
                _ => 0.25,
            };
            potential_savings += self.metrics.current_size_mb * compression_factor;
        }

        // LTO savings
        if self.config.enable_lto {
            potential_savings += self.metrics.current_size_mb * 0.12; // 12% savings
        }

        // Feature selection savings
        potential_savings += self.estimate_feature_selection_savings()?;

        self.metrics.optimized_size_mb =
            (self.metrics.current_size_mb - potential_savings).max(5.0);
        self.metrics.reduction_percentage =
            (potential_savings / self.metrics.current_size_mb) * 100.0;

        info!(
            "Estimated optimized size: {}MB ({:.1}% reduction)",
            self.metrics.optimized_size_mb, self.metrics.reduction_percentage
        );

        Ok(())
    }

    /// Estimate savings from feature selection
    fn estimate_feature_selection_savings(&self) -> Result<f64> {
        let savings = match &self.config.strategy {
            SizeOptimizationStrategy::Minimal => {
                self.metrics.current_size_mb * 0.50 // 50% savings with minimal features
            },
            SizeOptimizationStrategy::Balanced => {
                self.metrics.current_size_mb * 0.25 // 25% savings with balanced features
            },
            SizeOptimizationStrategy::MaximumCompressed => {
                self.metrics.current_size_mb * 0.15 // 15% savings with aggressive compression
            },
            SizeOptimizationStrategy::Custom {
                enabled_features, ..
            } => {
                // Estimate based on number of enabled features
                let feature_ratio = enabled_features.len() as f64 / 20.0; // Assume 20 total features
                self.metrics.current_size_mb * (1.0 - feature_ratio) * 0.40
            },
        };

        Ok(savings)
    }

    /// Generate optimization recommendations
    pub fn generate_optimization_recommendations(&self) -> Vec<String> {
        let mut recommendations = Vec::new();

        if self.metrics.optimized_size_mb > self.config.target_size_mb {
            let gap = self.metrics.optimized_size_mb - self.config.target_size_mb;
            warn!(
                "Target size not achievable with current optimizations. Gap: {}MB",
                gap
            );

            recommendations.push(format!(
                "Consider using MinimalStrategy to reduce size by additional {}MB",
                gap
            ));
        }

        if !self.config.enable_lto {
            recommendations
                .push("Enable Link-Time Optimization (LTO) for ~12% size reduction".to_string());
        }

        if self.config.compression_level < 7 {
            recommendations
                .push("Increase compression level to 7-9 for better size reduction".to_string());
        }

        recommendations
            .push("Consider feature-gating optional components (Unity, React Native)".to_string());
        recommendations.push("Use dynamic linking for common dependencies".to_string());
        recommendations.push("Implement lazy loading for infrequently used models".to_string());

        recommendations
    }

    /// Generate optimized build configuration
    pub fn generate_build_config(&self) -> String {
        let mut config_lines = vec![
            "[profile.release]".to_string(),
            "opt-level = 'z'".to_string(), // Optimize for size
            "codegen-units = 1".to_string(),
            "panic = 'abort'".to_string(),
        ];

        if self.config.enable_lto {
            config_lines.push("lto = true".to_string());
        }

        if self.config.strip_debug_symbols {
            config_lines.push("debug = false".to_string());
            config_lines.push("strip = true".to_string());
        }

        // Add feature configuration
        match &self.config.strategy {
            SizeOptimizationStrategy::Minimal => {
                config_lines.push("\n[features]".to_string());
                config_lines.push("default = []".to_string());
                config_lines.push("minimal = []".to_string());
            },
            SizeOptimizationStrategy::Custom {
                enabled_features, ..
            } => {
                config_lines.push("\n[features]".to_string());
                config_lines.push(format!("default = {:?}", enabled_features));
            },
            _ => {},
        }

        config_lines.join("\n")
    }

    /// Check if target size is achievable
    pub fn is_target_achievable(&self) -> bool {
        self.metrics.optimized_size_mb <= self.config.target_size_mb
    }

    /// Get current metrics
    pub fn get_metrics(&self) -> &SizeMetrics {
        &self.metrics
    }

    /// Update configuration
    pub fn update_config(&mut self, config: SizeOptimizerConfig) -> Result<()> {
        self.config = config;
        self.estimate_optimization_potential()?;
        Ok(())
    }

    /// Export optimization report
    pub fn export_optimization_report(&self) -> String {
        serde_json::to_string_pretty(&json!({
            "framework_size_analysis": {
                "current_size_mb": self.metrics.current_size_mb,
                "target_size_mb": self.metrics.target_size_mb,
                "optimized_size_mb": self.metrics.optimized_size_mb,
                "reduction_percentage": self.metrics.reduction_percentage,
                "target_achievable": self.is_target_achievable(),
                "component_breakdown": self.metrics.component_sizes
            },
            "optimization_config": self.config,
            "recommendations": self.generate_optimization_recommendations(),
            "build_config": self.generate_build_config()
        }))
        .unwrap_or_default()
    }
}

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

    #[test]
    fn test_size_optimizer_creation() {
        let config = SizeOptimizerConfig::default();
        let optimizer = SizeOptimizer::new(config).expect("Operation failed");

        assert_eq!(optimizer.config.target_size_mb, 50.0);
        assert!(optimizer.config.enable_lto);
        assert!(optimizer.config.strip_symbols);
    }

    #[test]
    fn test_size_analysis() {
        let config = SizeOptimizerConfig::default();
        let mut optimizer = SizeOptimizer::new(config).expect("Operation failed");

        let metrics = optimizer.analyze_current_size().expect("Operation failed");
        assert!(metrics.current_size_mb > 0.0);
        assert!(metrics.optimized_size_mb < metrics.current_size_mb);
    }

    #[test]
    fn test_target_achievement() {
        let mut config = SizeOptimizerConfig::default();
        config.strategy = SizeOptimizationStrategy::Minimal;

        let mut optimizer = SizeOptimizer::new(config).expect("Operation failed");
        optimizer.analyze_current_size().expect("Operation failed");

        // With minimal strategy, should be able to achieve target
        assert!(optimizer.is_target_achievable());
    }

    #[test]
    fn test_optimization_recommendations() {
        let config = SizeOptimizerConfig::default();
        let mut optimizer = SizeOptimizer::new(config).expect("Operation failed");
        optimizer.analyze_current_size().expect("Operation failed");

        let recommendations = optimizer.generate_optimization_recommendations();
        assert!(!recommendations.is_empty());
    }

    #[test]
    fn test_build_config_generation() {
        let config = SizeOptimizerConfig::default();
        let optimizer = SizeOptimizer::new(config).expect("Operation failed");

        let build_config = optimizer.generate_build_config();
        assert!(build_config.contains("opt-level = 'z'"));
        assert!(build_config.contains("lto = true"));
    }

    #[test]
    fn test_export_report() {
        let config = SizeOptimizerConfig::default();
        let mut optimizer = SizeOptimizer::new(config).expect("Operation failed");
        optimizer.analyze_current_size().expect("Operation failed");

        let report = optimizer.export_optimization_report();
        assert!(report.contains("framework_size_analysis"));
        assert!(report.contains("optimization_config"));
    }
}