//! Core Benchmark Types
//!
//! Defines the main benchmark structures and implementation wrappers.

use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};

use crate::{Backend, BenchmarkConfig, BenchmarkError, Metrics, OperationType, Result};

/// A function that executes a kernel implementation
pub type KernelFn = Arc<dyn Fn(&[Vec<f32>], &[usize]) -> Vec<f32> + Send + Sync>;

/// A named implementation of a kernel
#[derive(Clone)]
pub struct Implementation {
    /// Name of this implementation (e.g., "cubecl", "handcrafted")
    pub name: String,
    /// Description of the implementation
    pub description: String,
    /// The kernel function
    pub kernel: KernelFn,
    /// Backend this implementation targets
    pub backend: Backend,
}

impl fmt::Debug for Implementation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Implementation")
            .field("name", &self.name)
            .field("description", &self.description)
            .field("backend", &self.backend)
            .finish()
    }
}

impl Implementation {
    /// Create a new implementation
    pub fn new<F>(name: impl Into<String>, backend: Backend, kernel: F) -> Self
    where
        F: Fn(&[Vec<f32>], &[usize]) -> Vec<f32> + Send + Sync + 'static,
    {
        Self {
            name: name.into(),
            description: String::new(),
            kernel: Arc::new(kernel),
            backend,
        }
    }

    /// Set the description
    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
        self.description = desc.into();
        self
    }

    /// Execute the kernel
    pub fn execute(&self, inputs: &[Vec<f32>], shape: &[usize]) -> Vec<f32> {
        (self.kernel)(inputs, shape)
    }
}

/// Result of a single benchmark measurement
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Measurement {
    /// Duration of this measurement
    pub duration: Duration,
    /// Memory used (bytes) if available
    pub memory_bytes: Option<u64>,
    /// Throughput (elements/second) if applicable
    pub throughput: Option<f64>,
}

