//! Model Promotion Manager
//!
//! Tracks models awaiting operator approval for promotion to active status.
//! Compares new model metrics against currently active model to determine
//! if promotion should be offered.

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

use chrono::{DateTime, Utc};
use tokio::sync::RwLock;
use tracing::info;

/// A model that has completed training and is pending operator review.
#[derive(Debug, Clone)]
pub struct PendingModel {
    pub model_id: String,
    pub model_type: String,
    pub symbol: String,
    pub s3_path: String,
    pub new_metrics: HashMap<String, f64>,
    pub current_metrics: HashMap<String, f64>,
    pub trained_at: DateTime<Utc>,
    pub job_id: String,
}

/// The currently active model for a given (model_type, symbol) pair.
#[derive(Debug, Clone)]
pub struct ActiveModel {
    pub model_id: String,
    pub s3_path: String,
    pub metrics: HashMap<String, f64>,
    pub promoted_at: DateTime<Utc>,
}

/// Promotion status returned after registering a completion.
#[derive(Debug, Clone, PartialEq)]
pub enum PromotionStatus {
    /// New model is better -- awaiting approval.
    PendingPromotion,
    /// New model is not better than active.
    NoImprovement,
    /// No active model exists -- registered (auto-promote candidate).
    Registered,
    /// Error in comparison.
    Error(String),
}

/// Manages pending promotions and active model tracking.
pub struct PromotionManager {
    pending: Arc<RwLock<HashMap<String, PendingModel>>>,
    /// Maps (model_type, symbol) to the currently active model.
    active_models: Arc<RwLock<HashMap<(String, String), ActiveModel>>>,
}

impl Default for PromotionManager {
    fn default() -> Self {
        Self::new()
    }
}

impl PromotionManager {
    /// Create a new promotion manager with empty state.
    pub fn new() -> Self {
        Self {
            pending: Arc::new(RwLock::new(HashMap::new())),
            active_models: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a completed training job and determine promotion status.
    ///
    /// If no active model exists for this (model_type, symbol) pair, returns
    /// `Registered`. If the new model is better, stores it as pending and
    /// returns `PendingPromotion`. Otherwise returns `NoImprovement`.
    pub async fn register_completion(
        &self,
        job_id: &str,
        model_type: &str,
        symbol: &str,
        s3_path: &str,
        metrics: HashMap<String, f64>,
    ) -> PromotionStatus {
        let key = (model_type.to_string(), symbol.to_string());
        let active = self.active_models.read().await;

        let status = match active.get(&key) {
            None => {
                // No active model for this type+symbol -- first model
                drop(active);
                PromotionStatus::Registered
            }
            Some(active_model) => {
                let current_metrics = active_model.metrics.clone();
                drop(active);

                if Self::is_better(&metrics, &current_metrics) {
                    let model_id = uuid::Uuid::new_v4().to_string();
                    let pending_model = PendingModel {
                        model_id: model_id.clone(),
                        model_type: model_type.to_string(),
                        symbol: symbol.to_string(),
                        s3_path: s3_path.to_string(),
                        new_metrics: metrics,
                        current_metrics,
                        trained_at: Utc::now(),
                        job_id: job_id.to_string(),
                    };

                    let mut pending = self.pending.write().await;
                    pending.insert(model_id.clone(), pending_model);
                    info!(
                        model_id = %model_id,
                        model_type = %model_type,
                        symbol = %symbol,
                        "new model is better than active -- pending promotion"
                    );
                    PromotionStatus::PendingPromotion
                } else {
                    PromotionStatus::NoImprovement
                }
            }
        };

        status
    }

    /// Return all models currently pending promotion.
    pub async fn list_pending(&self) -> Vec<PendingModel> {
        let pending = self.pending.read().await;
        pending.values().cloned().collect()
    }

    /// Approve a pending model, promoting it to active status.
    ///
    /// Replaces the current active model for the same (model_type, symbol)
    /// and removes the entry from the pending map.
    pub async fn approve(&self, model_id: &str) -> anyhow::Result<()> {
        let mut pending = self.pending.write().await;

        let model = pending
            .remove(model_id)
            .ok_or_else(|| anyhow::anyhow!("pending model not found: {}", model_id))?;

        let key = (model.model_type.clone(), model.symbol.clone());
        let active = ActiveModel {
            model_id: model.model_id.clone(),
            s3_path: model.s3_path.clone(),
            metrics: model.new_metrics.clone(),
            promoted_at: Utc::now(),
        };

        let mut active_models = self.active_models.write().await;
        active_models.insert(key, active);

        info!(
            model_id = %model.model_id,
            model_type = %model.model_type,
            symbol = %model.symbol,
            s3_path = %model.s3_path,
            "model promoted to active"
        );

        Ok(())
    }

    /// Reject a pending model and remove it from the pending map.
    pub async fn reject(&self, model_id: &str, reason: &str) -> anyhow::Result<()> {
        let mut pending = self.pending.write().await;

        let model = pending
            .remove(model_id)
            .ok_or_else(|| anyhow::anyhow!("pending model not found: {}", model_id))?;

        info!(
            model_id = %model.model_id,
            model_type = %model.model_type,
            symbol = %model.symbol,
            reason = %reason,
            "model promotion rejected"
        );

        Ok(())
    }

    /// Expose the active_models map for test assertions (test-only).
    #[cfg(test)]
    pub fn active_models_for_test(&self) -> Arc<RwLock<HashMap<(String, String), ActiveModel>>> {
        Arc::clone(&self.active_models)
    }

    /// Compare new metrics against current active metrics.
    ///
    /// Returns true if the new model is better:
    /// - `best_val_loss`: lower is better
    /// - `sharpe_ratio`: higher is better
    ///
    /// If either metric is strictly better AND the other is not worse, returns true.
    /// If neither metric exists in both maps, returns false.
    fn is_better(new_metrics: &HashMap<String, f64>, current: &HashMap<String, f64>) -> bool {
        let new_loss = new_metrics.get("best_val_loss");
        let cur_loss = current.get("best_val_loss");

        let new_sharpe = new_metrics.get("sharpe_ratio");
        let cur_sharpe = current.get("sharpe_ratio");

        // Check if we have any comparable metrics
        let has_loss = new_loss.is_some() && cur_loss.is_some();
        let has_sharpe = new_sharpe.is_some() && cur_sharpe.is_some();

        if !has_loss && !has_sharpe {
            return false;
        }

        // For val_loss: lower is better; for sharpe: higher is better
        let loss_better = match (new_loss, cur_loss) {
            (Some(n), Some(c)) => *n < *c,
            _ => false,
        };
        let loss_not_worse = match (new_loss, cur_loss) {
            (Some(n), Some(c)) => *n <= *c,
            _ => true, // If missing, not considered worse
        };

        let sharpe_better = match (new_sharpe, cur_sharpe) {
            (Some(n), Some(c)) => *n > *c,
            _ => false,
        };
        let sharpe_not_worse = match (new_sharpe, cur_sharpe) {
            (Some(n), Some(c)) => *n >= *c,
            _ => true, // If missing, not considered worse
        };

        // Either metric is strictly better AND the other is not worse
        (loss_better && sharpe_not_worse) || (sharpe_better && loss_not_worse)
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_register_first_model_returns_registered() {
        let mgr = PromotionManager::new();
        let metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.05),
            ("sharpe_ratio".to_string(), 1.5),
        ]);

        let status = mgr
            .register_completion("job-1", "DQN", "ES.FUT", "s3://bucket/model.bin", metrics)
            .await;

        assert_eq!(status, PromotionStatus::Registered);
    }

