//! FoundationForge - Model Compression & Distillation Hub.
//!
//! This demo showcases model compression techniques including quantization,
//! pruning, and knowledge distillation for efficient model deployment.

pub mod distiller;
pub mod pruner;
pub mod quantizer;
pub mod sample_data;

use thiserror::Error;

use forge_shared::{
    BenchmarkConfig, BenchmarkResult, CompressionProgress, CompressionResult, CompressionStep,
    DistillationConfig, ExportConfig, ModelInfo, PipelineConfig, PruningConfig, QuantizationConfig,
    StudentConfig,
};

/// Errors that can occur in FoundationForge.
#[derive(Debug, Error)]
pub enum ForgeError {
    /// Invalid configuration.
    #[error("Invalid configuration: {0}")]
    InvalidConfig(String),

    /// Compression failed.
    #[error("Compression failed: {0}")]
    CompressionFailed(String),

    /// Model not loaded.
    #[error("Model not loaded")]
    ModelNotLoaded,

    /// Export failed.
    #[error("Export failed: {0}")]
    ExportFailed(String),

    /// Benchmark failed.
    #[error("Benchmark failed: {0}")]
    BenchmarkFailed(String),
}

/// Main FoundationForge system.
#[derive(Debug)]
pub struct FoundationForge {
    /// Original model info.
    original_model: ModelInfo,
    /// Current model info (after compression).
    current_model: ModelInfo,
    /// Quantizer.
    quantizer: quantizer::Quantizer,
    /// Pruner.
    pruner: pruner::Pruner,
    /// Distiller.
    distiller: distiller::Distiller,
    /// Applied compression steps.
    applied_steps: Vec<String>,
    /// Is model loaded.
    loaded: bool,
}

impl Default for FoundationForge {
    fn default() -> Self {
        Self::new(ModelInfo::default())
    }
}

impl FoundationForge {
    /// Create a new FoundationForge instance.
    #[must_use]
    pub fn new(model: ModelInfo) -> Self {
        Self {
            original_model: model.clone(),
            current_model: model,
            quantizer: quantizer::Quantizer::new(),
            pruner: pruner::Pruner::new(),
            distiller: distiller::Distiller::new(),
            applied_steps: Vec::new(),
            loaded: true,
        }
    }

