//! Benchmark Generator
//!
//! Automatic generation of performance benchmarks for model implementations.

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// Benchmark configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BenchmarkGeneratorConfig {
    /// Model name to benchmark
    pub model_name: String,
    /// Benchmark types to include
    pub benchmark_types: Vec<BenchmarkType>,
    /// Test configurations
    pub test_configs: HashMap<String, TestConfig>,
    /// Hardware targets
    pub hardware_targets: Vec<HardwareTarget>,
}

/// Type of benchmark to generate
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BenchmarkType {
    Latency,
    Throughput,
    Memory,
    Accuracy,
    Scalability,
    Comparative,
}

/// Test configuration for benchmarks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestConfig {
    pub batch_sizes: Vec<usize>,
    pub sequence_lengths: Vec<usize>,
    pub iterations: usize,
    pub warmup_iterations: usize,
}

/// Hardware target for benchmarks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum HardwareTarget {
    CPU,
    GPU,
    Metal,
    WASM,
}

/// Benchmark generator
pub struct BenchmarkGenerator {
    config: BenchmarkGeneratorConfig,
}

impl BenchmarkGenerator {
    /// Create a new benchmark generator
    pub fn new(config: BenchmarkGeneratorConfig) -> Self {
        Self { config }
    }

    /// Generate benchmark suite
    pub fn generate_benchmarks(&self, output_path: &Path) -> Result<()> {
        let mut benchmark_content = self.generate_header();

        for benchmark_type in &self.config.benchmark_types {
            match benchmark_type {
                BenchmarkType::Latency => {
                    benchmark_content.push_str(&self.generate_latency_benchmark())
                },
                BenchmarkType::Throughput => {
                    benchmark_content.push_str(&self.generate_throughput_benchmark())
                },
                BenchmarkType::Memory => {
                    benchmark_content.push_str(&self.generate_memory_benchmark())
                },
                BenchmarkType::Accuracy => {
                    benchmark_content.push_str(&self.generate_accuracy_benchmark())
                },
                BenchmarkType::Scalability => {
                    benchmark_content.push_str(&self.generate_scalability_benchmark())
                },
                BenchmarkType::Comparative => {
                    benchmark_content.push_str(&self.generate_comparative_benchmark())
                },
            }
        }

        std::fs::write(output_path, benchmark_content)?;
        Ok(())
    }

    /// Generate benchmark file header
    fn generate_header(&self) -> String {
        format!(
            "//! Performance Benchmarks for {}\n//!\n//! This file is auto-generated. Do not edit manually.\n\nuse criterion::{{black_box, criterion_group, criterion_main, Criterion}};\nuse super::{{{}Config, {}Model}};\nuse trustformers_core::tensor::Tensor;\nuse std::time::Duration;\n\n",
            self.config.model_name,
            self.config.model_name,
            self.config.model_name
        )
    }