    #[tokio::test]
    async fn test_register_better_model_returns_pending_promotion() {
        let mgr = PromotionManager::new();

        // Seed an active model
        {
            let mut active = mgr.active_models.write().await;
            active.insert(
                ("DQN".to_string(), "ES.FUT".to_string()),
                ActiveModel {
                    model_id: "old-model".to_string(),
                    s3_path: "s3://bucket/old.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.10),
                        ("sharpe_ratio".to_string(), 1.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a better model (lower val_loss, higher sharpe)
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.05),
            ("sharpe_ratio".to_string(), 2.0),
        ]);

        let status = mgr
            .register_completion("job-2", "DQN", "ES.FUT", "s3://bucket/new.bin", new_metrics)
            .await;

        assert_eq!(status, PromotionStatus::PendingPromotion);

        // Verify it's in the pending list
        let pending = mgr.list_pending().await;
        assert_eq!(pending.len(), 1);
        assert_eq!(pending[0].job_id, "job-2");
    }

    #[tokio::test]
    async fn test_register_worse_model_returns_no_improvement() {
        let mgr = PromotionManager::new();

        // Seed an active model with good metrics
        {
            let mut active = mgr.active_models.write().await;
            active.insert(
                ("DQN".to_string(), "ES.FUT".to_string()),
                ActiveModel {
                    model_id: "good-model".to_string(),
                    s3_path: "s3://bucket/good.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.02),
                        ("sharpe_ratio".to_string(), 3.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a worse model (higher val_loss, lower sharpe)
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);

        let status = mgr
            .register_completion("job-3", "DQN", "ES.FUT", "s3://bucket/bad.bin", new_metrics)
            .await;

        assert_eq!(status, PromotionStatus::NoImprovement);

        // Verify nothing pending
        let pending = mgr.list_pending().await;
        assert!(pending.is_empty());
    }

    #[tokio::test]
    async fn test_approve_moves_to_active() {
        let mgr = PromotionManager::new();

        // Seed an active model
        {
            let mut active = mgr.active_models.write().await;
            active.insert(
                ("DQN".to_string(), "ES.FUT".to_string()),
                ActiveModel {
                    model_id: "old-model".to_string(),
                    s3_path: "s3://bucket/old.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.10),
                        ("sharpe_ratio".to_string(), 1.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a better model
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.03),
            ("sharpe_ratio".to_string(), 2.5),
        ]);
        let status = mgr
            .register_completion("job-4", "DQN", "ES.FUT", "s3://bucket/better.bin", new_metrics)
            .await;
        assert_eq!(status, PromotionStatus::PendingPromotion);

        // Get the model_id from pending
        let pending = mgr.list_pending().await;
        assert_eq!(pending.len(), 1);
        let model_id = pending[0].model_id.clone();

        // Approve it
        mgr.approve(&model_id).await.unwrap();

        // Verify pending is empty
        let pending = mgr.list_pending().await;
        assert!(pending.is_empty());

        // Verify active was updated
        let active = mgr.active_models.read().await;
        let active_model = active
            .get(&("DQN".to_string(), "ES.FUT".to_string()))
            .unwrap();
        assert_eq!(active_model.model_id, model_id);
        assert_eq!(active_model.s3_path, "s3://bucket/better.bin");
    }

    #[tokio::test]
    async fn test_reject_removes_from_pending() {
        let mgr = PromotionManager::new();

        // Seed an active model
        {
            let mut active = mgr.active_models.write().await;
            active.insert(
                ("DQN".to_string(), "ES.FUT".to_string()),
                ActiveModel {
                    model_id: "old-model".to_string(),
                    s3_path: "s3://bucket/old.bin".to_string(),
                    metrics: HashMap::from([
                        ("best_val_loss".to_string(), 0.10),
                        ("sharpe_ratio".to_string(), 1.0),
                    ]),
                    promoted_at: Utc::now(),
                },
            );
        }

        // Register a better model
        let new_metrics = HashMap::from([
            ("best_val_loss".to_string(), 0.04),
            ("sharpe_ratio".to_string(), 2.0),
        ]);
        let status = mgr
            .register_completion("job-5", "DQN", "ES.FUT", "s3://bucket/candidate.bin", new_metrics)
            .await;
        assert_eq!(status, PromotionStatus::PendingPromotion);

        // Get the model_id from pending
        let pending = mgr.list_pending().await;
        assert_eq!(pending.len(), 1);
        let model_id = pending[0].model_id.clone();

        // Reject it
        mgr.reject(&model_id, "not ready for production")
            .await
            .unwrap();

        // Verify pending is empty
        let pending = mgr.list_pending().await;
        assert!(pending.is_empty());

        // Verify active was NOT changed
        let active = mgr.active_models.read().await;
        let active_model = active
            .get(&("DQN".to_string(), "ES.FUT".to_string()))
            .unwrap();
        assert_eq!(active_model.model_id, "old-model");
    }

    #[test]
    fn test_is_better_both_metrics_improved() {
        let new = HashMap::from([
            ("best_val_loss".to_string(), 0.03),
            ("sharpe_ratio".to_string(), 2.0),
        ]);
        let current = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        assert!(PromotionManager::is_better(&new, &current));
    }

    #[test]
    fn test_is_better_only_loss_improved() {
        let new = HashMap::from([
            ("best_val_loss".to_string(), 0.03),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        let current = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        assert!(PromotionManager::is_better(&new, &current));
    }

    #[test]
    fn test_is_better_only_sharpe_improved() {
        let new = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 2.0),
        ]);
        let current = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        assert!(PromotionManager::is_better(&new, &current));
    }

    #[test]
    fn test_is_better_no_metrics() {
        let new: HashMap<String, f64> = HashMap::new();
        let current: HashMap<String, f64> = HashMap::new();
        assert!(!PromotionManager::is_better(&new, &current));
    }

    #[test]
    fn test_is_better_worse_on_both() {
        let new = HashMap::from([
            ("best_val_loss".to_string(), 0.20),
            ("sharpe_ratio".to_string(), 0.5),
        ]);
        let current = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        assert!(!PromotionManager::is_better(&new, &current));
    }

    #[test]
    fn test_is_better_tradeoff_loss_better_sharpe_worse() {
        // Loss is better but sharpe is worse -- should NOT promote
        let new = HashMap::from([
            ("best_val_loss".to_string(), 0.03),
            ("sharpe_ratio".to_string(), 0.5),
        ]);
        let current = HashMap::from([
            ("best_val_loss".to_string(), 0.10),
            ("sharpe_ratio".to_string(), 1.0),
        ]);
        assert!(!PromotionManager::is_better(&new, &current));
    }
}
