//! Production PPO model — loads from actor/critic safetensors checkpoints.
//!
//! Uses 45-action FactoredAction space (5 exposure × 3 order × 3 urgency).

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

pub(crate) struct PPOModel {
    model_id: String,
    agent: Arc<RwLock<ml::ppo::PPO>>,
    feature_count: usize,
    loaded: bool,
}

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

impl PPOModel {
    /// Create new PPO model from checkpoint (actor + critic)
    ///
    /// Reads architecture config from the checkpoint's `_metadata.json` file
    /// (saved by `PPO::save_checkpoint`). Falls back to hardcoded defaults
    /// for legacy checkpoints without metadata.
    pub fn from_checkpoint(
        model_id: String,
        actor_path: &std::path::Path,
        critic_path: &std::path::Path,
    ) -> ml::MLResult<Self> {
        use ml::ppo::PPO;

        // Try to read architecture config from checkpoint metadata JSON
        let config = Self::config_from_metadata(actor_path).unwrap_or_else(|e| {
            tracing::info!(
                "PPO '{}': no metadata config ({}) — using hardcoded defaults",
                model_id, e
            );
            Self::default_config()
        });

        let state_dim = config.state_dim;
        info!(
            "PPO '{}' config: state_dim={}, num_actions={}, \
             policy_hidden_dims={:?}, value_hidden_dims={:?}",
            model_id, config.state_dim, config.num_actions,
            config.policy_hidden_dims, config.value_hidden_dims
        );

        // PPO::load_checkpoint reads config from <path>.json and weights from <path>.safetensors
        let actor_pathbuf = actor_path.to_path_buf();
        let agent = PPO::load_checkpoint(&actor_pathbuf)
            .map_err(|e| {
                ml::MLError::ModelError(format!("Failed to load PPO checkpoint: {}", e))
            })?;
        let _ = (critic_path, config); // Critic path loaded by load_checkpoint internally

        info!(
            "Loaded PPO model {} from actor={}, critic={}",
            model_id,
            actor_path.display(),
            critic_path.display()
        );

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

    /// Read PPOConfig from the metadata JSON saved alongside the actor checkpoint.
    fn config_from_metadata(actor_path: &std::path::Path) -> Result<ml::ppo::PPOConfig, String> {
        let metadata_path = if let Some(parent) = actor_path.parent() {
            if let Some(stem) = actor_path.file_stem() {
                parent.join(format!("{}_metadata.json", stem.to_string_lossy()))
            } else {
                parent.join("checkpoint_metadata.json")
            }
        } else {
            return Err("no parent directory".to_string());
        };

        let raw = std::fs::read_to_string(&metadata_path)
            .map_err(|e| format!("read {}: {}", metadata_path.display(), e))?;
        let json: serde_json::Value = serde_json::from_str(&raw)
            .map_err(|e| format!("parse {}: {}", metadata_path.display(), e))?;
        let cfg = json.get("config").ok_or("missing 'config' key")?;

        let state_dim = cfg.get("state_dim").and_then(|v| v.as_u64()).unwrap_or(54) as usize;
        let num_actions = cfg.get("num_actions").and_then(|v| v.as_u64()).unwrap_or(45) as usize;

        let parse_dims = |key: &str, default: &[usize]| -> Vec<usize> {
            cfg.get(key)
                .and_then(|v| v.as_array())
                .map(|arr| arr.iter().filter_map(|x| x.as_u64().map(|n| n as usize)).collect())
                .unwrap_or_else(|| default.to_vec())
        };

        let policy_hidden_dims = parse_dims("policy_hidden_dims", &[128, 64]);
        let value_hidden_dims = parse_dims("value_hidden_dims", &[512, 384, 256, 128, 64]);

        let mut config = Self::default_config();
        config.state_dim = state_dim;
        config.num_actions = num_actions;
        config.policy_hidden_dims = policy_hidden_dims;
        config.value_hidden_dims = value_hidden_dims;

        // Override training hyperparams from metadata if present
        if let Some(v) = cfg.get("clip_epsilon").and_then(|v| v.as_f64()) {
            config.clip_epsilon = v as f32;
        }
        if let Some(v) = cfg.get("entropy_coeff").and_then(|v| v.as_f64()) {
            config.entropy_coeff = v as f32;
        }
        if let Some(v) = cfg.get("max_grad_norm").and_then(|v| v.as_f64()) {
            config.max_grad_norm = v as f32;
        }

        Ok(config)
    }

    /// Default PPOConfig for legacy checkpoints without metadata.
    fn default_config() -> ml::ppo::PPOConfig {
        ml::ppo::PPOConfig {
            state_dim: 54,
            num_actions: 45,
            policy_hidden_dims: vec![128, 64],
            value_hidden_dims: vec![512, 384, 256, 128, 64],
            policy_learning_rate: 0.0003,
            value_learning_rate: 0.001,
            clip_epsilon: 0.2,
            value_loss_coeff: 0.5,
            entropy_coeff: 0.01,
            gae_config: ml::ppo::gae::GAEConfig {
                gamma: 0.99,
                lambda: 0.95,
                normalize_advantages: true,
            },
            batch_size: 64,
            mini_batch_size: 32,
            num_epochs: 10,
            max_grad_norm: 0.5,
            early_stopping_enabled: false,
            early_stopping_patience: 5,
            early_stopping_min_delta: 0.001,
            early_stopping_min_epochs: 10,
            max_position_absolute: 2.0,
            transaction_cost_bps: 0.10,
            cash_reserve_pct: 20.0,
            circuit_breaker_threshold: 5,
            use_lstm: false,
            lstm_hidden_dim: 128,
            lstm_num_layers: 1,
            lstm_sequence_length: 32,
            accumulation_steps: 1,
            clip_epsilon_high: None,
            use_symlog: true,
            use_adaptive_entropy: true,
            use_percentile_scaling: true,
        }
    }
}

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

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

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

