//! Ensemble Audit Logger for PostgreSQL
//!
//! This module implements production-ready audit logging for ensemble predictions,
//! storing every prediction with per-model attribution to PostgreSQL.
//!
//! ## Features
//! - Async PostgreSQL inserts with connection pooling
//! - Per-model vote tracking (DQN, PPO, MAMBA-2, TFT)
//! - Execution tracking (order_id, P&L attribution)
//! - A/B test metadata
//! - Feature snapshot for reproducibility
//! - High-performance batch inserts

use ml::ensemble::{EnsembleDecision, TradingAction};
use sqlx::PgPool;
use std::time::Instant;
use tracing::{debug, info};
use uuid::Uuid;

/// Prediction audit record for PostgreSQL
#[derive(Debug, Clone)]
pub struct EnsemblePredictionAudit {
    // Trading context
    pub symbol: String,
    pub account_id: Option<String>,
    pub strategy_id: Option<String>,

    // Ensemble decision
    pub ensemble_action: String, // BUY, SELL, HOLD
    pub ensemble_signal: f64,
    pub ensemble_confidence: f64,
    pub disagreement_rate: f64,

    // Per-model votes
    pub dqn_signal: Option<f64>,
    pub dqn_confidence: Option<f64>,
    pub dqn_weight: Option<f64>,
    pub dqn_vote: Option<String>,

    pub ppo_signal: Option<f64>,
    pub ppo_confidence: Option<f64>,
    pub ppo_weight: Option<f64>,
    pub ppo_vote: Option<String>,

    pub mamba2_signal: Option<f64>,
    pub mamba2_confidence: Option<f64>,
    pub mamba2_weight: Option<f64>,
    pub mamba2_vote: Option<String>,

    pub tft_signal: Option<f64>,
    pub tft_confidence: Option<f64>,
    pub tft_weight: Option<f64>,
    pub tft_vote: Option<String>,

    // Execution tracking
    pub order_id: Option<Uuid>,
    pub executed_price: Option<i64>,
    pub position_size: Option<i64>,

    // A/B testing
    pub ab_test_id: Option<Uuid>,
    pub ab_group: Option<String>,
    pub ab_variant: Option<String>,

    // Feature snapshot (JSON)
    pub feature_snapshot: Option<serde_json::Value>,

    // Model checkpoints
    pub dqn_checkpoint_id: Option<String>,
    pub ppo_checkpoint_id: Option<String>,
    pub mamba2_checkpoint_id: Option<String>,
    pub tft_checkpoint_id: Option<String>,

    // System context
    pub node_id: Option<String>,
    pub inference_latency_us: Option<i32>,
    pub aggregation_latency_us: Option<i32>,

    // Compliance
    pub user_id: Option<String>,
    pub session_id: Option<Uuid>,
    pub request_id: Option<Uuid>,

    // Metadata
    pub metadata: Option<serde_json::Value>,
}

