//! Plan 5 Task 2 — A.4 Regression Detection Hard-Stop.
//!
//! Watches every HEALTH_DIAG metric per epoch against the bands defined in
//! `config/metric-bands.toml`. After `2 × consecutive_epochs_for_warn`
//! consecutive epochs of any single metric being out-of-error-band, training
//! self-terminates with `CommonError::RegressionDetected`.
//!
//! Behaviour summary:
//!
//! * `update_and_check(metric, value)` is called once per (epoch, metric) tuple
//!   from `training_loop.rs` after the HEALTH_DIAG line is emitted.
//! * Unknown metric names return `None` (silent OK — covered by Invariant 7
//!   audit; documented behaviour, not a stub).
//! * In-band values reset both warn and error counters for that metric.
//! * Out-of-warn-band values increment the warn counter. Crossing
//!   `consecutive_epochs_for_warn` emits a `tracing::warn!` log line every
//!   subsequent epoch the streak persists.
//! * Out-of-error-band values increment the error counter. When the counter
//!   reaches `2 × consecutive_epochs_for_warn` (i.e. `consecutive_epochs_for_error`
//!   from the toml `[settings]` block), the function returns
//!   `Some(TerminationReason::RegressionError { .. })` and the trainer
//!   propagates it as an `Err(CommonError::RegressionDetected)`.
//!
//! Termination is clean — no SIGTERM is raised; the error itself is the
//! signal. The caller's `?` operator unwinds the call stack normally and the
//! trainer's `Drop` impl tears down CUDA resources.

use std::collections::HashMap;
use std::path::Path;

use anyhow::{anyhow, Context, Result};
use serde::Deserialize;

/// Per-metric warn/error band. Inclusive bounds: `value < low || value > high`
/// triggers the band.
#[derive(Debug, Clone, Copy, PartialEq, Deserialize)]
pub(crate) struct MetricBands {
    pub warn_low: f64,
    pub warn_high: f64,
    pub error_low: f64,
    pub error_high: f64,
}

/// Streak thresholds loaded from the `[settings]` block of metric-bands.toml.
#[derive(Debug, Clone, Copy, PartialEq, Deserialize)]
pub(crate) struct BandSettings {
    /// Number of consecutive out-of-warn-band epochs before logging a warning.
    pub consecutive_epochs_for_warn: u32,
    /// Number of consecutive out-of-error-band epochs before terminating
    /// training. Per the Plan 5 design the file always satisfies
    /// `consecutive_epochs_for_error == 2 * consecutive_epochs_for_warn` (the
    /// "2N" rule), but the registry simply respects whatever the file states.
    pub consecutive_epochs_for_error: u32,
}

impl Default for BandSettings {
    /// Conservative defaults if the toml is absent or has no `[settings]`
    /// block. Match `config/metric-bands.toml` shipped values.
    fn default() -> Self {
        Self {
            consecutive_epochs_for_warn: 3,
            consecutive_epochs_for_error: 6,
        }
    }
}

/// Reason the training loop should self-terminate. Carried back to
/// `training_loop.rs`, then translated into a `CommonError::RegressionDetected`
/// for clean unwinding.
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum TerminationReason {
    /// Metric stayed in the warn band for `>= consecutive_epochs_for_warn`
    /// epochs but never crossed into the error band. Currently unused as a
    /// termination trigger (warn streaks only log); kept for symmetry so the
    /// caller can pattern-match exhaustively if a future change wants to
    /// promote warn-band streaks to terminations.
    #[allow(dead_code)]
    RegressionWarn {
        metric: String,
        value: f64,
        band: MetricBands,
        consecutive: u32,
    },
    /// Metric stayed in the error band for `>= consecutive_epochs_for_error`
    /// epochs. The trainer must terminate.
    RegressionError {
        metric: String,
        value: f64,
        band: MetricBands,
        consecutive: u32,
    },
}

impl std::fmt::Display for TerminationReason {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::RegressionWarn { metric, value, band, consecutive } => write!(
                f,
                "WARN streak {consecutive}: metric `{metric}` = {value:.6e} \
                 outside warn band [{:.6e}, {:.6e}]",
                band.warn_low, band.warn_high
            ),
            Self::RegressionError { metric, value, band, consecutive } => write!(
                f,
                "TERMINATED_BY_REGRESSION: streak {consecutive}: metric `{metric}` \
                 = {value:.6e} outside error band [{:.6e}, {:.6e}]",
                band.error_low, band.error_high
            ),
        }
    }
}

