//! Public data types: drift scores, reports, and baseline snapshots.

use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

use crate::error::{RagDriftError, Result};

/// One of the five drift dimensions ragdrift monitors.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum DriftDimension {
    /// Tabular feature drift (KS / PSI per feature).
    Data,
    /// Embedding drift (MMD + sliced Wasserstein).
    Embedding,
    /// Response distribution drift (length, entropy, optional semantic).
    Response,
    /// Confidence score drift (KS + ECE delta).
    Confidence,
    /// Query-pattern drift (cluster assignment KL).
    Query,
}

impl DriftDimension {
    /// Stable lowercase string identifier used in metric names.
    pub fn as_str(&self) -> &'static str {
        match self {
            Self::Data => "data",
            Self::Embedding => "embedding",
            Self::Response => "response",
            Self::Confidence => "confidence",
            Self::Query => "query",
        }
    }
}

/// A scalar drift score for one dimension.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DriftScore {
    /// Which of the five dimensions this score is for.
    pub dimension: DriftDimension,
    /// Numeric drift magnitude. Larger means more drift; semantics vary by method.
    pub score: f64,
    /// Threshold above which the score is considered drifted.
    pub threshold: f64,
    /// True iff `score > threshold`.
    pub exceeded: bool,
    /// Free-form algorithm identifier (e.g. `"ks"`, `"mmd+sw"`).
    pub method: String,
}

impl DriftScore {
    /// Construct a [`DriftScore`] and set `exceeded` from `score > threshold`.
    pub fn new(
        dimension: DriftDimension,
        score: f64,
        threshold: f64,
        method: impl Into<String>,
    ) -> Self {
        Self {
            dimension,
            score,
            threshold,
            exceeded: score > threshold,
            method: method.into(),
        }
    }
}

/// A combined drift report across all configured dimensions.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DriftReport {
    /// One entry per detector that ran successfully.
    pub scores: Vec<DriftScore>,
    /// Unix timestamp (seconds) when the report was generated.
    pub timestamp: i64,
    /// Number of baseline samples that fed the report.
    pub sample_size_baseline: usize,
    /// Number of current samples that fed the report.
    pub sample_size_current: usize,
}

impl DriftReport {
    /// True iff any score exceeds its threshold.
    pub fn any_exceeded(&self) -> bool {
        self.scores.iter().any(|s| s.exceeded)
    }

    /// Return the score for a specific dimension if present.
    pub fn score_for(&self, dim: DriftDimension) -> Option<&DriftScore> {
        self.scores.iter().find(|s| s.dimension == dim)
    }

    /// How many of the configured detectors flagged drift.
    pub fn exceeded_count(&self) -> usize {
        self.scores.iter().filter(|s| s.exceeded).count()
    }

    /// Highest-magnitude score across all dimensions, if any. Useful for
    /// alerting policies that page on a single combined number.
    pub fn worst_score(&self) -> Option<&DriftScore> {
        self.scores.iter().max_by(|a, b| {
            a.score
                .partial_cmp(&b.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        })
    }

    /// Multi-line human-readable rollup. Stable, machine-grep-able format
    /// suitable for cron-job emails, Slack messages, log lines, etc.
    /// One line per detector plus a header. Format example:
    ///
    /// ```text
    /// ragdrift report (n_baseline=1024, n_current=1024, exceeded=2/3):
    ///   [DRIFT] embedding   score=  1.0324 threshold=0.100 method=mmd+sw
    ///   [ ok  ] data        score=  0.0124 threshold=0.250 method=ks+psi
    ///   [DRIFT] confidence  score=  1.0000 threshold=0.150 method=ks
    /// ```
    pub fn summary(&self) -> String {
        let total = self.scores.len();
        let bad = self.exceeded_count();
        let mut out = format!(
            "ragdrift report (n_baseline={}, n_current={}, exceeded={}/{}):",
            self.sample_size_baseline, self.sample_size_current, bad, total
        );
        for s in &self.scores {
            let flag = if s.exceeded { "DRIFT" } else { " ok  " };
            out.push_str(&format!(
                "\n  [{}] {:11} score={:8.4} threshold={:5.3} method={}",
                flag,
                s.dimension.as_str(),
                s.score,
                s.threshold,
                s.method
            ));
        }
        out
    }
}

/// Persisted baseline against which future samples are compared.
///
/// Only the fields a caller has are populated. Save with
/// [`BaselineSnapshot::save`] and reload with [`BaselineSnapshot::load`].
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct BaselineSnapshot {
    /// Reference embedding matrix (n_samples x dim).
    pub embeddings: Option<Array2<f32>>,
    /// Reference tabular features (n_samples x n_features).
    pub features: Option<Array2<f64>>,
    /// Reference response lengths in tokens.
    pub response_lengths: Option<Array1<f64>>,
    /// Reference confidence scores in `[0, 1]`.
    pub confidence_scores: Option<Array1<f64>>,
    /// Reference query embeddings (n_samples x dim) for cluster fitting.
    pub query_embeddings: Option<Array2<f32>>,
    /// Free-form metadata (model name, index name, capture window).
    pub metadata: std::collections::BTreeMap<String, String>,
}

impl BaselineSnapshot {
    /// Save the snapshot to a JSON file.
    pub fn save(&self, path: &std::path::Path) -> Result<()> {
        let file = std::fs::File::create(path)?;
        serde_json::to_writer(std::io::BufWriter::new(file), self)?;
        Ok(())
    }

