//! xLSTM block -- pre-LayerNorm + sLSTM/mLSTM cell + residual connection.
//!
//! Each block wraps either an sLSTM or mLSTM cell with:
//! 1. Pre-LayerNorm on the input
//! 2. Cell forward pass
//! 3. Residual connection (if input dim == hidden dim)

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{gpu_add, gpu_layer_norm, gpu_ones, GpuTensor};
use super::slstm::SLSTMCell;
use super::mlstm::MLSTMCell;

/// Type of cell in a block.
#[derive(Debug)]
enum CellType {
    #[allow(clippy::upper_case_acronyms)]
    SLSTM(SLSTMCell),
    #[allow(clippy::upper_case_acronyms)]
    MLSTM(MLSTMCell),
}

/// An xLSTM block: `LayerNorm` -> cell -> residual.
#[derive(Debug)]
pub struct XLSTMBlock {
    ln_gamma: GpuTensor,
    ln_beta: GpuTensor,
    cell: CellType,
    hidden_dim: usize,
    /// Flat size of cell state for mLSTM (`num_heads` * `head_dim^2`), or `hidden_dim` for sLSTM.
    c_flat_size: usize,
    has_residual: bool,
}

impl XLSTMBlock {
    /// Create an sLSTM block.
    pub fn new_slstm(
        input_dim: usize,
        hidden_dim: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let ln_gamma = gpu_ones(&[hidden_dim], stream)?;
        let ln_beta = GpuTensor::zeros(&[hidden_dim], stream)?;
        let cell = SLSTMCell::new(input_dim, hidden_dim, stream)?;

        Ok(Self {
            ln_gamma,
            ln_beta,
            cell: CellType::SLSTM(cell),
            hidden_dim,
            c_flat_size: hidden_dim,
            has_residual: input_dim == hidden_dim,
        })
    }

    /// Create an mLSTM block.
    pub fn new_mlstm(
        input_dim: usize,
        hidden_dim: usize,
        num_heads: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<Self, MLError> {
        let ln_gamma = gpu_ones(&[hidden_dim], stream)?;
        let ln_beta = GpuTensor::zeros(&[hidden_dim], stream)?;
        let head_dim = hidden_dim / num_heads;
        let cell = MLSTMCell::new(input_dim, hidden_dim, num_heads, stream)?;

        Ok(Self {
            ln_gamma,
            ln_beta,
            cell: CellType::MLSTM(cell),
            hidden_dim,
            c_flat_size: num_heads * head_dim * head_dim,
            has_residual: input_dim == hidden_dim,
        })
    }

    /// Forward pass for a single timestep.
    pub fn forward(
        &self,
        x: &GpuTensor,
        state: Option<(&GpuTensor, &GpuTensor)>,
    ) -> Result<(GpuTensor, GpuTensor), MLError> {
        // Pre-LayerNorm
        let x_norm = gpu_layer_norm(x, &self.ln_gamma, &self.ln_beta, 1e-5)?;

        let (h, c) = match &self.cell {
            CellType::SLSTM(cell) => cell.forward(&x_norm, state)?,
            CellType::MLSTM(cell) => cell.forward(&x_norm, state)?,
        };

        // Residual connection
        let h_out = if self.has_residual {
            gpu_add(&h, x)?
        } else {
            h
        };

        Ok((h_out, c))
    }

    /// Hidden dimension of this block.
    pub const fn hidden_dim(&self) -> usize {
        self.hidden_dim
    }

    /// Flat cell state size.
    pub const fn c_flat_size(&self) -> usize {
        self.c_flat_size
    }
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_slstm_block_forward() {
        let stream = test_stream();
        let block = XLSTMBlock::new_slstm(16, 16, &stream).unwrap();

        let x = GpuTensor::randn(&[2, 16], 1.0, &stream).unwrap();
        let (h, c) = block.forward(&x, None).unwrap();
        assert_eq!(h.shape, vec![2, 16]);
        assert_eq!(c.shape, vec![2, 16]);
    }

    #[test]
    fn test_mlstm_block_forward() {
        let stream = test_stream();
        let block = XLSTMBlock::new_mlstm(16, 16, 2, &stream).unwrap();

        let x = GpuTensor::randn(&[2, 16], 1.0, &stream).unwrap();
        let (h, c) = block.forward(&x, None).unwrap();
        assert_eq!(h.shape, vec![2, 16]);
        // mLSTM c_flat: 2 heads * 8 * 8 = 128
        assert_eq!(c.shape, vec![2, 128]);
    }

    #[test]
    fn test_block_residual_connection() {
        let stream = test_stream();
        let block = XLSTMBlock::new_slstm(16, 16, &stream).unwrap();
        assert!(block.has_residual);
    }

    #[test]
    fn test_block_no_residual_when_dims_differ() {
        let stream = test_stream();
        let block = XLSTMBlock::new_slstm(8, 16, &stream).unwrap();
        assert!(!block.has_residual);
    }
}
