//! Production DQN model — loads from safetensors checkpoints with 5-exposure DQN actions.
//!
//! Config is read from checkpoint metadata (hidden_dims, Rainbow flags, etc.) so
//! production automatically matches whatever architecture was trained.

use std::sync::Arc;
use tokio::sync::RwLock;
use ml::{Features, MLModel, ModelMetadata, ModelPrediction, ModelType};

#[allow(missing_debug_implementations)]
pub(crate) struct DQNModel {
    model_id: String,
    dqn: Arc<RwLock<ml::dqn::DQN>>,
    feature_count: usize,
    loaded: bool,
}

impl std::fmt::Debug for DQNModel {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("DQNModel")
            .field("model_id", &self.model_id)
            .field("feature_count", &self.feature_count)
            .field("loaded", &self.loaded)
            .finish_non_exhaustive()
    }
}

impl DQNModel {
    pub fn from_checkpoint(
        model_id: String,
        checkpoint_path: &std::path::Path,
    ) -> ml::MLResult<Self> {
        use ml::dqn::{DQNConfig, DQN};

        // Resolve safetensors path
        let st_path = if checkpoint_path.extension().is_some_and(|ext| ext == "safetensors") {
            checkpoint_path.to_path_buf()
        } else {
            let candidate = checkpoint_path.with_extension("safetensors");
            if candidate.exists() {
                candidate
            } else {
                return Err(ml::MLError::CheckpointError(format!(
                    "No safetensors checkpoint found at {} \
                     (legacy JSON not supported for DQN production inference)",
                    checkpoint_path.display()
                )));
            }
        };

        // Read architecture from checkpoint metadata — no hardcoded dims
        let mut config = DQNConfig::from_safetensors_file(&st_path)?;
        // Disable ALL exploration for production inference
        config.epsilon_start = 0.0;
        config.epsilon_end = 0.0;
        config.warmup_steps = 0; // No random warmup period
        config.noisy_epsilon_floor = 0.0; // No noisy-net exploration floor
        // NoisyNets always enabled — noisy_epsilon_floor=0.0 above disables exploration noise
        config.use_count_bonus = false; // No UCB exploration bonus on Q-values
        tracing::info!(
            "DQN '{}' config from checkpoint: state_dim={}, num_actions={}, \
             hidden_dims={:?}, dueling={}, iqn={}",
            model_id, config.state_dim, config.num_actions, config.hidden_dims,
            true, config.use_iqn
        );

        let state_dim = config.state_dim;
        let mut dqn = DQN::new(config)?;
        dqn.load_from_safetensors(&st_path.to_string_lossy())?;

        tracing::info!("DQN model '{}' loaded from {}", model_id, st_path.display());

        Ok(Self {
            model_id,
            dqn: Arc::new(RwLock::new(dqn)),
            feature_count: state_dim,
            loaded: true,
        })
    }
}

#[async_trait::async_trait]
impl MLModel for DQNModel {
    fn name(&self) -> &str {
        &self.model_id
    }

    fn model_type(&self) -> ModelType {
        ModelType::DQN
    }