    /// Generate latency benchmark
    fn generate_latency_benchmark(&self) -> String {
        format!(
            "// ========== Latency Benchmarks ==========\n\nfn bench_{}_latency(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_latency\");\n    group.warm_up_time(Duration::from_millis(500));\n    group.measurement_time(Duration::from_secs(3));\n    \n    let batch_sizes = vec![1, 4, 8, 16];\n    let sequence_length = 512;\n    \n    for batch_size in batch_sizes {{\n        group.bench_with_input(\n            format!(\"batch_{{}}\", batch_size),\n            &batch_size,\n            |b, &batch_size| {{\n                let input = Tensor::randn(&[batch_size, sequence_length]);\n                b.iter(|| {{\n                    black_box(model.forward(black_box(&input)).expect(\"Forward pass failed\"))\n                }});\n            }},\n        );\n    }}\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }

    /// Generate throughput benchmark
    fn generate_throughput_benchmark(&self) -> String {
        format!(
            "// ========== Throughput Benchmarks ==========\n\nfn bench_{}_throughput(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_throughput\");\n    group.warm_up_time(Duration::from_millis(500));\n    group.measurement_time(Duration::from_secs(5));\n    group.throughput(criterion::Throughput::Elements(1000));\n    \n    let batch_size = 8;\n    let sequence_length = 512;\n    let input = Tensor::randn(&[batch_size, sequence_length]);\n    \n    group.bench_function(\"tokens_per_second\", |b| {{\n        b.iter(|| {{\n            for _ in 0..100 {{\n                black_box(model.forward(black_box(&input)).expect(\"Forward pass failed\"));\n            }}\n        }});\n    }});\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }

    /// Generate memory benchmark
    fn generate_memory_benchmark(&self) -> String {
        format!(
            "// ========== Memory Benchmarks ==========\n\nfn bench_{}_memory(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_memory\");\n    group.warm_up_time(Duration::from_millis(100));\n    group.measurement_time(Duration::from_secs(2));\n    \n    let sequence_lengths = vec![128, 256, 512, 1024];\n    let batch_size = 4;\n    \n    for seq_len in sequence_lengths {{\n        group.bench_with_input(\n            format!(\"seq_len_{{}}\", seq_len),\n            &seq_len,\n            |b, &seq_len| {{\n                b.iter(|| {{\n                    let input = Tensor::randn(&[batch_size, seq_len]);\n                    black_box(model.forward(black_box(&input)).expect(\"Forward pass failed\"))\n                }});\n            }},\n        );\n    }}\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }

    /// Generate accuracy benchmark
    fn generate_accuracy_benchmark(&self) -> String {
        format!(
            "// ========== Accuracy Benchmarks ==========\n\nfn bench_{}_accuracy(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_accuracy\");\n    group.warm_up_time(Duration::from_millis(200));\n    group.measurement_time(Duration::from_secs(2));\n    \n    // Test numerical precision\n    let input = Tensor::ones(&[4, 512]);\n    \n    group.bench_function(\"numerical_precision\", |b| {{\n        b.iter(|| {{\n            let output = model.forward(black_box(&input)).expect(\"Forward pass failed\");\n            \n            // Verify output consistency\n            match &output {{\n                Tensor::F32(arr) => {{\n                    let sum: f32 = arr.iter().sum();\n                    black_box(sum);\n                }}\n                _ => panic!(\"Expected F32 tensor\"),\n            }}\n        }});\n    }});\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }

    /// Generate scalability benchmark
    fn generate_scalability_benchmark(&self) -> String {
        format!(
            "// ========== Scalability Benchmarks ==========\n\nfn bench_{}_scalability(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_scalability\");\n    group.warm_up_time(Duration::from_millis(500));\n    group.measurement_time(Duration::from_secs(3));\n    \n    // Test different configurations\n    let test_cases = vec![\n        (1, 128),   // Small: 1 batch, 128 sequence\n        (4, 256),   // Medium: 4 batch, 256 sequence\n        (8, 512),   // Large: 8 batch, 512 sequence\n        (16, 256),  // Wide: 16 batch, 256 sequence\n    ];\n    \n    for (batch_size, seq_len) in test_cases {{\n        group.bench_with_input(\n            format!(\"batch_{{}}x{{}}\", batch_size, seq_len),\n            &(batch_size, seq_len),\n            |b, &(batch_size, seq_len)| {{\n                let input = Tensor::randn(&[batch_size, seq_len]);\n                b.iter(|| {{\n                    black_box(model.forward(black_box(&input)).expect(\"Forward pass failed\"))\n                }});\n            }},\n        );\n    }}\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }

    /// Generate comparative benchmark
    fn generate_comparative_benchmark(&self) -> String {
        format!(
            "// ========== Comparative Benchmarks ==========\n\nfn bench_{}_comparative(c: &mut Criterion) {{\n    let config = {}Config::default();\n    let model = {}Model::new(config).expect(\"Failed to create model\");\n    \n    let mut group = c.benchmark_group(\"{}_comparative\");\n    group.warm_up_time(Duration::from_millis(300));\n    group.measurement_time(Duration::from_secs(2));\n    \n    let input = Tensor::randn(&[8, 512]);\n    \n    group.bench_function(\"baseline\", |b| {{\n        b.iter(|| {{\n            black_box(model.forward(black_box(&input)).expect(\"Forward pass failed\"))\n        }});\n    }});\n    \n    // Add comparisons with other models or configurations here\n    // This would typically compare against reference implementations\n    \n    group.finish();\n}}\n\n",
            self.config.model_name.to_lowercase(),
            self.config.model_name,
            self.config.model_name,
            self.config.model_name.to_lowercase()
        )
    }
}

impl BenchmarkGenerator {
    /// Generate the main benchmark group registration
    pub fn generate_main(&self) -> String {
        let function_names: Vec<String> = self
            .config
            .benchmark_types
            .iter()
            .map(|bt| {
                format!(
                    "bench_{}_{}",
                    self.config.model_name.to_lowercase(),
                    match bt {
                        BenchmarkType::Latency => "latency",
                        BenchmarkType::Throughput => "throughput",
                        BenchmarkType::Memory => "memory",
                        BenchmarkType::Accuracy => "accuracy",
                        BenchmarkType::Scalability => "scalability",
                        BenchmarkType::Comparative => "comparative",
                    }
                )
            })
            .collect();

        format!(
            "criterion_group!(benches, {});\ncriterion_main!(benches);\n",
            function_names.join(", ")
        )
    }
}

/// Predefined benchmark configurations
pub struct BenchmarkTemplates;

impl BenchmarkTemplates {
    /// Get comprehensive benchmark configuration
    pub fn comprehensive(model_name: String) -> BenchmarkGeneratorConfig {
        let mut test_configs = HashMap::new();
        test_configs.insert(
            "default".to_string(),
            TestConfig {
                batch_sizes: vec![1, 4, 8, 16],
                sequence_lengths: vec![128, 256, 512, 1024],
                iterations: 100,
                warmup_iterations: 10,
            },
        );

        BenchmarkGeneratorConfig {
            model_name,
            benchmark_types: vec![
                BenchmarkType::Latency,
                BenchmarkType::Throughput,
                BenchmarkType::Memory,
                BenchmarkType::Accuracy,
                BenchmarkType::Scalability,
                BenchmarkType::Comparative,
            ],
            test_configs,
            hardware_targets: vec![HardwareTarget::CPU],
        }
    }

    /// Get performance-focused benchmark configuration
    pub fn performance(model_name: String) -> BenchmarkGeneratorConfig {
        let mut test_configs = HashMap::new();
        test_configs.insert(
            "performance".to_string(),
            TestConfig {
                batch_sizes: vec![1, 8, 16, 32],
                sequence_lengths: vec![256, 512, 1024],
                iterations: 200,
                warmup_iterations: 20,
            },
        );

        BenchmarkGeneratorConfig {
            model_name,
            benchmark_types: vec![
                BenchmarkType::Latency,
                BenchmarkType::Throughput,
                BenchmarkType::Scalability,
            ],
            test_configs,
            hardware_targets: vec![HardwareTarget::CPU, HardwareTarget::GPU],
        }
    }

    /// Get accuracy-focused benchmark configuration
    pub fn accuracy(model_name: String) -> BenchmarkGeneratorConfig {
        let mut test_configs = HashMap::new();
        test_configs.insert(
            "accuracy".to_string(),
            TestConfig {
                batch_sizes: vec![4, 8],
                sequence_lengths: vec![512],
                iterations: 50,
                warmup_iterations: 5,
            },
        );

        BenchmarkGeneratorConfig {
            model_name,
            benchmark_types: vec![BenchmarkType::Accuracy, BenchmarkType::Comparative],
            test_configs,
            hardware_targets: vec![HardwareTarget::CPU],
        }
    }
}
