#![allow(unsafe_code)] // Required for safetensors f32-to-u8 slice reinterpretation
//! UnifiedTrainable trait implementation for DQN model
//!
//! This adapter wraps the DQN implementation to provide a unified
//! training interface compatible with the ML training orchestration system.

use ml_core::native_types::NativeDevice;
#[cfg(test)]
use ml_core::device::MlDevice;
use std::collections::HashMap as StdHashMap;

use crate::dqn::{Experience, DQN, DQNConfig};
use crate::training::unified_trainer::{
    checkpoint, CheckpointMetadata, TrainingMetrics, UnifiedTrainable,
};
use crate::MLError;

/// Adapter that wraps DQN to implement UnifiedTrainable trait
pub struct DQNTrainableAdapter {
    /// Underlying DQN model
    dqn: DQN,
    /// Configuration
    config: DQNConfig,
    /// Device descriptor
    device: NativeDevice,
    /// Current learning rate
    learning_rate: f64,
    /// Latest training metrics
    latest_metrics: TrainingMetrics,
    /// Current training step
    current_step: usize,
    /// Loss history for validation
    loss_history: Vec<f64>,
}

impl std::fmt::Debug for DQNTrainableAdapter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DQNTrainableAdapter")
            .field("config", &self.config)
            .field("device", &self.device)
            .field("learning_rate", &self.learning_rate)
            .field("latest_metrics", &self.latest_metrics)
            .field("current_step", &self.current_step)
            .field("loss_history_len", &self.loss_history.len())
            .finish_non_exhaustive()
    }
}

impl DQNTrainableAdapter {
    /// Create new DQN trainable adapter
    pub fn new(config: DQNConfig) -> Result<Self, MLError> {
        let learning_rate = config.learning_rate;
        let dqn = DQN::new(config.clone())?;
        let device = NativeDevice::Cuda(0);

        Ok(Self {
            dqn,
            config,
            device,
            learning_rate,
            latest_metrics: TrainingMetrics::default(),
            current_step: 0,
            loss_history: Vec::new(),
        })
    }

    /// Create new DQN trainable adapter on a specific device.
    ///
    /// Shares the caller's device instead of creating a new cuBLAS handle.
    #[cfg(test)]
    pub(crate) fn new_on_device(config: DQNConfig, ml_device: MlDevice) -> Result<Self, MLError> {
        let learning_rate = config.learning_rate;
        let dqn = DQN::new_on_device(config.clone(), ml_device)?;
        let device = NativeDevice::Cuda(0);

        Ok(Self {
            dqn,
            config,
            device,
            learning_rate,
            latest_metrics: TrainingMetrics::default(),
            current_step: 0,
            loss_history: Vec::new(),
        })
    }

    /// Get underlying DQN model (for action selection, etc.)
    pub fn model(&self) -> &DQN {
        &self.dqn
    }

    /// Get mutable reference to underlying DQN model
    pub fn model_mut(&mut self) -> &mut DQN {
        &mut self.dqn
    }

    /// Store experience in replay buffer
    pub fn store_experience(&mut self, experience: Experience) -> Result<(), MLError> {
        self.dqn.store_experience(experience)
    }

    /// Train on a batch of experiences
    ///
    /// This is a convenience method that combines forward, backward, and optimizer_step
    pub fn train_batch(&mut self, experiences: Vec<Experience>) -> Result<f64, MLError> {
        let result = self.dqn.train_step(Some(super::replay_buffer_type::BatchSample::uniform(experiences)))?;
        let stream = self.dqn.cuda_stream();
        let loss_host = result.loss_gpu
            .to_host(stream)?;
        let loss = loss_host.first().copied().unwrap_or(0.0);
        self.current_step += 1;
        self.loss_history.push(loss as f64);
        Ok(loss as f64)
    }

    /// Get epsilon value for exploration
    pub fn epsilon(&self) -> f32 {
        self.dqn.get_epsilon()
    }

    /// Check if ready for training (enough replay buffer samples)
    pub fn can_train(&self) -> bool {
        self.dqn.can_train()
    }
}