    async fn predict(&self, features: &Features) -> ml::MLResult<ModelPrediction> {
        if !self.loaded {
            return Err(ml::MLError::InferenceError(format!(
                "DQN model '{}' has not been loaded -- refusing to predict with random weights",
                self.model_id
            )));
        }

        // Read lock: select_action_inference(&self) is a pure forward pass with no
        // mutations (no epsilon exploration, no step counters, no action tracking).
        // This allows concurrent predict() calls from multiple tokio tasks.
        let dqn = self.dqn.read().await;

        let dim = self.feature_count;
        let state_vec: Vec<f32> = if features.values.len() >= dim {
            features.values.iter().take(dim).map(|&v| v as f32).collect()
        } else {
            let mut padded: Vec<f32> = features.values.iter().map(|&v| v as f32).collect();
            padded.resize(dim, 0.0);
            padded
        };

        // 5-exposure DQN action (Short100/Short50/Flat/Long50/Long100)
        // OrderRouter wraps exposure into FactoredAction with (Market, Normal) defaults.
        // Returns (action, confidence) where confidence is softmax probability of
        // the selected action over Q-values, clamped to [0.5, 0.95].
        let (action, confidence) = dqn
            .select_action_inference(&state_vec)
            .map_err(|e| ml::MLError::InferenceError(format!("DQN prediction failed: {}", e)))?;

        // Map exposure (-1.0..+1.0) to prediction value (0.0..1.0)
        let prediction_value = ((action.target_exposure() + 1.0) / 2.0).clamp(0.0, 1.0);

        if !prediction_value.is_finite() {
            return Err(ml::MLError::InferenceError(format!(
                "DQN model '{}' produced non-finite prediction", self.model_id
            )));
        }

        Ok(ModelPrediction {
            value: prediction_value,
            confidence: f64::from(confidence),
            metadata: std::collections::HashMap::new(),
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64,
            model_id: self.model_id.clone(),
        })
    }

    fn get_confidence(&self) -> f64 {
        0.85
    }

    fn is_ready(&self) -> bool {
        self.loaded && self.dqn.try_read().is_ok()
    }

    fn get_metadata(&self) -> ModelMetadata {
        ModelMetadata {
            model_type: ModelType::DQN,
            version: "2.0.0".to_string(),
            features_used: self.feature_count,
            memory_usage_mb: 74.0,
            additional_metadata: std::collections::HashMap::new(),
        }
    }
}

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

    /// Construct an unloaded DQN model for testing guards.
    fn unloaded_dqn() -> DQNModel {
        let config = ml::dqn::DQNConfig::default();
        let dqn = ml::dqn::DQN::new(config).expect("DQN creation should succeed");
        DQNModel {
            model_id: "dqn-test".to_string(),
            dqn: Arc::new(RwLock::new(dqn)),
            feature_count: 54,
            loaded: false,
        }
    }

    #[tokio::test]
    async fn predict_refuses_unloaded_model() {
        let model = unloaded_dqn();
        let features = ml::Features {
            values: vec![0.1; 54],
            names: vec!["f".into(); 54],
            timestamp: 0,
            symbol: None,
        };
        let result = model.predict(&features).await;
        assert!(result.is_err(), "predict must fail on unloaded DQN model");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(err_msg.contains("not been loaded"), "Should mention not loaded: {err_msg}");
    }

    #[test]
    fn is_ready_false_when_unloaded() {
        let model = unloaded_dqn();
        assert!(!model.is_ready(), "DQN must report is_ready=false when not loaded");
    }

    #[test]
    fn from_checkpoint_rejects_missing_file() {
        let result = DQNModel::from_checkpoint(
            "dqn-test".to_string(),
            std::path::Path::new("/tmp/nonexistent_dqn.safetensors"),
        );
        assert!(result.is_err());
    }

    #[test]
    fn metadata_reports_dqn_type() {
        let config = ml::dqn::DQNConfig::default();
        let dqn = ml::dqn::DQN::new(config).expect("DQN creation should succeed");
        let model = DQNModel {
            model_id: "dqn-test".to_string(),
            dqn: Arc::new(RwLock::new(dqn)),
            feature_count: 54,
            loaded: true,
        };
        let meta = model.get_metadata();
        assert_eq!(meta.model_type, ModelType::DQN);
        assert_eq!(meta.version, "2.0.0");
        assert_eq!(meta.features_used, 54);
    }

    #[test]
    fn is_ready_true_when_loaded() {
        let config = ml::dqn::DQNConfig::default();
        let dqn = ml::dqn::DQN::new(config).expect("DQN creation should succeed");
        let model = DQNModel {
            model_id: "dqn-test".to_string(),
            dqn: Arc::new(RwLock::new(dqn)),
            feature_count: 54,
            loaded: true,
        };
        assert!(model.is_ready(), "DQN must report is_ready=true when loaded");
    }
}