impl EnsemblePredictionAudit {
    /// Create audit record from ensemble decision
    pub fn from_decision(decision: &EnsembleDecision, symbol: String) -> Self {
        let ensemble_action = match decision.action {
            TradingAction::Buy => "BUY",
            TradingAction::Sell => "SELL",
            TradingAction::Hold => "HOLD",
        }
        .to_string();

        // Extract per-model votes
        let (dqn_signal, dqn_confidence, dqn_weight, dqn_vote) = decision
            .model_votes
            .get("DQN")
            .map(|v| {
                (
                    Some(v.signal),
                    Some(v.confidence),
                    Some(v.weight),
                    Some(action_to_string(v.signal)),
                )
            })
            .unwrap_or((None, None, None, None));

        let (ppo_signal, ppo_confidence, ppo_weight, ppo_vote) = decision
            .model_votes
            .get("PPO")
            .map(|v| {
                (
                    Some(v.signal),
                    Some(v.confidence),
                    Some(v.weight),
                    Some(action_to_string(v.signal)),
                )
            })
            .unwrap_or((None, None, None, None));

        let (mamba2_signal, mamba2_confidence, mamba2_weight, mamba2_vote) = decision
            .model_votes
            .get("MAMBA2")
            .map(|v| {
                (
                    Some(v.signal),
                    Some(v.confidence),
                    Some(v.weight),
                    Some(action_to_string(v.signal)),
                )
            })
            .unwrap_or((None, None, None, None));

        let (tft_signal, tft_confidence, tft_weight, tft_vote) = decision
            .model_votes
            .get("TFT")
            .map(|v| {
                (
                    Some(v.signal),
                    Some(v.confidence),
                    Some(v.weight),
                    Some(action_to_string(v.signal)),
                )
            })
            .unwrap_or((None, None, None, None));

        Self {
            symbol,
            account_id: None,
            strategy_id: None,
            ensemble_action,
            ensemble_signal: decision.signal,
            ensemble_confidence: decision.confidence,
            disagreement_rate: decision.disagreement_rate,
            dqn_signal,
            dqn_confidence,
            dqn_weight,
            dqn_vote,
            ppo_signal,
            ppo_confidence,
            ppo_weight,
            ppo_vote,
            mamba2_signal,
            mamba2_confidence,
            mamba2_weight,
            mamba2_vote,
            tft_signal,
            tft_confidence,
            tft_weight,
            tft_vote,
            order_id: None,
            executed_price: None,
            position_size: None,
            ab_test_id: None,
            ab_group: None,
            ab_variant: None,
            feature_snapshot: None,
            dqn_checkpoint_id: None,
            ppo_checkpoint_id: None,
            mamba2_checkpoint_id: None,
            tft_checkpoint_id: None,
            node_id: None,
            inference_latency_us: None,
            aggregation_latency_us: None,
            user_id: None,
            session_id: None,
            request_id: None,
            metadata: None,
        }
    }

    /// Set execution tracking details
    pub fn with_execution(
        mut self,
        order_id: Uuid,
        executed_price: i64,
        position_size: i64,
    ) -> Self {
        self.order_id = Some(order_id);
        self.executed_price = Some(executed_price);
        self.position_size = Some(position_size);
        self
    }

    /// Set A/B test metadata
    pub fn with_ab_test(mut self, test_id: Uuid, group: String, variant: Option<String>) -> Self {
        self.ab_test_id = Some(test_id);
        self.ab_group = Some(group);
        self.ab_variant = variant;
        self
    }

    /// Set feature snapshot
    pub fn with_features(mut self, features: serde_json::Value) -> Self {
        self.feature_snapshot = Some(features);
        self
    }

    /// Set model checkpoints
    pub fn with_checkpoints(
        mut self,
        dqn: Option<String>,
        ppo: Option<String>,
        mamba2: Option<String>,
        tft: Option<String>,
    ) -> Self {
        self.dqn_checkpoint_id = dqn;
        self.ppo_checkpoint_id = ppo;
        self.mamba2_checkpoint_id = mamba2;
        self.tft_checkpoint_id = tft;
        self
    }

    /// Set latency metrics
    pub fn with_latency(mut self, inference_us: i32, aggregation_us: i32) -> Self {
        self.inference_latency_us = Some(inference_us);
        self.aggregation_latency_us = Some(aggregation_us);
        self
    }

    /// Set compliance context
    pub fn with_compliance(mut self, user_id: String, session_id: Uuid, request_id: Uuid) -> Self {
        self.user_id = Some(user_id);
        self.session_id = Some(session_id);
        self.request_id = Some(request_id);
        self
    }
}

/// Convert signal to action string
fn action_to_string(signal: f64) -> String {
    if signal > 0.3 {
        "BUY".to_string()
    } else if signal < -0.3 {
        "SELL".to_string()
    } else {
        "HOLD".to_string()
    }
}

/// Ensemble audit logger with PostgreSQL backend
#[derive(Debug, Clone)]
pub struct EnsembleAuditLogger {
    pool: PgPool,
}

