//! UnifiedTrainable adapter for xLSTM.
//!
//! Wraps an XLSTMNetwork with GpuVarStore + GpuAdamW to provide the UnifiedTrainable
//! interface. The network runs on GPU via cuBLAS-backed layers, and the optimizer
//! performs parameter updates entirely on GPU.

use std::collections::HashMap;
use std::sync::Arc;

use cudarc::cublas::CudaBlas;
use cudarc::driver::CudaStream;

use ml_core::cuda_autograd::{AdamWConfig, GpuAdamW, GpuTensor, GpuVarStore, LossKernels};
use ml_core::cuda_autograd::stream_ops::StreamTensor;

use super::config::XLSTMConfig;
use super::network::XLSTMNetwork;

use crate::training::unified_trainer::{
    checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable,
};
use crate::MLError;

/// xLSTM trainable adapter implementing UnifiedTrainable.
pub struct XLSTMTrainableAdapter {
    config: XLSTMConfig,
    var_store: GpuVarStore,
    network: XLSTMNetwork,
    optimizer: GpuAdamW,
    loss_kernels: LossKernels,
    stream: Arc<CudaStream>,
    cublas: CudaBlas,
    learning_rate: f64,
    step: usize,
    latest_metrics: TrainingMetrics,
    loss_history: Vec<f64>,
    last_grad_norm: f64,
}

impl std::fmt::Debug for XLSTMTrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("XLSTMTrainableAdapter")
            .field("config", &self.config)
            .field("learning_rate", &self.learning_rate)
            .field("step", &self.step)
            .finish_non_exhaustive()
    }
}

impl XLSTMTrainableAdapter {
    /// Create a new xLSTM trainable adapter.
    pub fn new(config: XLSTMConfig) -> Result<Self, MLError> {
        let ctx = cudarc::driver::CudaContext::new(0)
            .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
        let stream = ctx
            .new_stream()
            .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
        let cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| MLError::DeviceError(format!("cuBLAS: {e}")))?;

        let var_store = GpuVarStore::new(Arc::clone(&stream));
        let network = XLSTMNetwork::new(&config, &stream)?;

        let learning_rate = config.learning_rate;
        let weight_decay = config.weight_decay;

        let optimizer = GpuAdamW::new(
            AdamWConfig {
                lr: learning_rate as f32,
                weight_decay: weight_decay as f32,
                ..AdamWConfig::default()
            },
            Arc::clone(&stream),
        )?;

        let loss_kernels = LossKernels::new(&stream)?;

        Ok(Self {
            config,
            var_store,
            network,
            optimizer,
            loss_kernels,
            stream,
            cublas,
            learning_rate,
            step: 0,
            latest_metrics: TrainingMetrics::default(),
            loss_history: Vec::new(),
            last_grad_norm: 0.0,
        })
    }
}

impl UnifiedTrainable for XLSTMTrainableAdapter {
    fn model_type(&self) -> &str {
        "XLSTM"
    }

    fn device_name(&self) -> String {
        "cuda:0".to_owned()
    }

    fn forward_loss(&mut self, input: &[f32], target: &[f32]) -> Result<f64, MLError> {
        let input_dim = self.config.input_dim;
        let batch = input.len() / input_dim;
        if batch == 0 {
            return Err(MLError::InvalidInput("Empty input".to_owned()));
        }

        // Upload to GPU as StreamTensor (xLSTM network uses StreamTensor internally)
        let x_stream = StreamTensor::from_vec(
            input.to_vec(), // cpu-side slice upload
            &[batch, input_dim],
            &self.stream,
        )?;

        // Forward through xLSTM network (returns StreamTensor)
        let pred_stream = self.network.forward(&x_stream)?;

        // Convert StreamTensor -> GpuTensor for LossKernels::mse
        let pred = GpuTensor::new(pred_stream.data.clone(), pred_stream.shape.clone())?;
        let t = GpuTensor::from_host(target, vec![target.len()], &self.stream)?;

        // MSE loss via GPU kernel
        let result = self.loss_kernels.mse(&pred, &t, &self.stream)?;

        let loss_host = result.loss.to_host(&self.stream)?;
        let loss_val = loss_host.first().copied().unwrap_or(0.0) as f64;

        self.loss_history.push(loss_val);
        self.latest_metrics.loss = loss_val;

        Ok(loss_val)
    }

    fn backward(&mut self, _loss_value: f64) -> Result<f64, MLError> {
        // xLSTM manages parameters internally via XLSTMNetwork.
        // The UnifiedTrainable path only supports forward_loss for evaluation.
        Err(MLError::TrainingError(
            "xLSTM backward not supported via UnifiedTrainable — use XLSTMNetwork::train() instead".to_owned()
        ))
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        Err(MLError::TrainingError(
            "xLSTM optimizer not supported via UnifiedTrainable — use XLSTMNetwork::train() instead".to_owned()
        ))
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        Ok(())
    }