impl UnifiedTrainable for DQNTrainableAdapter {
    fn model_type(&self) -> &str {
        "DQN"
    }

    fn device_name(&self) -> String {
        self.device.to_string()
    }

    fn forward_loss(&mut self, _input: &[f32], _target: &[f32]) -> Result<f64, MLError> {
        // DQN uses train_step() for full training iteration (RL, not supervised)
        // This placeholder computes a dummy loss for the unified training interface
        Ok(self.loss_history.last().copied().unwrap_or(0.0))
    }

    fn backward(&mut self, _loss_value: f64) -> Result<f64, MLError> {
        // DQN backward pass is integrated into train_step()
        // Return last known gradient norm
        Ok(self.latest_metrics.grad_norm.unwrap_or(0.0))
    }

    fn optimizer_step(&mut self) -> Result<(), MLError> {
        // The DQN's train_step method already handles optimizer step internally
        Ok(())
    }

    fn zero_grad(&mut self) -> Result<(), MLError> {
        // Gradients are automatically zeroed in DQN's train_step
        Ok(())
    }

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

    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError> {
        self.learning_rate = lr;
        tracing::debug!("DQN learning rate change requested but dynamic LR not yet supported");
        Ok(())
    }

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

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

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

        metrics.learning_rate = self.learning_rate;

        // Add DQN-specific metrics
        metrics
            .custom_metrics
            .insert("epsilon".to_owned(), self.dqn.get_epsilon() as f64);
        metrics.custom_metrics.insert(
            "training_steps".to_owned(),
            self.dqn.get_training_steps() as f64,
        );
        if let Ok(buffer_size) = self.dqn.get_replay_buffer_size() {
            metrics
                .custom_metrics
                .insert("replay_buffer_size".to_owned(), buffer_size as f64);
        }