    /// Apply a compression pipeline.
    pub fn compress(
        &mut self,
        pipeline: &PipelineConfig,
        progress_callback: Option<Box<dyn Fn(CompressionProgress) + Send>>,
    ) -> Result<CompressionResult, ForgeError> {
        if !self.loaded {
            return Err(ForgeError::ModelNotLoaded);
        }

        let start_time = std::time::Instant::now();
        let total_steps = pipeline.steps.len();

        for (step_idx, step) in pipeline.steps.iter().enumerate() {
            let step_name = match step {
                CompressionStep::Quantize(_) => "Quantization",
                CompressionStep::Prune(_) => "Pruning",
                CompressionStep::Distill(_) => "Distillation",
                CompressionStep::RemoveLayers(_) => "Layer Removal",
                CompressionStep::PruneVocab { .. } => "Vocabulary Pruning",
                CompressionStep::PruneHeads { .. } => "Head Pruning",
            };

            if let Some(ref callback) = progress_callback {
                callback(CompressionProgress {
                    step: step_idx + 1,
                    total_steps,
                    step_name: step_name.to_string(),
                    step_progress: 0.0,
                    current_size: self.current_model.size_bytes,
                    original_size: self.original_model.size_bytes,
                    compression_ratio: self.original_model.size_bytes as f32
                        / self.current_model.size_bytes.max(1) as f32,
                    current_metric: None,
                    original_metric: None,
                    elapsed_seconds: start_time.elapsed().as_secs_f64(),
                });
            }

            match step {
                CompressionStep::Quantize(config) => {
                    self.apply_quantization(config)?;
                }
                CompressionStep::Prune(config) => {
                    self.apply_pruning(config)?;
                }
                CompressionStep::Distill(config) => {
                    self.apply_distillation(config)?;
                }
                CompressionStep::RemoveLayers(layers) => {
                    self.remove_layers(layers)?;
                }
                CompressionStep::PruneVocab { keep_tokens } => {
                    self.prune_vocabulary(*keep_tokens)?;
                }
                CompressionStep::PruneHeads { target_heads } => {
                    self.prune_heads(*target_heads)?;
                }
            }

            self.applied_steps.push(step_name.to_string());

            if let Some(ref callback) = progress_callback {
                callback(CompressionProgress {
                    step: step_idx + 1,
                    total_steps,
                    step_name: step_name.to_string(),
                    step_progress: 1.0,
                    current_size: self.current_model.size_bytes,
                    original_size: self.original_model.size_bytes,
                    compression_ratio: self.original_model.size_bytes as f32
                        / self.current_model.size_bytes.max(1) as f32,
                    current_metric: None,
                    original_metric: None,
                    elapsed_seconds: start_time.elapsed().as_secs_f64(),
                });
            }
        }

        Ok(CompressionResult {
            original_model: self.original_model.clone(),
            compressed_model: self.current_model.clone(),
            compression_ratio: self.original_model.size_bytes as f32
                / self.current_model.size_bytes.max(1) as f32,
            speedup_ratio: Some(self.estimate_speedup()),
            original_metric: 5.0, // Example perplexity
            compressed_metric: 5.2,
            metric_degradation: 0.04,
            compression_time: start_time.elapsed().as_secs_f64(),
            steps_applied: self.applied_steps.clone(),
            output_path: pipeline.export.output_path.clone(),
        })
    }

    /// Apply quantization.
    fn apply_quantization(&mut self, config: &QuantizationConfig) -> Result<(), ForgeError> {
        let compression_factor = self.quantizer.quantize(&self.current_model, config);

        // Update model info
        self.current_model.size_bytes =
            (self.current_model.size_bytes as f32 / compression_factor) as u64;
        self.current_model.precision_bits = match config.precision {
            forge_shared::QuantPrecision::Int2 => 2,
            forge_shared::QuantPrecision::Int3 => 3,
            forge_shared::QuantPrecision::Int4 => 4,
            forge_shared::QuantPrecision::Int8 => 8,
            forge_shared::QuantPrecision::FP8 => 8,
            forge_shared::QuantPrecision::FP16 => 16,
            forge_shared::QuantPrecision::BF16 => 16,
            forge_shared::QuantPrecision::Mixed => 8,
        };

        Ok(())
    }

    /// Apply pruning.
    fn apply_pruning(&mut self, config: &PruningConfig) -> Result<(), ForgeError> {
        let remaining_params = self.pruner.prune(&self.current_model, config);

        // Update model info
        self.current_model.num_parameters =
            (self.current_model.num_parameters as f32 * remaining_params) as u64;
        self.current_model.size_bytes =
            (self.current_model.size_bytes as f32 * remaining_params) as u64;

        Ok(())
    }

    /// Apply distillation.
    fn apply_distillation(&mut self, config: &DistillationConfig) -> Result<(), ForgeError> {
        let student_info = self.distiller.distill(&self.current_model, config);

        // Replace current model with student
        self.current_model = student_info;

        Ok(())
    }

    /// Remove specific layers.
    fn remove_layers(&mut self, layers: &[usize]) -> Result<(), ForgeError> {
        let remaining_layers = self.current_model.num_layers - layers.len();
        let reduction_factor = remaining_layers as f32 / self.current_model.num_layers as f32;

        self.current_model.num_layers = remaining_layers;
        self.current_model.num_parameters =
            (self.current_model.num_parameters as f32 * reduction_factor) as u64;
        self.current_model.size_bytes =
            (self.current_model.size_bytes as f32 * reduction_factor) as u64;

        Ok(())
    }