impl EnsembleAuditLogger {
    /// Create new audit logger with PostgreSQL pool
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    /// Log ensemble prediction to PostgreSQL
    ///
    /// This is an async operation that inserts the prediction audit record
    /// into the `ensemble_predictions` table.
    pub async fn log_prediction(
        &self,
        audit: EnsemblePredictionAudit,
    ) -> Result<Uuid, sqlx::Error> {
        let start = Instant::now();

        let id = Uuid::new_v4();

        sqlx::query!(
            r#"
            INSERT INTO ensemble_predictions (
                id, symbol, account_id, strategy_id,
                ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
                dqn_signal, dqn_confidence, dqn_weight, dqn_vote,
                ppo_signal, ppo_confidence, ppo_weight, ppo_vote,
                mamba2_signal, mamba2_confidence, mamba2_weight, mamba2_vote,
                tft_signal, tft_confidence, tft_weight, tft_vote,
                order_id, executed_price, position_size,
                ab_test_id, ab_group, ab_variant,
                feature_snapshot,
                dqn_checkpoint_id, ppo_checkpoint_id, mamba2_checkpoint_id, tft_checkpoint_id,
                node_id, inference_latency_us, aggregation_latency_us,
                user_id, session_id, request_id,
                metadata
            ) VALUES (
                $1, $2, $3, $4,
                $5, $6, $7, $8,
                $9, $10, $11, $12,
                $13, $14, $15, $16,
                $17, $18, $19, $20,
                $21, $22, $23, $24,
                $25, $26, $27,
                $28, $29, $30,
                $31,
                $32, $33, $34, $35,
                $36, $37, $38,
                $39, $40, $41,
                $42
            )
            "#,
            id,
            audit.symbol,
            audit.account_id,
            audit.strategy_id,
            audit.ensemble_action,
            audit.ensemble_signal,
            audit.ensemble_confidence,
            audit.disagreement_rate,
            audit.dqn_signal,
            audit.dqn_confidence,
            audit.dqn_weight,
            audit.dqn_vote,
            audit.ppo_signal,
            audit.ppo_confidence,
            audit.ppo_weight,
            audit.ppo_vote,
            audit.mamba2_signal,
            audit.mamba2_confidence,
            audit.mamba2_weight,
            audit.mamba2_vote,
            audit.tft_signal,
            audit.tft_confidence,
            audit.tft_weight,
            audit.tft_vote,
            audit.order_id,
            audit.executed_price,
            audit.position_size,
            audit.ab_test_id,
            audit.ab_group,
            audit.ab_variant,
            audit.feature_snapshot,
            audit.dqn_checkpoint_id,
            audit.ppo_checkpoint_id,
            audit.mamba2_checkpoint_id,
            audit.tft_checkpoint_id,
            audit.node_id,
            audit.inference_latency_us,
            audit.aggregation_latency_us,
            audit.user_id,
            audit.session_id,
            audit.request_id,
            audit.metadata,
        )
        .execute(&self.pool)
        .await?;

        let elapsed = start.elapsed();
        debug!(
            "Logged ensemble prediction {} in {:.2}ms",
            id,
            elapsed.as_secs_f64() * 1000.0
        );

        Ok(id)
    }

    /// Update P&L for a prediction after trade closes
    pub async fn update_pnl(
        &self,
        prediction_id: Uuid,
        pnl: i64,
        commission: i64,
        slippage_bps: i32,
    ) -> Result<(), sqlx::Error> {
        sqlx::query!(
            r#"
            UPDATE ensemble_predictions
            SET pnl = $2, commission = $3, slippage_bps = $4
            WHERE id = $1
            "#,
            prediction_id,
            pnl,
            commission,
            slippage_bps,
        )
        .execute(&self.pool)
        .await?;

        debug!(
            "Updated P&L for prediction {}: {} cents",
            prediction_id, pnl
        );
        Ok(())
    }

