//! UnifiedTrainable adapter for MAMBA-2
//!
//! Wraps `Mamba2SSM` with the `UnifiedTrainable` trait for standardized training
//! orchestration. Uses a local wrapper struct to satisfy the orphan rule (trait
//! in ml-core, type in ml-supervised).

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

use cudarc::driver::CudaStream;

use ml_supervised::gpu_tensor::GpuTensor as Mamba2Tensor;

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

/// Wrapper adapter for Mamba2SSM that implements UnifiedTrainable.
///
/// Required because `UnifiedTrainable` is defined in ml-core and `Mamba2SSM`
/// in ml-supervised -- orphan rule prevents direct `impl` in ml.
pub struct Mamba2TrainableAdapter {
    /// Underlying MAMBA-2 model
    pub model: Mamba2SSM,
    /// Cached input from forward_loss for SPSA backward
    cached_input: Option<Mamba2Tensor>,
    /// Cached target from forward_loss for SPSA backward
    cached_target: Option<Mamba2Tensor>,
    /// Cached loss tensor from forward_loss for SPSA backward
    cached_loss: Option<Mamba2Tensor>,
}

impl Mamba2TrainableAdapter {
    /// Create a new adapter wrapping a Mamba2SSM model.
    ///
    /// Accepts `&Arc<CudaStream>` since Mamba2SSM is GPU-native.
    pub fn new(config: Mamba2Config, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let model = Mamba2SSM::new(config, stream)?;
        Ok(Self { model, cached_input: None, cached_target: None, cached_loss: None })
    }

    /// Create from NativeDevice (opens a stream from ordinal).
    pub fn from_native_device(config: Mamba2Config, device: &ml_core::native_types::NativeDevice) -> Result<Self, MLError> {
        let ordinal = match device {
            ml_core::native_types::NativeDevice::Cuda(ord) => *ord,
            _ => 0,
        };
        let ctx = cudarc::driver::CudaContext::new(ordinal)
            .map_err(|e| MLError::DeviceError(format!("CUDA context: {e}")))?;
        let stream = ctx
            .new_stream()
            .map_err(|e| MLError::DeviceError(format!("CUDA stream: {e}")))?;
        Self::new(config, &stream)
    }

    /// Create adapter from an existing Mamba2SSM
    pub fn from_model(model: Mamba2SSM) -> Self {
        Self { model, cached_input: None, cached_target: None, cached_loss: None }
    }
}

impl std::fmt::Debug for Mamba2TrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Mamba2TrainableAdapter")
            .field("model_type", &"MAMBA-2")
            .field("step", &self.model.step_count)
            .finish()
    }
}

impl UnifiedTrainable for Mamba2TrainableAdapter {
    fn model_type(&self) -> &str {
        "MAMBA-2"
    }

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

    fn forward_loss(&mut self, input: &[f32], target: &[f32]) -> Result<f64, MLError> {
        use ml_supervised::gpu_tensor::{GpuTensor, gpu_sub, gpu_sqr, gpu_mean_all};
        let stream = &self.model.stream;
        let x = GpuTensor::from_vec(input.to_vec(), &[1, input.len()], stream)?; // cpu-side: initial data upload
        let t = GpuTensor::from_vec(target.to_vec(), &[1, target.len()], stream)?; // cpu-side: initial data upload
        let output = self.model.forward(&x)?;
        let diff = gpu_sub(&output, &t)?;
        let sq = gpu_sqr(&diff)?;
        let loss_tensor = self.model.compute_loss(&output, &t)?;
        let loss = gpu_mean_all(&sq)?; // gpu-exit: 1 scalar loss

        // Cache input/target on GPU for SPSA backward
        self.cached_input = Some(x);
        self.cached_target = Some(t);
        self.cached_loss = Some(loss_tensor);

        Ok(loss as f64)
    }