/// On-disk shape of `metric-bands.toml` — `[settings]` block + flat
/// `[bands]` table mapping metric names → band quadruples.
#[derive(Debug, Clone, Deserialize, Default)]
struct MetricBandsTomlSchema {
    #[serde(default)]
    settings: Option<BandSettings>,
    #[serde(default)]
    bands: HashMap<String, MetricBands>,
}

/// Per-metric registry that owns:
/// * the loaded bands (immutable after construction),
/// * the running streak counters (mutated each epoch),
/// * the streak-thresholds (immutable after construction).
#[derive(Debug, Clone)]
pub(crate) struct MetricBandsRegistry {
    bands: HashMap<String, MetricBands>,
    consecutive_warn: HashMap<String, u32>,
    consecutive_error: HashMap<String, u32>,
    settings: BandSettings,
}

impl MetricBandsRegistry {
    /// Construct a registry with explicit bands + settings (used by tests
    /// that don't want a toml file on disk).
    pub(crate) fn new(bands: HashMap<String, MetricBands>, settings: BandSettings) -> Self {
        Self {
            bands,
            consecutive_warn: HashMap::new(),
            consecutive_error: HashMap::new(),
            settings,
        }
    }

    /// Construct an empty registry (no bands; `update_and_check` returns
    /// `None` for every input). Used as the default field value when no
    /// bands toml is configured — callers can opt-in by replacing the
    /// registry via `DQNTrainer::set_metric_bands_registry`.
    pub(crate) fn empty() -> Self {
        Self::new(HashMap::new(), BandSettings::default())
    }

    /// Clear consecutive-warn and consecutive-error streak counters for every
    /// known metric. Bands and settings are preserved.
    ///
    /// Called from `DQNTrainer::reset_for_fold` so different walk-forward
    /// folds — which are independent training runs over disjoint date ranges
    /// — never share streaks. Without this, a fold-2 epoch in the error band
    /// could trigger termination because fold-0 and fold-1 each contributed
    /// some out-of-band epochs that accumulated into the count, even though
    /// neither fold individually crossed the 2N threshold.
    pub(crate) fn reset_streaks(&mut self) {
        self.consecutive_warn.clear();
        self.consecutive_error.clear();
    }

    /// Load bands + settings from a toml file. The `[bands]` section may be
    /// empty (returns an empty registry); a missing file is a hard error.
    pub(crate) fn load_from_toml(path: &Path) -> Result<Self> {
        let raw = std::fs::read_to_string(path).with_context(|| {
            format!(
                "MetricBandsRegistry: failed to read metric-bands.toml at {}",
                path.display()
            )
        })?;
        let parsed: MetricBandsTomlSchema = toml::from_str(&raw).with_context(|| {
            format!(
                "MetricBandsRegistry: failed to parse metric-bands.toml at {}",
                path.display()
            )
        })?;
        let settings = parsed.settings.unwrap_or_default();
        // Sanity-check bands: low <= high on every metric (otherwise the
        // band semantics are undefined and silently never trigger).
        for (name, band) in &parsed.bands {
            if !(band.warn_low <= band.warn_high) || !(band.error_low <= band.error_high) {
                return Err(anyhow!(
                    "MetricBandsRegistry: band for `{name}` has low > high \
                     (warn=[{}, {}], error=[{}, {}]) — would never trigger",
                    band.warn_low, band.warn_high, band.error_low, band.error_high
                ));
            }
            if band.error_low > band.warn_low || band.error_high < band.warn_high {
                return Err(anyhow!(
                    "MetricBandsRegistry: band for `{name}` has error band tighter \
                     than warn band (warn=[{}, {}], error=[{}, {}]) — invariant \
                     violation: error band must envelop warn band",
                    band.warn_low, band.warn_high, band.error_low, band.error_high
                ));
            }
        }
        Ok(Self::new(parsed.bands, settings))
    }

    /// Number of metrics currently watched (zero when registry is empty).
    pub(crate) fn band_count(&self) -> usize {
        self.bands.len()
    }

