//! TMA-Accelerated Operations
//!
//! Provides Tensor Memory Accelerator (TMA) accelerated operations.

use std::fmt;
use tracing::{debug, instrument};

use super::tile_kernel::TileConfig;
use crate::error::{IntegrationError, Result};

#[cfg(all(target_os = "macos", feature = "tma-metal"))]
use tensor_accelerator::metal::{MetalTensor, MetalTensorDevice};

/// Configuration for TMA operations
#[derive(Debug, Clone)]
pub struct TmaConfig {
    /// Enable async memory copies
    pub async_copy: bool,
    /// Tile size for M dimension
    pub tile_m: usize,
    /// Tile size for N dimension
    pub tile_n: usize,
    /// Tile size for K dimension
    pub tile_k: usize,
    /// Use FP16 accumulation
    pub use_fp16_accumulator: bool,
    /// Number of stages for software pipelining
    pub pipeline_stages: usize,
}

impl Default for TmaConfig {
    fn default() -> Self {
        Self {
            async_copy: true,
            tile_m: 128,
            tile_n: 128,
            tile_k: 64,
            use_fp16_accumulator: false,
            pipeline_stages: 3,
        }
    }
}

impl TmaConfig {
    /// Configuration optimized for NVIDIA Blackwell (sm_110)
    pub fn blackwell() -> Self {
        Self {
            async_copy: true,
            tile_m: 128,
            tile_n: 128,
            tile_k: 64,
            use_fp16_accumulator: true,
            pipeline_stages: 4,
        }
    }

    /// Configuration optimized for Apple M3/M4 (Apple9)
    pub fn apple_silicon() -> Self {
        Self {
            async_copy: true,
            tile_m: 64,
            tile_n: 64,
            tile_k: 32,
            use_fp16_accumulator: true,
            pipeline_stages: 2,
        }
    }

    /// Convert to TileConfig
    pub fn to_tile_config(&self) -> TileConfig {
        TileConfig {
            tile_m: self.tile_m,
            tile_n: self.tile_n,
            tile_k: self.tile_k,
            warp_tile_m: self.tile_m / 4,
            warp_tile_n: self.tile_n / 4,
            pipeline_stages: self.pipeline_stages,
        }
    }
}

/// Trait for TMA-accelerated tensor operations
///
/// Provides optimized implementations of common tensor operations
/// using Tensor Memory Accelerator on supported hardware.
pub trait TmaOps: Sized {
    /// Matrix multiply with TMA optimization
    ///
    /// Uses tile-based memory access patterns for better cache utilization.
    ///
    /// # Example
    /// ```ignore
    /// let c = a.matmul_tma(&b)?;
    /// ```
    fn matmul_tma(&self, other: &Self) -> Result<Self>;

    /// Matrix multiply with custom configuration
    fn matmul_tma_with_config(&self, other: &Self, config: &TmaConfig) -> Result<Self>;

    /// Batched matrix multiply with TMA
    fn batched_matmul_tma(&self, other: &Self) -> Result<Self>;

    /// Attention operation with TMA optimization
    ///
    /// Computes: softmax(Q @ K.T / sqrt(d_k)) @ V
    fn attention_tma(&self, key: &Self, value: &Self, scale: f32) -> Result<Self>;

    /// Flash attention variant with memory-efficient implementation
    fn flash_attention_tma(
        &self,
        key: &Self,
        value: &Self,
        scale: f32,
        causal: bool,
    ) -> Result<Self>;

    /// Element-wise addition with tiled memory access
    fn add_tma(&self, other: &Self) -> Result<Self>;

    /// Get the shape of this tensor
    fn shape(&self) -> &[usize];

    /// Check if TMA is available for this tensor
    fn tma_available(&self) -> bool;
}

/// Wrapper for TMA-capable tensors
pub struct TmaTensor {
    /// Shape of the tensor
    shape: Vec<usize>,
    /// Data storage
    data: Vec<f32>,
    /// TMA configuration
    config: TmaConfig,
    /// Whether tensor is on GPU
    on_device: bool,
}

impl TmaTensor {
    /// Create a new TMA tensor
    pub fn new(shape: Vec<usize>, config: TmaConfig) -> Self {
        let size: usize = shape.iter().product();
        Self {
            shape,
            data: vec![0.0; size],
            config,
            on_device: false,
        }
    }