    fn backward(&mut self, _loss_value: f64) -> Result<f64, MLError> {
        // Delegate to Mamba2SSM's SPSA backward_pass (proper gradient estimation)
        let input = self.cached_input.take().ok_or_else(|| {
            MLError::TrainingError("backward() called without prior forward_loss()".to_owned())
        })?;
        let target = self.cached_target.take().ok_or_else(|| {
            MLError::TrainingError("backward() missing cached target".to_owned())
        })?;
        let loss = self.cached_loss.take().ok_or_else(|| {
            MLError::TrainingError("backward() missing cached loss".to_owned())
        })?;

        self.model.backward_pass(&loss, &input, &target)?;
        self.model.optimizer_step()?;
        self.model.step_count += 1;

        Ok(0.0) // SPSA doesn't produce a meaningful gradient norm
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        // Optimizer step is fused into backward() for SPSA efficiency
        Ok(())
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        self.model.gradients.clear();
        self.cached_input = None;
        self.cached_target = None;
        self.cached_loss = None;
        Ok(())
    }

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

    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        if lr <= 0.0 || lr > 1.0 {
            return Err(MLError::ValidationError {
                message: format!("Invalid learning rate: {}. Must be in range (0.0, 1.0]", lr),
            });
        }
        self.model.config.learning_rate = lr;
        Ok(())
    }

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

    fn collect_metrics(&self) -> TrainingMetrics {
        let perf_metrics = self.model.get_performance_metrics();
        let mut custom_metrics = HashMap::new();
        for (key, value) in perf_metrics.iter() {
            custom_metrics.insert(key.clone(), *value);
        }

        TrainingMetrics {
            loss: self
                .model
                .metadata
                .training_history
                .back()
                .map(|e| e.loss)
                .unwrap_or(0.0),
            val_loss: self
                .model
                .metadata
                .training_history
                .back()
                .map(|e| e.loss),
            accuracy: self
                .model
                .metadata
                .training_history
                .back()
                .map(|e| e.accuracy),
            learning_rate: self.model.config.learning_rate,
            grad_norm: None,
            custom_metrics,
        }
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        let metadata = CheckpointMetadata {
            model_type: "MAMBA-2".to_owned(),
            version: self.model.metadata.version.clone(),
            epoch: self.model.metadata.training_history.len(),
            step: self.model.step_count,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&self.model.config)
                .map_err(|e| MLError::ModelError(format!("Failed to serialize config: {}", e)))?,
            metrics: self.collect_metrics(),
        };

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

        Ok(format!("{}.json", checkpoint_path))
    }

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

        self.model.step_count = metadata.step;
        self.model.is_trained = true;

        Ok(metadata)
    }
}

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

    #[test]
    fn test_mamba2_adapter_creation() -> anyhow::Result<()> {
        let config = Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            batch_size: 4,
            seq_len: 32,
            ..Default::default()
        };
        let adapter = Mamba2TrainableAdapter::from_native_device(config, &NativeDevice::Cuda(0))?;

        assert_eq!(adapter.model_type(), "MAMBA-2");
        assert!(adapter.device_name().contains("Cuda"));
        assert_eq!(adapter.get_step(), 0);
        assert!(adapter.get_learning_rate() > 0.0);

        Ok(())
    }

    #[test]
    fn test_mamba2_learning_rate_validation() -> anyhow::Result<()> {
        let config = Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            ..Default::default()
        };
        let mut adapter = Mamba2TrainableAdapter::from_native_device(config, &NativeDevice::Cuda(0))?;

        assert!(adapter.set_learning_rate(1e-3).is_ok());
        assert_eq!(adapter.get_learning_rate(), 1e-3);

        assert!(adapter.set_learning_rate(0.0).is_err());
        assert!(adapter.set_learning_rate(-0.1).is_err());
        assert!(adapter.set_learning_rate(1.5).is_err());

        Ok(())
    }

    #[test]
    fn test_mamba2_metrics_collection() -> anyhow::Result<()> {
        let config = Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            ..Default::default()
        };
        let adapter = Mamba2TrainableAdapter::from_native_device(config, &NativeDevice::Cuda(0))?;

        let metrics = adapter.collect_metrics();
        assert!(metrics.loss >= 0.0);
        assert_eq!(metrics.learning_rate, adapter.get_learning_rate());
        assert!(!metrics.custom_metrics.is_empty());
        assert!(metrics.custom_metrics.contains_key("total_inferences"));
        assert!(metrics.custom_metrics.contains_key("model_parameters"));

        Ok(())
    }

    #[test]
    fn test_mamba2_checkpoint_roundtrip() -> anyhow::Result<()> {
        let config = Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            batch_size: 4,
            seq_len: 32,
            ..Default::default()
        };
        let ml_dev = ml_core::device::MlDevice::cuda(0)?;
        let stream = std::sync::Arc::clone(ml_dev.cuda_stream()?);
        let adapter = Mamba2TrainableAdapter::new(config.clone(), &stream)?;

        let temp_dir = tempfile::tempdir()?;
        let checkpoint_path = temp_dir.path().join("mamba2_test_checkpoint");
        let checkpoint_path_str = checkpoint_path.to_str().ok_or_else(|| {
            MLError::ModelError("Invalid path".to_owned())
        })?;

        let checkpoint_str = adapter.save_checkpoint(checkpoint_path_str)?;
        assert!(checkpoint_str.ends_with(".json"));

        let metadata_path = format!("{}.json", checkpoint_path_str);
        assert!(
            std::path::Path::new(&metadata_path).exists(),
            "Metadata file should exist at {}",
            metadata_path
        );

        let mut loaded_adapter = Mamba2TrainableAdapter::from_native_device(config.clone(), &NativeDevice::Cuda(0))?;
        loaded_adapter.load_checkpoint(checkpoint_path_str)?;
        assert!(loaded_adapter.model.is_trained);

        Ok(())
    }

    #[test]
    fn test_mamba2_zero_grad() -> anyhow::Result<()> {
        let config = Mamba2Config {
            d_model: 64,
            d_state: 16,
            num_layers: 2,
            ..Default::default()
        };
        let mut adapter = Mamba2TrainableAdapter::from_native_device(config, &NativeDevice::Cuda(0))?;

        adapter.zero_grad()?;
        assert!(adapter.model.gradients.is_empty());

        Ok(())
    }
}