    /// Settings copy (consumers may want to log them at startup).
    pub(crate) fn settings(&self) -> BandSettings {
        self.settings
    }

    /// Update the per-metric streak counters with this epoch's value and
    /// return `Some(TerminationReason::RegressionError { .. })` iff the
    /// metric crossed the `consecutive_epochs_for_error` threshold.
    ///
    /// Unknown metric names return `None` immediately. NaN / infinite values
    /// are treated as out-of-band on both sides.
    pub(crate) fn update_and_check(
        &mut self,
        metric: &str,
        value: f64,
    ) -> Option<TerminationReason> {
        let bands = self.bands.get(metric)?;
        let bands = *bands; // copy to drop the borrow before mutating maps

        let in_error_band = value.is_finite()
            && value >= bands.error_low
            && value <= bands.error_high;

        if !in_error_band {
            let counter = self.consecutive_error.entry(metric.to_owned()).or_insert(0);
            *counter += 1;
            if *counter >= self.settings.consecutive_epochs_for_error {
                let consecutive = *counter;
                // Reset so that if the caller decides not to terminate (and a
                // future change re-enters the loop) the streak starts over.
                // For the production path this is moot — the trainer returns
                // Err right after — but it keeps the registry in a sane state
                // for unit tests that probe behaviour after a termination.
                self.consecutive_error.insert(metric.to_owned(), 0);
                self.consecutive_warn.insert(metric.to_owned(), 0);
                return Some(TerminationReason::RegressionError {
                    metric: metric.to_owned(),
                    value,
                    band: bands,
                    consecutive,
                });
            }
        } else {
            self.consecutive_error.insert(metric.to_owned(), 0);
        }

        let in_warn_band = value.is_finite()
            && value >= bands.warn_low
            && value <= bands.warn_high;

        if !in_warn_band {
            let counter = self.consecutive_warn.entry(metric.to_owned()).or_insert(0);
            *counter += 1;
            if *counter >= self.settings.consecutive_epochs_for_warn {
                tracing::warn!(
                    target: "dqn::regression",
                    metric = metric,
                    value = value,
                    consecutive = *counter,
                    warn_low = bands.warn_low,
                    warn_high = bands.warn_high,
                    "Metric `{metric}` outside warn band for {} consecutive epochs (value={:.6e})",
                    *counter, value,
                );
            }
        } else {
            self.consecutive_warn.insert(metric.to_owned(), 0);
        }

        None
    }
}

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

    fn band(warn_low: f64, warn_high: f64, error_low: f64, error_high: f64) -> MetricBands {
        MetricBands { warn_low, warn_high, error_low, error_high }
    }

    fn bands_one(name: &str, b: MetricBands) -> HashMap<String, MetricBands> {
        let mut m = HashMap::new();
        m.insert(name.to_owned(), b);
        m
    }

    #[test]
    fn unknown_metric_returns_none() {
        let mut reg = MetricBandsRegistry::new(
            bands_one("known", band(0.0, 1.0, -1.0, 2.0)),
            BandSettings { consecutive_epochs_for_warn: 3, consecutive_epochs_for_error: 6 },
        );
        for v in [-99.0, 0.5, 1e30] {
            assert!(reg.update_and_check("not_in_toml", v).is_none());
        }
    }

    #[test]
    fn in_band_resets_streaks() {
        let mut reg = MetricBandsRegistry::new(
            bands_one("m", band(0.0, 1.0, -1.0, 2.0)),
            BandSettings { consecutive_epochs_for_warn: 3, consecutive_epochs_for_error: 6 },
        );
        // 5 out-of-error-band epochs (1 short of termination)
        for _ in 0..5 {
            assert!(reg.update_and_check("m", 1e9).is_none());
        }
        // One in-band epoch — resets the error streak
        assert!(reg.update_and_check("m", 0.5).is_none());
        // Another 5 out-of-error-band epochs must still NOT terminate (streak reset)
        for _ in 0..5 {
            assert!(reg.update_and_check("m", 1e9).is_none());
        }
    }

    #[test]
    fn reset_streaks_clears_consecutive_counters() {
        // P5T5 Phase E: walk-forward folds are independent training runs;
        // reset_streaks must clear cross-fold accumulation so a fold-N
        // out-of-band epoch can't trip termination off counts inherited
        // from folds 0..N-1.
        let mut reg = MetricBandsRegistry::new(
            bands_one("m", band(-1.0, 1.0, -100.0, 100.0)),
            BandSettings { consecutive_epochs_for_warn: 3, consecutive_epochs_for_error: 6 },
        );
        // Accumulate 5 out-of-error-band epochs (1 short of termination)
        for _ in 0..5 {
            assert!(reg.update_and_check("m", 1e9).is_none());
        }
        // Fold boundary
        reg.reset_streaks();
        // Streak must now be zero — 5 more out-of-band epochs still no
        // termination (would have terminated at the 1st without the reset)
        for i in 0..5 {
            let result = reg.update_and_check("m", 1e9);
            assert!(
                result.is_none(),
                "post-reset epoch {i}: expected None (streak < 6), got {result:?}",
            );
        }
        // 6th post-reset epoch terminates
        let term = reg
            .update_and_check("m", 1e9)
            .expect("expected RegressionError on 6th post-reset out-of-band epoch");
        assert!(matches!(term, TerminationReason::RegressionError { .. }));
    }

    #[test]
    fn six_consecutive_error_epochs_terminate() {
        let mut reg = MetricBandsRegistry::new(
            bands_one("train_loss", band(1e-6, 10.0, 0.0, 1000.0)),
            BandSettings { consecutive_epochs_for_warn: 3, consecutive_epochs_for_error: 6 },
        );
        // 5 epochs of saturated loss — streak below threshold
        for i in 0..5 {
            let result = reg.update_and_check("train_loss", 1e15);
            assert!(
                result.is_none(),
                "epoch {i}: expected None (streak < 6), got {result:?}",
            );
        }
        // 6th consecutive epoch — must terminate
        let term = reg
            .update_and_check("train_loss", 1e15)
            .expect("expected RegressionError on 6th out-of-band epoch");
        match term {
            TerminationReason::RegressionError { metric, consecutive, .. } => {
                assert_eq!(metric, "train_loss");
                assert_eq!(consecutive, 6);
            }
            other => panic!("expected RegressionError, got {other:?}"),
        }
    }

    #[test]
    fn nan_treated_as_out_of_band() {
        let mut reg = MetricBandsRegistry::new(
            bands_one("m", band(-100.0, 100.0, -1000.0, 1000.0)),
            BandSettings { consecutive_epochs_for_warn: 2, consecutive_epochs_for_error: 4 },
        );
        for _ in 0..3 {
            assert!(reg.update_and_check("m", f64::NAN).is_none());
        }
        let term = reg
            .update_and_check("m", f64::NAN)
            .expect("NaN must count as out-of-band");
        assert!(matches!(term, TerminationReason::RegressionError { .. }));
    }

    #[test]
    fn warn_band_does_not_terminate_alone() {
        // Warn band [0, 1], error band [-100, 100]. value=5 is in warn-only.
        let mut reg = MetricBandsRegistry::new(
            bands_one("m", band(0.0, 1.0, -100.0, 100.0)),
            BandSettings { consecutive_epochs_for_warn: 2, consecutive_epochs_for_error: 4 },
        );
        for _ in 0..20 {
            assert!(reg.update_and_check("m", 5.0).is_none());
        }
    }

    #[test]
    fn loading_invalid_toml_rejected() {
        // Construct an inverted band directly — MetricBandsRegistry::new would
        // accept it, but load_from_toml's validator rejects.
        let toml_text = r#"
[settings]
consecutive_epochs_for_warn = 3
consecutive_epochs_for_error = 6

[bands.bad]
warn_low = 10.0
warn_high = -10.0
error_low = -100.0
error_high = 100.0
"#;
        let dir = std::env::temp_dir();
        let path = dir.join("metric-bands-invalid-test.toml");
        std::fs::write(&path, toml_text).unwrap();
        let err = MetricBandsRegistry::load_from_toml(&path).unwrap_err();
        let msg = format!("{err:#}");
        assert!(
            msg.contains("low > high") || msg.contains("invariant violation"),
            "unexpected error message: {msg}",
        );
        let _ = std::fs::remove_file(&path);
    }
}