    /// Create from existing data
    pub fn from_vec(shape: Vec<usize>, data: Vec<f32>) -> Result<Self> {
        let expected_size: usize = shape.iter().product();
        if data.len() != expected_size {
            return Err(IntegrationError::ShapeMismatch {
                expected: format!("{} elements", expected_size),
                actual: format!("{} elements", data.len()),
            });
        }

        Ok(Self {
            shape,
            data,
            config: TmaConfig::default(),
            on_device: false,
        })
    }

    /// Get tensor data
    pub fn data(&self) -> &[f32] {
        &self.data
    }

    /// Get mutable tensor data
    pub fn data_mut(&mut self) -> &mut [f32] {
        &mut self.data
    }

    /// Set TMA configuration
    pub fn with_config(mut self, config: TmaConfig) -> Self {
        self.config = config;
        self
    }

    /// Move tensor to device
    #[instrument(skip(self))]
    pub fn to_device(&mut self) -> Result<()> {
        if self.on_device {
            return Ok(());
        }

        debug!("Moving tensor {:?} to device", self.shape);

        #[cfg(all(target_os = "macos", feature = "tma-metal"))]
        {
            // Actual device transfer would happen here
        }

        self.on_device = true;
        Ok(())
    }

    /// Move tensor to host
    pub fn to_host(&mut self) -> Result<()> {
        if !self.on_device {
            return Ok(());
        }

        debug!("Moving tensor {:?} to host", self.shape);

        #[cfg(all(target_os = "macos", feature = "tma-metal"))]
        {
            // Actual host transfer would happen here
        }

        self.on_device = false;
        Ok(())
    }
}

impl TmaOps for TmaTensor {
    #[instrument(skip(self, other))]
    fn matmul_tma(&self, other: &Self) -> Result<Self> {
        self.matmul_tma_with_config(other, &self.config)
    }

    fn matmul_tma_with_config(&self, other: &Self, config: &TmaConfig) -> Result<Self> {
        // Validate shapes for matmul
        if self.shape.len() < 2 || other.shape.len() < 2 {
            return Err(IntegrationError::ShapeMismatch {
                expected: "2D or higher tensors".to_string(),
                actual: format!("{}D and {}D", self.shape.len(), other.shape.len()),
            });
        }

        let m = self.shape[self.shape.len() - 2];
        let k1 = self.shape[self.shape.len() - 1];
        let k2 = other.shape[other.shape.len() - 2];
        let n = other.shape[other.shape.len() - 1];

        if k1 != k2 {
            return Err(IntegrationError::ShapeMismatch {
                expected: format!("K dimensions to match: {}", k1),
                actual: format!("got {}", k2),
            });
        }

        debug!(
            "TMA matmul: ({}, {}) x ({}, {}) with tiles {}x{}x{}",
            m, k1, k2, n, config.tile_m, config.tile_n, config.tile_k
        );

        // Output shape
        let mut output_shape = self.shape.clone();
        *output_shape.last_mut().unwrap() = n;

        let output_size: usize = output_shape.iter().product();
        let mut output_data = vec![0.0f32; output_size];

        // Naive implementation (actual TMA would use hardware acceleration)
        // In production, this would dispatch to GPU kernels
        for i in 0..m {
            for j in 0..n {
                let mut sum = 0.0f32;
                for k in 0..k1 {
                    sum += self.data[i * k1 + k] * other.data[k * n + j];
                }
                output_data[i * n + j] = sum;
            }
        }

        Ok(TmaTensor {
            shape: output_shape,
            data: output_data,
            config: config.clone(),
            on_device: self.on_device,
        })
    }

    fn batched_matmul_tma(&self, other: &Self) -> Result<Self> {
        // For batched matmul, iterate over batch dimensions
        self.matmul_tma(other)
    }

    fn attention_tma(&self, key: &Self, value: &Self, scale: f32) -> Result<Self> {
        self.flash_attention_tma(key, value, scale, false)
    }