/// Result of running a benchmark
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct BenchmarkResult {
    /// Name of the benchmark
    pub name: String,
    /// Operation type
    pub operation: OperationType,
    /// Implementation name
    pub implementation: String,
    /// Backend used
    pub backend: Backend,
    /// Input shape
    pub shape: Vec<usize>,
    /// Total number of elements
    pub num_elements: usize,
    /// All measurements
    pub measurements: Vec<Measurement>,
    /// Computed metrics
    pub metrics: Metrics,
    /// Timestamp when benchmark was run
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

impl BenchmarkResult {
    /// Create a new benchmark result from measurements
    pub fn from_measurements(
        name: String,
        operation: OperationType,
        implementation: String,
        backend: Backend,
        shape: Vec<usize>,
        measurements: Vec<Measurement>,
    ) -> Self {
        let num_elements = shape.iter().product();
        let metrics = Metrics::from_measurements(&measurements);

        Self {
            name,
            operation,
            implementation,
            backend,
            shape,
            num_elements,
            measurements,
            metrics,
            timestamp: chrono::Utc::now(),
        }
    }

    /// Get the mean duration
    pub fn mean_duration(&self) -> Duration {
        self.metrics.timing.mean
    }

    /// Get throughput in GFLOPS (assuming 2 FLOPs per element for most ops)
    pub fn gflops(&self) -> f64 {
        let flops = self.num_elements as f64 * 2.0; // Approximate
        let seconds = self.metrics.timing.mean.as_secs_f64();
        flops / seconds / 1e9
    }

    /// Get throughput in GB/s
    pub fn bandwidth_gbps(&self) -> f64 {
        let bytes = self.num_elements as f64 * 4.0 * 2.0; // Read + write, f32
        let seconds = self.metrics.timing.mean.as_secs_f64();
        bytes / seconds / 1e9
    }
}

/// Builder for kernel benchmarks
#[derive(Debug)]
pub struct KernelBenchmark {
    /// Name of the benchmark
    name: String,
    /// Operation being benchmarked
    operation: OperationType,
    /// Shapes to benchmark
    shapes: Vec<Vec<usize>>,
    /// Implementation A (typically CubeCL)
    impl_a: Option<Implementation>,
    /// Implementation B (typically hand-crafted)
    impl_b: Option<Implementation>,
    /// Additional implementations
    other_impls: Vec<Implementation>,
}

impl KernelBenchmark {
    /// Create a new kernel benchmark
    pub fn new(name: impl Into<String>, operation: OperationType) -> Self {
        Self {
            name: name.into(),
            operation,
            shapes: Vec::new(),
            impl_a: None,
            impl_b: None,
            other_impls: Vec::new(),
        }
    }

    /// Add shapes to benchmark
    pub fn with_shapes(mut self, shapes: Vec<Vec<usize>>) -> Self {
        self.shapes = shapes;
        self
    }

    /// Add a single shape
    pub fn add_shape(mut self, shape: Vec<usize>) -> Self {
        self.shapes.push(shape);
        self
    }

    /// Set implementation A (typically CubeCL or reference)
    pub fn implementation_a(mut self, impl_a: Implementation) -> Self {
        self.impl_a = Some(impl_a);
        self
    }

    /// Set implementation B (typically hand-crafted or optimized)
    pub fn implementation_b(mut self, impl_b: Implementation) -> Self {
        self.impl_b = Some(impl_b);
        self
    }

    /// Add additional implementation for comparison
    pub fn add_implementation(mut self, impl_: Implementation) -> Self {
        self.other_impls.push(impl_);
        self
    }

    /// Get all implementations
    pub fn implementations(&self) -> Vec<&Implementation> {
        let mut impls = Vec::new();
        if let Some(ref a) = self.impl_a {
            impls.push(a);
        }
        if let Some(ref b) = self.impl_b {
            impls.push(b);
        }
        impls.extend(self.other_impls.iter());
        impls
    }

    /// Get the benchmark name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Get the operation type
    pub fn operation(&self) -> &OperationType {
        &self.operation
    }

    /// Run the benchmark with the given configuration
    pub fn run(&self, config: &BenchmarkConfig) -> Result<Vec<BenchmarkResult>> {
        let mut results = Vec::new();

        // Validate
        if self.shapes.is_empty() {
            return Err(BenchmarkError::ConfigError("No shapes specified".into()));
        }

        let impls = self.implementations();
        if impls.is_empty() {
            return Err(BenchmarkError::ConfigError(
                "No implementations specified".into(),
            ));
        }

        // Run benchmarks for each shape and implementation
        for shape in &self.shapes {
            // Generate test inputs
            let inputs = self.generate_inputs(shape);

            for impl_ in &impls {
                if !impl_.backend.is_available() && !config.skip_unavailable_backends {
                    return Err(BenchmarkError::BackendNotAvailable(
                        impl_.backend.name().to_string(),
                    ));
                }

                if !impl_.backend.is_available() {
                    continue;
                }

                let measurements = self.run_implementation(impl_, &inputs, shape, config)?;

                results.push(BenchmarkResult::from_measurements(
                    self.name.clone(),
                    self.operation.clone(),
                    impl_.name.clone(),
                    impl_.backend,
                    shape.clone(),
                    measurements,
                ));
            }
        }

        Ok(results)
    }

    /// Generate random inputs for a given shape
    fn generate_inputs(&self, shape: &[usize]) -> Vec<Vec<f32>> {
        let num_elements: usize = shape.iter().product();

        // Generate 1-2 input tensors depending on operation type
        let num_inputs = if self.operation.is_elementwise() {
            2
        } else {
            1
        };

        (0..num_inputs)
            .map(|_| {
                (0..num_elements)
                    .map(|i| ((i as f32 * 0.001) % 1.0) - 0.5)
                    .collect()
            })
            .collect()
    }

    /// Run a single implementation
    fn run_implementation(
        &self,
        impl_: &Implementation,
        inputs: &[Vec<f32>],
        shape: &[usize],
        config: &BenchmarkConfig,
    ) -> Result<Vec<Measurement>> {
        use crate::kernels::memory::MemorySnapshot;

        // Warmup
        for _ in 0..config.warmup.iterations {
            let _ = impl_.execute(inputs, shape);
        }

        // Actual measurements
        let mut measurements = Vec::with_capacity(config.iterations);

        for _ in 0..config.iterations {
            // Take memory snapshot before if tracking enabled
            let mem_before = if config.track_memory {
                Some(MemorySnapshot::now())
            } else {
                None
            };

            let start = Instant::now();
            let _ = impl_.execute(inputs, shape);
            let duration = start.elapsed();

            // Take memory snapshot after and compute delta
            let memory_bytes = if config.track_memory {
                let mem_after = MemorySnapshot::now();
                match (mem_before, mem_after.total()) {
                    (Some(_before), Some(after)) => {
                        // Report current memory usage (not delta, which could be negative)
                        Some(after)
                    }
                    (None, Some(after)) => Some(after),
                    _ => None,
                }
            } else {
                None
            };

            measurements.push(Measurement {
                duration,
                memory_bytes,
                throughput: None, // Computed later in Metrics
            });
        }

        Ok(measurements)
    }
}

/// Standard benchmark shapes for common operations
pub mod shapes {
    /// Small shapes for quick testing
    pub fn small() -> Vec<Vec<usize>> {
        vec![vec![64, 64], vec![128, 128], vec![256, 256]]
    }

    /// Medium shapes for typical workloads
    pub fn medium() -> Vec<Vec<usize>> {
        vec![vec![512, 512], vec![1024, 1024], vec![2048, 2048]]
    }

    /// Large shapes for stress testing
    pub fn large() -> Vec<Vec<usize>> {
        vec![vec![4096, 4096], vec![8192, 8192]]
    }

    /// LLM-specific shapes (batch, seq_len, hidden_dim)
    pub fn llm() -> Vec<Vec<usize>> {
        vec![
            vec![1, 512, 4096],  // Single sequence
            vec![8, 512, 4096],  // Small batch
            vec![32, 512, 4096], // Medium batch
            vec![1, 2048, 4096], // Long sequence
            vec![8, 2048, 4096], // Long batch
        ]
    }

    /// Attention shapes (batch, heads, seq_len, head_dim)
    pub fn attention() -> Vec<Vec<usize>> {
        vec![
            vec![1, 32, 512, 128],
            vec![8, 32, 512, 128],
            vec![1, 32, 2048, 128],
            vec![8, 32, 2048, 128],
        ]
    }
}

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

    fn dummy_kernel(inputs: &[Vec<f32>], shape: &[usize]) -> Vec<f32> {
        let size: usize = shape.iter().product();
        vec![0.0; size]
    }

    #[test]
    fn test_implementation_creation() {
        let impl_ = Implementation::new("test", Backend::Cpu, dummy_kernel)
            .with_description("A test implementation");

        assert_eq!(impl_.name, "test");
        assert_eq!(impl_.backend, Backend::Cpu);
    }

    #[test]
    fn test_benchmark_builder() {
        let bench = KernelBenchmark::new("test_add", OperationType::Add)
            .with_shapes(vec![vec![64, 64], vec![128, 128]])
            .implementation_a(Implementation::new("impl_a", Backend::Cpu, dummy_kernel))
            .implementation_b(Implementation::new("impl_b", Backend::Cpu, dummy_kernel));

        assert_eq!(bench.implementations().len(), 2);
    }

    #[test]
    fn test_run_benchmark() {
        let bench = KernelBenchmark::new("test_add", OperationType::Add)
            .with_shapes(vec![vec![64, 64]])
            .implementation_a(Implementation::new("cpu", Backend::Cpu, dummy_kernel));

        let config = BenchmarkConfig {
            iterations: 3,
            warmup: crate::WarmupConfig {
                iterations: 1,
                ..Default::default()
            },
            skip_unavailable_backends: true,
            ..Default::default()
        };

        let results = bench.run(&config).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].measurements.len(), 3);
    }

    #[test]
    fn test_shape_generators() {
        assert!(!shapes::small().is_empty());
        assert!(!shapes::medium().is_empty());
        assert!(!shapes::llm().is_empty());
    }
}