        let agent = self.agent.read().await;

        let dim = self.feature_count;
        let state_vec: Vec<f32> = if features.values.len() >= dim {
            features.values[0..dim].iter().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
        };

        // 45-action FactoredAction (5 exposure × 3 order × 3 urgency)
        // act_with_log_prob returns (action, log_prob, value) — act() drops log_prob
        let (action, log_prob, _value) = agent
            .act_with_log_prob(&state_vec)
            .map_err(|e| ml::MLError::InferenceError(format!("PPO 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!(
                "PPO model '{}' produced non-finite prediction", self.model_id
            )));
        }

        // Confidence from log probability: exp(log_prob) ∈ (0, 1]
        let confidence = (log_prob.exp().clamp(0.5, 0.95)) as f64;

        Ok(ModelPrediction {
            value: prediction_value,
            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.agent.try_read().is_ok()
    }

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

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

    /// Construct an unloaded PPO model for testing guards.
    fn unloaded_ppo() -> PPOModel {
        use ml::ppo::PPO;

        let config = PPOModel::default_config();
        let agent = PPO::new(config)
            .expect("PPO creation should succeed");
        PPOModel {
            model_id: "ppo-test".to_string(),
            agent: Arc::new(RwLock::new(agent)),
            feature_count: 54,
            loaded: false,
        }
    }

    #[tokio::test]
    async fn predict_refuses_unloaded_model() {
        let model = unloaded_ppo();
        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 PPO 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_ppo();
        assert!(!model.is_ready(), "PPO must report is_ready=false when not loaded");
    }

    #[test]
    fn metadata_reports_ppo_type() {
        let model = unloaded_ppo();
        let meta = model.get_metadata();
        assert_eq!(meta.model_type, ModelType::PPO);
        assert_eq!(meta.version, "1.0.0");
        assert_eq!(meta.features_used, 54);
    }

    #[tokio::test]
    async fn predict_output_clamped_when_loaded() {
        use ml::ppo::PPO;

        let config = PPOModel::default_config();
        let agent = PPO::new(config)
            .expect("PPO creation should succeed");
        let model = PPOModel {
            model_id: "ppo-test".to_string(),
            agent: Arc::new(RwLock::new(agent)),
            feature_count: 54,
            loaded: true,
        };
        let features = ml::Features {
            values: vec![0.5; 54],
            names: vec!["f".into(); 54],
            timestamp: 0,
            symbol: None,
        };
        let result = model.predict(&features).await;
        assert!(result.is_ok(), "predict should succeed on loaded PPO: {:?}", result.err());
        let pred = result.unwrap();
        assert!((0.0..=1.0).contains(&pred.value), "prediction must be in [0,1], got {}", pred.value);
    }
}