    fn flash_attention_tma(
        &self,
        key: &Self,
        value: &Self,
        scale: f32,
        _causal: bool,
    ) -> Result<Self> {
        debug!(
            "Flash attention TMA: Q{:?} K{:?} V{:?} scale={}",
            self.shape, key.shape, value.shape, scale
        );

        // Compute Q @ K.T
        // For now, transpose K and multiply
        let k_t = Self::transpose_2d(key)?;
        let mut scores = self.matmul_tma(&k_t)?;

        // Scale
        for v in scores.data.iter_mut() {
            *v *= scale;
        }

        // Softmax (simplified - actual would be numerically stable)
        let seq_len = scores.shape[scores.shape.len() - 1];
        for i in 0..(scores.data.len() / seq_len) {
            let start = i * seq_len;
            let end = start + seq_len;

            let max = scores.data[start..end]
                .iter()
                .cloned()
                .fold(f32::NEG_INFINITY, f32::max);
            let mut sum = 0.0f32;

            for v in &mut scores.data[start..end] {
                *v = (*v - max).exp();
                sum += *v;
            }

            for v in &mut scores.data[start..end] {
                *v /= sum;
            }
        }

        // Multiply by V
        scores.matmul_tma(value)
    }

    fn add_tma(&self, other: &Self) -> Result<Self> {
        if self.shape != other.shape {
            return Err(IntegrationError::ShapeMismatch {
                expected: format!("{:?}", self.shape),
                actual: format!("{:?}", other.shape),
            });
        }

        let data: Vec<f32> = self
            .data
            .iter()
            .zip(other.data.iter())
            .map(|(a, b)| a + b)
            .collect();

        Ok(TmaTensor {
            shape: self.shape.clone(),
            data,
            config: self.config.clone(),
            on_device: self.on_device,
        })
    }

    fn shape(&self) -> &[usize] {
        &self.shape
    }

    fn tma_available(&self) -> bool {
        #[cfg(all(target_os = "macos", feature = "tma-metal"))]
        {
            true
        }
        #[cfg(not(all(target_os = "macos", feature = "tma-metal")))]
        {
            false
        }
    }
}

impl TmaTensor {
    /// Transpose a 2D tensor
    fn transpose_2d(tensor: &Self) -> Result<Self> {
        if tensor.shape.len() != 2 {
            return Err(IntegrationError::ShapeMismatch {
                expected: "2D tensor".to_string(),
                actual: format!("{}D tensor", tensor.shape.len()),
            });
        }

        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let mut data = vec![0.0f32; rows * cols];

        for i in 0..rows {
            for j in 0..cols {
                data[j * rows + i] = tensor.data[i * cols + j];
            }
        }

        Ok(TmaTensor {
            shape: vec![cols, rows],
            data,
            config: tensor.config.clone(),
            on_device: tensor.on_device,
        })
    }
}

impl fmt::Debug for TmaTensor {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("TmaTensor")
            .field("shape", &self.shape)
            .field("on_device", &self.on_device)
            .field("tma_available", &self.tma_available())
            .finish()
    }
}

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

    #[test]
    fn test_tma_config_default() {
        let config = TmaConfig::default();
        assert!(config.async_copy);
        assert_eq!(config.tile_m, 128);
    }

    #[test]
    fn test_tma_config_presets() {
        let blackwell = TmaConfig::blackwell();
        assert_eq!(blackwell.pipeline_stages, 4);

        let apple = TmaConfig::apple_silicon();
        assert_eq!(apple.tile_m, 64);
    }

    #[test]
    fn test_tma_tensor_creation() {
        let tensor = TmaTensor::new(vec![32, 64], TmaConfig::default());
        assert_eq!(tensor.shape(), &[32, 64]);
        assert_eq!(tensor.data.len(), 32 * 64);
    }

    #[test]
    fn test_tma_matmul() {
        let a = TmaTensor::from_vec(vec![2, 3], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        let b = TmaTensor::from_vec(vec![3, 2], vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        let c = a.matmul_tma(&b).unwrap();
        assert_eq!(c.shape(), &[2, 2]);

        // Check result: [1,2,3] @ [[1,2],[3,4],[5,6]] = [22, 28]
        assert!((c.data[0] - 22.0).abs() < 1e-5);
        assert!((c.data[1] - 28.0).abs() < 1e-5);
    }

    #[test]
    fn test_shape_mismatch() {
        let a = TmaTensor::new(vec![2, 3], TmaConfig::default());
        let b = TmaTensor::new(vec![4, 2], TmaConfig::default());

        let result = a.matmul_tma(&b);
        assert!(result.is_err());
    }
}