        metrics
    }

    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError> {
        // Create checkpoint metadata
        let metadata = CheckpointMetadata {
            model_type: "DQN".to_owned(),
            version: "1.0.0".to_owned(),
            epoch: 0, // DQN doesn't use epochs
            step: self.current_step,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&self.config).map_err(|e| {
                MLError::SerializationError {
                    reason: format!("Failed to serialize config: {}", e),
                }
            })?,
            metrics: self.collect_metrics(),
        };

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

        // Save model weights via GpuVarStore export
        let safetensors_path = format!("{}.safetensors", checkpoint_path);
        let vars = self.dqn.get_q_network_vars();
        let host_data = vars.export_to_host().map_err(|e| { // gpu-exit: checkpoint save
            MLError::CheckpointError(format!("Failed to export vars: {}", e))
        })?;

        // Use safetensors crate to write (host data is BTreeMap<name -> (shape, Vec<f32>)>)
        let mut tensor_map = StdHashMap::new();
        for (name, (shape, data)) in &host_data {
            let bytes: &[u8] = unsafe {
                std::slice::from_raw_parts(
                    data.as_ptr() as *const u8,
                    data.len() * std::mem::size_of::<f32>(),
                )
            };
            let view = safetensors::tensor::TensorView::new(
                safetensors::Dtype::F32,
                shape.clone(),
                bytes,
            ).map_err(|e| MLError::CheckpointError(format!("TensorView for {}: {}", name, e)))?;
            tensor_map.insert(name.as_str(), view);
        }

        let arch_metadata = self.config.checkpoint_metadata();
        safetensors::serialize_to_file(
            tensor_map,
            Some(arch_metadata),
            std::path::Path::new(&safetensors_path),
        )
        .map_err(|e| MLError::CheckpointError(format!("Failed to save safetensors: {}", e)))?;

        tracing::info!(
            "Saved DQN checkpoint to {} (step {})",
            checkpoint_path,
            self.current_step
        );

        Ok(checkpoint_path.to_string())
    }

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

        // Validate model type
        if metadata.model_type != "DQN" {
            return Err(MLError::CheckpointError(format!(
                "Invalid model type in checkpoint: expected DQN, got {}",
                metadata.model_type
            )));
        }

        // Load model weights from safetensors with architecture validation
        let safetensors_path = format!("{}.safetensors", checkpoint_path);
        let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| {
            MLError::CheckpointError(format!("Failed to read safetensors file: {}", e))
        })?;

        // Validate architecture hash before loading any weights
        let (_header_size, st_metadata) =
            safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| {
                MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e))
            })?;
        self.config
            .validate_checkpoint_metadata(st_metadata.metadata())?;

        // Load tensors from safetensors and import into VarStore
        let st = safetensors::SafeTensors::deserialize(&raw_bytes).map_err(|e| {
            MLError::CheckpointError(format!("Failed to deserialize safetensors: {}", e))
        })?;

        let mut host_data = std::collections::BTreeMap::new();
        for (name, tensor_view) in st.tensors() {
            let shape: Vec<usize> = tensor_view.shape().to_vec(); // cpu-side shape clone
            let raw_bytes = tensor_view.data();
            let f32_data: Vec<f32> = unsafe {
                std::slice::from_raw_parts(
                    raw_bytes.as_ptr() as *const f32,
                    raw_bytes.len() / std::mem::size_of::<f32>(),
                )
            }.to_vec(); // cpu-side bytes→f32 clone
            host_data.insert(name.clone(), (shape, f32_data));
        }

        // Import checkpoint data into the DQN's VarStore via direct CudaSlice writes.
        let vars = self.dqn.get_q_network_vars();
        let stream = self.dqn.cuda_stream();
        for (name, (_shape, f32_data)) in &host_data {
            if let Some(param) = vars.get(name) {
                if param.data.len() == f32_data.len() {
                    // Direct HtoD memcpy into the existing CudaSlice
                    stream.memcpy_htod(f32_data, &mut param.data.clone())
                        .map_err(|e| MLError::CheckpointError(
                            format!("Failed to upload checkpoint param '{}': {}", name, e)
                        ))?;
                }
            }
        }

        // Restore training state
        self.current_step = metadata.step;
        self.latest_metrics = metadata.metrics.clone();

        tracing::info!(
            "Loaded DQN checkpoint from {} (step {})",
            checkpoint_path,
            metadata.step
        );

        Ok(metadata)
    }
}

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

    static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

    fn cuda_device() -> MlDevice {
        SHARED_CUDA
            .get_or_init(|| MlDevice::cuda(0).expect("CUDA device required"))
            .clone()
    }

    #[test]
    fn test_dqn_adapter_creation() -> anyhow::Result<()> {
        let config = DQNConfig::emergency_safe_defaults();
        let adapter = DQNTrainableAdapter::new_on_device(config, cuda_device())?;

        assert_eq!(adapter.model_type(), "DQN");
        assert_eq!(adapter.get_step(), 0);

        Ok(())
    }

    #[test]
    fn test_dqn_adapter_metrics() -> anyhow::Result<()> {
        let config = DQNConfig::emergency_safe_defaults();
        let adapter = DQNTrainableAdapter::new_on_device(config, cuda_device())?;

        let metrics = adapter.collect_metrics();
        assert!(metrics.custom_metrics.contains_key("epsilon"));
        assert!(metrics.custom_metrics.contains_key("training_steps"));

        Ok(())
    }

    #[test]
    fn test_dqn_adapter_checkpoint_metadata() -> anyhow::Result<()> {
        let config = DQNConfig::emergency_safe_defaults();
        let adapter = DQNTrainableAdapter::new_on_device(config, cuda_device())?;

        // Test metadata serialization
        let metadata = CheckpointMetadata {
            model_type: "DQN".to_owned(),
            version: "1.0.0".to_owned(),
            epoch: 0,
            step: 100,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::to_value(&adapter.config)?,
            metrics: adapter.collect_metrics(),
        };

        let json = serde_json::to_string(&metadata)?;
        let deserialized: CheckpointMetadata = serde_json::from_str(&json)?;

        assert_eq!(deserialized.model_type, "DQN");
        assert_eq!(deserialized.step, 100);

        Ok(())
    }
}
