//! Constants shared with Python, generated from it: `flood_config.json` is
//! written by `uv run python -m src.ml.export --config`, and
//! tests/test_model_export.py fails while it differs from the Python source.

use serde::Deserialize;
use std::collections::BTreeMap;
use std::sync::OnceLock;

#[derive(Debug, Deserialize)]
pub struct Stage {
    pub stage: u32,
    pub level: f64,
    pub discharge_cms: f64,
}

#[derive(Debug, Deserialize)]
pub struct DamConfig {
    pub dam_id: String,
    pub report_hour: u32,
    pub ffill_limit_h: usize,
    pub serving_days: i64,
}

#[derive(Debug, Deserialize)]
pub struct NotifyConfig {
    pub clear_margin_m: f64,
    pub capacity_guard_min_pct: f64,
    pub capacity_guard_exempt: Vec<String>,
    pub outlook_on: f64,
    pub outlook_off: f64,
    pub outlook_horizon: u32,
    pub feed_stale_after_h: f64,
    pub station_names: BTreeMap<String, String>,
}

#[derive(Debug, Deserialize)]
pub struct FloodConfig {
    thresholds: BTreeMap<String, (f64, f64)>,
    pub flood_stages: BTreeMap<String, Vec<Stage>>,
    upstream_leads: BTreeMap<String, Vec<(String, i64)>>,
    pub dam_stations: Vec<String>,
    pub monsoon_months: Vec<u32>,
    pub ffill_limit_h: usize,
    pub basin_anchor: String,
    pub forward_rain_h: usize,
    pub horizons: Vec<u32>,
    pub stale_after_h: f64,
    pub heuristic_sigma: f64,
    pub heuristic_version: String,
    pub latest_window_h: i64,
    pub hii_fill_stations: Vec<String>,
    pub hii_exact_mirrors: Vec<String>,
    pub hii_min_overlap_hours: usize,
    pub dam: DamConfig,
    pub notify: NotifyConfig,
}

impl FloodConfig {
    /// (warning, danger) for a station, or the "*" default.
    pub fn thresholds(&self, station: &str) -> (f64, f64) {
        self.thresholds
            .get(station)
            .or_else(|| self.thresholds.get("*"))
            .copied()
            .unwrap_or((f64::INFINITY, f64::INFINITY))
    }

    /// (upstream station, lead hours) pairs, in Python's order.
    pub fn upstream_leads(&self, station: &str) -> &[(String, i64)] {
        self.upstream_leads.get(station).map_or(&[], Vec::as_slice)
    }

    /// The stations the models are built for (Python's `UPSTREAM_LEADS` keys).
    pub fn forecast_stations(&self) -> impl Iterator<Item = &str> {
        self.upstream_leads.keys().map(String::as_str)
    }

    pub fn is_dam_station(&self, station: &str) -> bool {
        self.dam_stations.iter().any(|s| s == station)
    }
}

pub fn config() -> &'static FloodConfig {
    static CONFIG: OnceLock<FloodConfig> = OnceLock::new();
    CONFIG.get_or_init(|| {
        serde_json::from_str(include_str!("../flood_config.json"))
            .expect("flood_config.json is generated by src/ml/export.py and must parse")
    })
}

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

    #[test]
    fn generated_config_parses_and_has_the_basin() {
        let c = config();
        assert_eq!(c.thresholds("P.1"), (3.7, 4.2));
        assert_eq!(c.thresholds("P.999"), (3.0, 4.5));
        assert_eq!(c.upstream_leads("P.1")[0], ("P.103".to_string(), 1));
        assert!(c.upstream_leads("P.82").is_empty());
        assert_eq!(c.forecast_stations().count(), 16);
        assert_eq!(c.horizons, vec![6, 12, 24]);
        assert!(c.is_dam_station("P.1") && !c.is_dam_station("P.20"));
    }
}