    /// Batch insert predictions (for high-throughput scenarios)
    pub async fn log_predictions_batch(
        &self,
        audits: Vec<EnsemblePredictionAudit>,
    ) -> Result<Vec<Uuid>, sqlx::Error> {
        if audits.is_empty() {
            return Ok(Vec::new());
        }

        let start = Instant::now();
        let mut ids = Vec::with_capacity(audits.len());

        let mut tx = self.pool.begin().await?;

        for audit in audits {
            let id = Uuid::new_v4();
            ids.push(id);

            sqlx::query!(
                r#"
                INSERT INTO ensemble_predictions (
                    id, symbol, account_id, strategy_id,
                    ensemble_action, ensemble_signal, ensemble_confidence, disagreement_rate,
                    dqn_signal, dqn_confidence, dqn_weight, dqn_vote,
                    ppo_signal, ppo_confidence, ppo_weight, ppo_vote,
                    mamba2_signal, mamba2_confidence, mamba2_weight, mamba2_vote,
                    tft_signal, tft_confidence, tft_weight, tft_vote,
                    order_id, executed_price, position_size,
                    ab_test_id, ab_group, ab_variant,
                    feature_snapshot,
                    dqn_checkpoint_id, ppo_checkpoint_id, mamba2_checkpoint_id, tft_checkpoint_id,
                    node_id, inference_latency_us, aggregation_latency_us,
                    user_id, session_id, request_id,
                    metadata
                ) VALUES (
                    $1, $2, $3, $4,
                    $5, $6, $7, $8,
                    $9, $10, $11, $12,
                    $13, $14, $15, $16,
                    $17, $18, $19, $20,
                    $21, $22, $23, $24,
                    $25, $26, $27,
                    $28, $29, $30,
                    $31,
                    $32, $33, $34, $35,
                    $36, $37, $38,
                    $39, $40, $41,
                    $42
                )
                "#,
                id,
                audit.symbol,
                audit.account_id,
                audit.strategy_id,
                audit.ensemble_action,
                audit.ensemble_signal,
                audit.ensemble_confidence,
                audit.disagreement_rate,
                audit.dqn_signal,
                audit.dqn_confidence,
                audit.dqn_weight,
                audit.dqn_vote,
                audit.ppo_signal,
                audit.ppo_confidence,
                audit.ppo_weight,
                audit.ppo_vote,
                audit.mamba2_signal,
                audit.mamba2_confidence,
                audit.mamba2_weight,
                audit.mamba2_vote,
                audit.tft_signal,
                audit.tft_confidence,
                audit.tft_weight,
                audit.tft_vote,
                audit.order_id,
                audit.executed_price,
                audit.position_size,
                audit.ab_test_id,
                audit.ab_group,
                audit.ab_variant,
                audit.feature_snapshot,
                audit.dqn_checkpoint_id,
                audit.ppo_checkpoint_id,
                audit.mamba2_checkpoint_id,
                audit.tft_checkpoint_id,
                audit.node_id,
                audit.inference_latency_us,
                audit.aggregation_latency_us,
                audit.user_id,
                audit.session_id,
                audit.request_id,
                audit.metadata,
            )
            .execute(&mut *tx)
            .await?;
        }

        tx.commit().await?;

        let elapsed = start.elapsed();
        info!(
            "Batch logged {} ensemble predictions in {:.2}ms ({:.2} predictions/sec)",
            ids.len(),
            elapsed.as_secs_f64() * 1000.0,
            ids.len() as f64 / elapsed.as_secs_f64()
        );

        Ok(ids)
    }

    /// Record model performance attribution (rolling window metrics)
    pub async fn record_model_performance(
        &self,
        model_id: &str,
        symbol: &str,
        window_hours: i32,
        total_predictions: i32,
        correct_predictions: i32,
        accuracy: f64,
        total_pnl: i64,
        sharpe_ratio: Option<f64>,
        avg_weight: f64,
        avg_confidence: f64,
    ) -> Result<Uuid, sqlx::Error> {
        let id = Uuid::new_v4();

        sqlx::query!(
            r#"
            INSERT INTO model_performance_attribution (
                id, model_id, symbol, window_hours,
                total_predictions, correct_predictions, accuracy,
                total_pnl, sharpe_ratio,
                avg_weight, avg_confidence
            ) VALUES (
                $1, $2, $3, $4,
                $5, $6, $7,
                $8, $9,
                $10, $11
            )
            "#,
            id,
            model_id,
            symbol,
            window_hours,
            total_predictions,
            correct_predictions,
            accuracy,
            total_pnl,
            sharpe_ratio,
            avg_weight,
            avg_confidence,
        )
        .execute(&self.pool)
        .await?;

        debug!(
            "Recorded performance for {} on {}: accuracy={:.2}%, Sharpe={:?}",
            model_id,
            symbol,
            accuracy * 100.0,
            sharpe_ratio
        );

        Ok(id)
    }

    /// Get top performing models (last 24 hours)
    pub async fn get_top_models_24h(
        &self,
        symbol: Option<&str>,
        limit: i32,
    ) -> Result<Vec<ModelPerformanceSummary>, sqlx::Error> {
        let results = sqlx::query_as!(
            ModelPerformanceSummary,
            r#"
            SELECT
                model_id,
                total_predictions,
                accuracy,
                sharpe_ratio,
                total_pnl,
                avg_weight
            FROM get_top_models_24h($1, $2)
            "#,
            symbol,
            limit,
        )
        .fetch_all(&self.pool)
        .await?;

        Ok(results)
    }