    /// Prune vocabulary.
    fn prune_vocabulary(&mut self, keep_tokens: usize) -> Result<(), ForgeError> {
        if let Some(vocab_size) = self.current_model.vocab_size
            && keep_tokens < vocab_size
        {
            self.current_model.vocab_size = Some(keep_tokens);
            // Embedding table reduction
            let embedding_reduction =
                (self.current_model.hidden_dim * (vocab_size - keep_tokens)) as u64 * 2;
            self.current_model.size_bytes = self
                .current_model
                .size_bytes
                .saturating_sub(embedding_reduction);
        }

        Ok(())
    }

    /// Prune attention heads.
    fn prune_heads(&mut self, target_heads: usize) -> Result<(), ForgeError> {
        // Simplified: reduce hidden dimension proportionally
        let head_reduction = target_heads as f32 / 32.0; // Assume 32 heads originally
        self.current_model.size_bytes =
            (self.current_model.size_bytes as f32 * (0.7 + 0.3 * head_reduction)) as u64;

        Ok(())
    }

    /// Estimate speedup from compression.
    fn estimate_speedup(&self) -> f32 {
        let size_ratio =
            self.original_model.size_bytes as f32 / self.current_model.size_bytes.max(1) as f32;
        let param_ratio = self.original_model.num_parameters as f32
            / self.current_model.num_parameters.max(1) as f32;

        // Speedup is roughly proportional to parameter reduction
        (size_ratio * 0.3 + param_ratio * 0.7).sqrt()
    }

    /// Create a student model for distillation.
    pub fn create_student(&self, config: &StudentConfig) -> ModelInfo {
        ModelInfo {
            name: format!("{}-student", self.current_model.name),
            architecture: self.current_model.architecture,
            num_parameters: self.estimate_student_params(config),
            num_layers: config.num_layers,
            hidden_dim: config.hidden_dim,
            vocab_size: self.current_model.vocab_size,
            image_size: self.current_model.image_size,
            precision_bits: 16,
            size_bytes: self.estimate_student_params(config) * 2, // FP16
        }
    }

    /// Estimate student model parameters.
    fn estimate_student_params(&self, config: &StudentConfig) -> u64 {
        let layer_params = config.hidden_dim * config.intermediate_dim * 2; // FFN
        let attn_params = config.hidden_dim * config.hidden_dim * 4; // Q, K, V, O
        let params_per_layer = (layer_params + attn_params) as u64;

        let embedding_params =
            self.current_model.vocab_size.unwrap_or(32000) as u64 * config.hidden_dim as u64;

        config.num_layers as u64 * params_per_layer + embedding_params * 2
    }

    /// Benchmark the compressed model.
    pub fn benchmark(&self, _config: &BenchmarkConfig) -> Result<BenchmarkResult, ForgeError> {
        if !self.loaded {
            return Err(ForgeError::ModelNotLoaded);
        }

        // Simulated benchmark results
        let latency_base = 50.0; // ms
        let speedup = self.estimate_speedup();

        Ok(BenchmarkResult {
            perplexity: Some(5.2),
            accuracy: None,
            avg_latency_ms: Some(latency_base / speedup as f64),
            p50_latency_ms: Some(latency_base / speedup as f64 * 0.95),
            p99_latency_ms: Some(latency_base / speedup as f64 * 1.5),
            throughput: Some(1000.0 * speedup as f64 / latency_base),
            memory_bytes: Some(self.current_model.size_bytes),
            peak_memory_bytes: Some((self.current_model.size_bytes as f64 * 1.2) as u64),
        })
    }

    /// Export the compressed model.
    pub fn export(&self, config: &ExportConfig) -> Result<String, ForgeError> {
        if !self.loaded {
            return Err(ForgeError::ModelNotLoaded);
        }

        // Simulated export
        let format_ext = match config.format {
            forge_shared::ExportFormat::ONNX => "onnx",
            forge_shared::ExportFormat::TFLite => "tflite",
            forge_shared::ExportFormat::CoreML => "mlmodel",
            forge_shared::ExportFormat::GGUF => "gguf",
            forge_shared::ExportFormat::SafeTensors => "safetensors",
            forge_shared::ExportFormat::PyTorch => "pt",
            forge_shared::ExportFormat::Custom => "bin",
        };

        Ok(format!("{}.{}", config.output_path, format_ext))
    }