    /// Load a snapshot from a JSON file.
    pub fn load(path: &std::path::Path) -> Result<Self> {
        let file = std::fs::File::open(path)?;
        let snap = serde_json::from_reader(std::io::BufReader::new(file))?;
        Ok(snap)
    }
}

/// Validate that two 2D arrays share the same column count.
pub(crate) fn check_same_cols<T>(
    baseline: &ndarray::ArrayView2<'_, T>,
    current: &ndarray::ArrayView2<'_, T>,
) -> Result<()> {
    if baseline.ncols() != current.ncols() {
        return Err(RagDriftError::DimensionMismatch {
            baseline: baseline.shape().to_vec(),
            current: current.shape().to_vec(),
        });
    }
    Ok(())
}

/// Validate that an array has at least `needed` samples along the leading axis.
pub(crate) fn check_min_samples(got: usize, needed: usize) -> Result<()> {
    if got < needed {
        return Err(RagDriftError::InsufficientSamples { needed, got });
    }
    Ok(())
}

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

    #[test]
    fn drift_score_exceeded_flag_set() {
        let s = DriftScore::new(DriftDimension::Data, 0.3, 0.25, "ks");
        assert!(s.exceeded);
        let s = DriftScore::new(DriftDimension::Data, 0.1, 0.25, "ks");
        assert!(!s.exceeded);
    }

    #[test]
    fn drift_dimension_as_str_is_snake_case() {
        assert_eq!(DriftDimension::Embedding.as_str(), "embedding");
        assert_eq!(DriftDimension::Query.as_str(), "query");
    }

    #[test]
    fn report_any_exceeded() {
        let r = DriftReport {
            scores: vec![
                DriftScore::new(DriftDimension::Data, 0.0, 0.25, "ks"),
                DriftScore::new(DriftDimension::Embedding, 0.5, 0.25, "mmd+sw"),
            ],
            timestamp: 0,
            sample_size_baseline: 10,
            sample_size_current: 10,
        };
        assert!(r.any_exceeded());
        assert_eq!(r.score_for(DriftDimension::Embedding).unwrap().score, 0.5);
    }

    fn sample_report() -> DriftReport {
        DriftReport {
            scores: vec![
                DriftScore::new(DriftDimension::Data, 0.05, 0.25, "ks"),
                DriftScore::new(DriftDimension::Embedding, 1.2, 0.1, "mmd+sw"),
                DriftScore::new(DriftDimension::Confidence, 0.8, 0.15, "ks"),
            ],
            timestamp: 0,
            sample_size_baseline: 1024,
            sample_size_current: 1024,
        }
    }

    #[test]
    fn report_exceeded_count() {
        let r = sample_report();
        assert_eq!(r.exceeded_count(), 2);
    }

    #[test]
    fn report_worst_score_picks_largest() {
        let r = sample_report();
        let worst = r.worst_score().unwrap();
        assert_eq!(worst.dimension, DriftDimension::Embedding);
        assert!((worst.score - 1.2).abs() < 1e-9);
    }

    #[test]
    fn report_worst_score_none_when_empty() {
        let r = DriftReport {
            scores: vec![],
            timestamp: 0,
            sample_size_baseline: 0,
            sample_size_current: 0,
        };
        assert!(r.worst_score().is_none());
    }

    #[test]
    fn report_summary_format_is_grep_able() {
        let r = sample_report();
        let s = r.summary();
        // Header line.
        assert!(s.starts_with("ragdrift report ("));
        assert!(s.contains("exceeded=2/3"));
        assert!(s.contains("n_baseline=1024"));
        // Per-detector lines.
        assert!(s.contains("[DRIFT] embedding"));
        assert!(s.contains("[DRIFT] confidence"));
        assert!(s.contains("[ ok  ] data"));
        // 1 header + 3 detector lines = 4 lines total.
        assert_eq!(s.lines().count(), 4);
    }

    #[test]
    fn snapshot_round_trip() {
        let dir = tempdir();
        let path = dir.join("snap.json");
        let snap = BaselineSnapshot {
            response_lengths: Some(Array1::from(vec![1.0, 2.0, 3.0])),
            ..Default::default()
        };
        snap.save(&path).unwrap();
        let loaded = BaselineSnapshot::load(&path).unwrap();
        assert_eq!(loaded.response_lengths.unwrap().len(), 3);
    }

    fn tempdir() -> std::path::PathBuf {
        let dir = std::env::temp_dir().join(format!(
            "ragdrift-test-{}-{}",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }
}