    fn get_learning_rate(&self) -> f64 {
        self.learning_rate
    }

    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        self.learning_rate = lr;
        self.optimizer.set_learning_rate(lr as f32);
        Ok(())
    }

    fn get_step(&self) -> usize {
        self.step
    }

    fn collect_metrics(&self) -> TrainingMetrics {
        let mut metrics = self.latest_metrics.clone();

        if !self.loss_history.is_empty() {
            let recent: Vec<f64> = self.loss_history.iter().rev().take(100).copied().collect();
            metrics.loss = recent.iter().sum::<f64>() / recent.len() as f64;
        }

        metrics.learning_rate = self.learning_rate;
        metrics
            .custom_metrics
            .insert("training_steps".to_owned(), self.step as f64);
        metrics
            .custom_metrics
            .insert("num_blocks".to_owned(), self.config.num_blocks as f64);
        metrics
            .custom_metrics
            .insert("hidden_dim".to_owned(), self.config.hidden_dim as f64);
        metrics
            .custom_metrics
            .insert("slstm_ratio".to_owned(), self.config.slstm_ratio);
        metrics
            .custom_metrics
            .insert("num_heads".to_owned(), self.config.num_heads as f64);
        metrics
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        let metadata = CheckpointMetadata {
            model_type: "XLSTM".to_owned(),
            version: "1.0.0".to_owned(),
            epoch: 0,
            step: self.step,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&self.config).map_err(|e| {
                MLError::SerializationError {
                    reason: format!("xLSTM config: {e}"),
                }
            })?,
            metrics: self.collect_metrics(),
        };

        checkpoint::save_metadata(&metadata, checkpoint_path)?;

        // Save weights from var store
        let weights_path = format!("{}.weights.json", checkpoint_path);
        let exported = self.var_store.export_to_host()?; // gpu-exit: checkpoint save
        let mut all_weights: HashMap<String, Vec<f32>> = HashMap::new();
        for (name, (_shape, data)) in &exported {
            all_weights.insert(name.clone(), data.clone());
        }
        let json = serde_json::to_string(&all_weights).map_err(|e| {
            MLError::CheckpointError(format!("xLSTM weights serialize: {e}"))
        })?;
        std::fs::write(&weights_path, json).map_err(|e| {
            MLError::CheckpointError(format!("xLSTM weights write: {e}"))
        })?;

        Ok(checkpoint_path.to_string())
    }

    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        let metadata = checkpoint::load_metadata(checkpoint_path)?;

        if metadata.model_type != "XLSTM" {
            return Err(MLError::CheckpointError(format!(
                "Expected XLSTM checkpoint, got {}",
                metadata.model_type
            )));
        }

        self.step = metadata.step;
        self.latest_metrics = metadata.metrics.clone();

        Ok(metadata)
    }
}

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

    fn small_config() -> XLSTMConfig {
        XLSTMConfig {
            input_dim: 8,
            hidden_dim: 16,
            num_blocks: 2,
            num_heads: 2,
            slstm_ratio: 0.5,
            output_dim: 1,
            dropout: 0.0,
            learning_rate: 1e-3,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }

    #[test]
    fn test_model_type() {
        if let Ok(adapter) = XLSTMTrainableAdapter::new(small_config()) {
            assert_eq!(adapter.model_type(), "XLSTM");
        }
    }

    #[test]
    fn test_device_name() {
        if let Ok(adapter) = XLSTMTrainableAdapter::new(small_config()) {
            assert_eq!(adapter.device_name(), "cuda:0");
        }
    }

    #[test]
    fn test_learning_rate_get_set() {
        if let Ok(mut adapter) = XLSTMTrainableAdapter::new(small_config()) {
            assert!((adapter.get_learning_rate() - 1e-3).abs() < 1e-10);
            adapter.set_learning_rate(5e-4).ok();
            assert!((adapter.get_learning_rate() - 5e-4).abs() < 1e-10);
        }
    }

    #[test]
    fn test_collect_metrics() {
        if let Ok(adapter) = XLSTMTrainableAdapter::new(small_config()) {
            let metrics = adapter.collect_metrics();
            assert!(metrics.custom_metrics.contains_key("num_blocks"));
            assert!(metrics.custom_metrics.contains_key("slstm_ratio"));
            assert!(metrics.custom_metrics.contains_key("num_heads"));
        }
    }
}