    /// Get original model info.
    #[must_use]
    pub fn original_model(&self) -> &ModelInfo {
        &self.original_model
    }

    /// Get current model info.
    #[must_use]
    pub fn current_model(&self) -> &ModelInfo {
        &self.current_model
    }

    /// Get compression ratio.
    #[must_use]
    pub fn compression_ratio(&self) -> f32 {
        self.original_model.size_bytes as f32 / self.current_model.size_bytes.max(1) as f32
    }

    /// Get applied steps.
    #[must_use]
    pub fn applied_steps(&self) -> &[String] {
        &self.applied_steps
    }
}

/// Run the demo.
pub fn run_demo() -> Result<CompressionResult, ForgeError> {
    let model = forge_shared::sample_model_info();
    let mut forge = FoundationForge::new(model);

    let pipeline = forge_shared::sample_pipeline_config();
    forge.compress(&pipeline, None)
}

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

    #[test]
    fn test_forge_creation() {
        let forge = FoundationForge::default();
        assert!(forge.loaded);
    }

    #[test]
    fn test_compress_quantization() {
        let model = forge_shared::sample_model_info();
        let mut forge = FoundationForge::new(model);

        let pipeline = PipelineConfig {
            name: "quant-only".to_string(),
            steps: vec![CompressionStep::Quantize(QuantizationConfig::default())],
            ..Default::default()
        };

        let result = forge.compress(&pipeline, None);
        assert!(result.is_ok());

        let result = result.unwrap();
        assert!(result.compression_ratio > 1.0);
    }

    #[test]
    fn test_compress_pruning() {
        let model = forge_shared::sample_model_info();
        let mut forge = FoundationForge::new(model);

        let pipeline = PipelineConfig {
            name: "prune-only".to_string(),
            steps: vec![CompressionStep::Prune(PruningConfig::default())],
            ..Default::default()
        };

        let result = forge.compress(&pipeline, None);
        assert!(result.is_ok());
    }

    #[test]
    fn test_compress_distillation() {
        let model = forge_shared::sample_model_info();
        let mut forge = FoundationForge::new(model);

        let pipeline = PipelineConfig {
            name: "distill-only".to_string(),
            steps: vec![CompressionStep::Distill(DistillationConfig::default())],
            ..Default::default()
        };

        let result = forge.compress(&pipeline, None);
        assert!(result.is_ok());
    }

    #[test]
    fn test_full_pipeline() {
        let model = forge_shared::sample_model_info();
        let mut forge = FoundationForge::new(model);

        let pipeline = forge_shared::sample_pipeline_config();
        let result = forge.compress(&pipeline, None);

        assert!(result.is_ok());
        let result = result.unwrap();
        assert_eq!(result.steps_applied.len(), 2);
    }

    #[test]
    fn test_create_student() {
        let forge = FoundationForge::default();
        let student_config = StudentConfig::default();

        let student = forge.create_student(&student_config);
        assert!(student.num_layers < forge.current_model().num_layers);
    }

    #[test]
    fn test_benchmark() {
        let forge = FoundationForge::default();
        let config = BenchmarkConfig::default();

        let result = forge.benchmark(&config);
        assert!(result.is_ok());
    }

    #[test]
    fn test_export() {
        let forge = FoundationForge::default();
        let config = ExportConfig::default();

        let result = forge.export(&config);
        assert!(result.is_ok());
        assert!(result.unwrap().contains("safetensors"));
    }

    #[test]
    fn test_run_demo() {
        let result = run_demo();
        assert!(result.is_ok());
    }
}