    /// Get high disagreement events (last 24 hours)
    pub async fn get_high_disagreement_events_24h(
        &self,
        symbol: Option<&str>,
        disagreement_threshold: f64,
        limit: i32,
    ) -> Result<Vec<HighDisagreementEvent>, sqlx::Error> {
        let results = sqlx::query_as!(
            HighDisagreementEvent,
            r#"
            SELECT
                prediction_timestamp as "timestamp",
                symbol,
                ensemble_action,
                ensemble_confidence,
                disagreement_rate,
                dqn_vote,
                ppo_vote,
                mamba2_vote,
                tft_vote
            FROM get_high_disagreement_events_24h($1, $2, $3)
            "#,
            symbol,
            disagreement_threshold,
            limit,
        )
        .fetch_all(&self.pool)
        .await?;

        Ok(results)
    }
}

/// Model performance summary
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct ModelPerformanceSummary {
    pub model_id: Option<String>,
    pub total_predictions: Option<i32>, // Function returns INTEGER (i32)
    pub accuracy: Option<f64>,          // Function returns FLOAT (f64)
    pub sharpe_ratio: Option<f64>,      // Function returns FLOAT (f64)
    pub total_pnl: Option<i64>,         // Function returns BIGINT (i64)
    pub avg_weight: Option<f64>,        // Function returns FLOAT (f64)
}

/// High disagreement event
#[derive(Debug, Clone, sqlx::FromRow)]
pub struct HighDisagreementEvent {
    pub timestamp: Option<chrono::DateTime<chrono::Utc>>,
    pub symbol: Option<String>,
    pub ensemble_action: Option<String>,
    pub ensemble_confidence: Option<f64>,
    pub disagreement_rate: Option<f64>,
    pub dqn_vote: Option<String>,
    pub ppo_vote: Option<String>,
    pub mamba2_vote: Option<String>,
    pub tft_vote: Option<String>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use ml::ensemble::{ModelVote, TradingAction};
    use std::collections::HashMap;

    #[test]
    fn test_audit_from_decision() {
        let mut model_votes = HashMap::new();
        model_votes.insert(
            "DQN".to_string(),
            ModelVote::new("DQN".to_string(), 0.8, 0.9, 0.33),
        );
        model_votes.insert(
            "PPO".to_string(),
            ModelVote::new("PPO".to_string(), 0.7, 0.85, 0.33),
        );
        model_votes.insert(
            "TFT".to_string(),
            ModelVote::new("TFT".to_string(), 0.6, 0.8, 0.34),
        );

        let decision = EnsembleDecision::new(TradingAction::Buy, 0.85, 0.7, 0.15, model_votes);

        let audit = EnsemblePredictionAudit::from_decision(&decision, "ES.FUT".to_string());

        assert_eq!(audit.symbol, "ES.FUT");
        assert_eq!(audit.ensemble_action, "BUY");
        assert_eq!(audit.ensemble_signal, 0.7);
        assert_eq!(audit.ensemble_confidence, 0.85);
        assert_eq!(audit.disagreement_rate, 0.15);

        // Check DQN vote
        assert_eq!(audit.dqn_signal, Some(0.8));
        assert_eq!(audit.dqn_confidence, Some(0.9));
        assert_eq!(audit.dqn_weight, Some(0.33));
        assert_eq!(audit.dqn_vote, Some("BUY".to_string()));
    }

    #[test]
    fn test_audit_builder_pattern() {
        let decision = EnsembleDecision::new(TradingAction::Hold, 0.6, 0.1, 0.4, HashMap::new());

        let audit = EnsemblePredictionAudit::from_decision(&decision, "NQ.FUT".to_string())
            .with_execution(Uuid::new_v4(), 15000, 10)
            .with_ab_test(
                Uuid::new_v4(),
                "treatment".to_string(),
                Some("ensemble_v2".to_string()),
            )
            .with_latency(42, 8);

        assert_eq!(audit.symbol, "NQ.FUT");
        assert!(audit.order_id.is_some());
        assert!(audit.ab_test_id.is_some());
        assert_eq!(audit.ab_group, Some("treatment".to_string()));
        assert_eq!(audit.inference_latency_us, Some(42));
        assert_eq!(audit.aggregation_latency_us, Some(8));
    }
}
