Northern-Thailand-Ping-River-Monitor/rust/collector/src/service.rs

//! The long-running collection loop and the one-shot commands.

use crate::config::{Config, USER_AGENT};
use crate::dams::DamsClient;
use crate::db;
use crate::forecasting::{self, Forecaster, RunInfo};
use crate::hii::HiiClient;
use crate::notify;
use crate::rain::OpenMeteoClient;
use crate::rid::{self, CycleOutcome, RidClient};
use crate::stations::Stations;
use crate::util::{bangkok_now, floor_hour};
use anyhow::{Result, bail};
use chrono::{Duration, NaiveDateTime};
use ping_forecast::ModelStore;
use ping_forecast::rain::RainSeries;
use serde_json::json;
use sqlx::PgPool;
use std::sync::{Arc, RwLock};

pub struct Collector {
    pub cfg: Config,
    pub pool: PgPool,
    /// Shared with the website (the station CRUD API edits it).
    pub stations: Arc<RwLock<Stations>>,
    /// The last status written, shared with /scraping/status.
    pub status: Arc<RwLock<serde_json::Value>>,
    /// POST /scrape/trigger: poll now.
    pub trigger: Arc<tokio::sync::Notify>,
    pub rid: RidClient,
    hii: HiiClient,
    dams: DamsClient,
    rain: OpenMeteoClient,
    http: reqwest::Client,
}

#[derive(Default)]
struct Stats {
    total_runs: u64,
    successful_runs: u64,
    failed_runs: u64,
    last_run: Option<NaiveDateTime>,
    next_run: Option<NaiveDateTime>,
    feed_current: Option<bool>,
    latest_reading: Option<NaiveDateTime>,
    readings_saved: usize,
    forecast: Option<RunInfo>,
}

fn iso(t: Option<NaiveDateTime>) -> serde_json::Value {
    t.map_or(serde_json::Value::Null, |t| {
        json!(t.format("%Y-%m-%dT%H:%M:%S").to_string())
    })
}

impl Collector {
    pub fn new(cfg: Config, pool: PgPool, stations: Stations) -> Result<Self> {
        let http = reqwest::Client::builder()
            .user_agent(USER_AGENT)
            .connect_timeout(std::time::Duration::from_secs(20))
            .build()?;
        Ok(Self {
            rid: RidClient::new(http.clone(), cfg.rid_api_url.clone()),
            hii: HiiClient::new(http.clone(), cfg.hii_api_base.clone()),
            dams: DamsClient::new(http.clone(), cfg.dams_url.clone()),
            rain: OpenMeteoClient::new(http.clone(), cfg.openmeteo_url.clone()),
            http,
            cfg,
            pool,
            stations: Arc::new(RwLock::new(stations)),
            status: Arc::new(RwLock::new(json!({"is_running": false}))),
            trigger: Arc::new(tokio::sync::Notify::new()),
        })
    }

    /// The station mapping as of now (the API may edit it between cycles).
    pub fn stations(&self) -> Stations {
        self.stations.read().unwrap().clone()
    }

    /// The website/API state over this collector's pool, stations and status.
    pub fn web_state(&self, in_process: bool) -> crate::web::Shared {
        crate::web::AppState::new(
            self.cfg.clone(),
            self.pool.clone(),
            self.http.clone(),
            self.stations.clone(),
            in_process.then(|| self.status.clone()),
            in_process.then(|| self.trigger.clone()),
        )
    }

    /// `ping-collector serve`: the website and API only; /scraping/status
    /// reads the status file the collector service writes.
    pub async fn serve(&self, listen: String) -> Result<()> {
        db::ensure_schema(&self.pool).await?;
        crate::web::serve(self.web_state(false), listen, shutdown_signal()).await
    }

    /// HII rain + water level, RID dams (today and yesterday), Open-Meteo.
    /// Each source is failure-isolated.
    pub async fn collect_hourly(&self, now: NaiveDateTime) {
        if self.cfg.hii_enabled {
            match self.hii.fetch_rain(self.cfg.hii_basin_code).await {
                Ok(r) => match db::save_hii_rain(&self.pool, &r, now).await {
                    Ok(n) => tracing::info!("HII rainfall: {n} rows saved"),
                    Err(e) => tracing::error!("HII rainfall save failed: {e:#}"),
                },
                Err(e) => tracing::error!("HII rainfall fetch failed: {e:#}"),
            }
            match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await {
                Ok(r) => match db::save_hii_waterlevel(&self.pool, &r, now).await {
                    Ok(n) => tracing::info!("HII water level: {n} rows saved"),
                    Err(e) => tracing::error!("HII water level save failed: {e:#}"),
                },
                Err(e) => tracing::error!("HII water level fetch failed: {e:#}"),
            }
        }
        if self.cfg.dams_enabled {
            for date in [None, Some(now.date() - Duration::days(1))] {
                match self.dams.fetch_day(date).await {
                    Ok(r) => match db::save_dams(&self.pool, &r, now).await {
                        Ok(n) => tracing::info!("RID dams {date:?}: {n} dam-day rows saved"),
                        Err(e) => tracing::error!("RID dams save failed: {e:#}"),
                    },
                    Err(e) => tracing::error!("RID dams fetch {date:?} failed: {e:#}"),
                }
            }
        }
        if self.cfg.rain_enabled {
            match self.rain.fetch_forecast(5, 2).await {
                Ok(r) => match db::save_rain(&self.pool, &r).await {
                    Ok(n) => tracing::info!("Open-Meteo rain: {n} hourly rows upserted"),
                    Err(e) => tracing::error!("Open-Meteo save failed: {e:#}"),
                },
                Err(e) => tracing::error!("Open-Meteo fetch failed: {e:#}"),
            }
        }
    }

    fn write_status(&self, stats: &Stats, running: bool) {
        let body = json!({
            "is_running": running,
            "total_runs": stats.total_runs,
            "successful_runs": stats.successful_runs,
            "failed_runs": stats.failed_runs,
            "last_run": iso(stats.last_run),
            "next_run": iso(stats.next_run),
            "feed_current": stats.feed_current,
            "latest_reading": iso(stats.latest_reading),
            "readings_saved": stats.readings_saved,
            "collector": concat!("ping-collector ", env!("CARGO_PKG_VERSION")),
            "forecast": stats.forecast.as_ref().map(|f| json!({
                "last_run": iso(Some(f.at)),
                "rows": f.rows,
                "model_rows": f.model_rows,
                "models": f.models,
                "notifications_sent": f.notified,
            })),
        });
        *self.status.write().unwrap() = body.clone();
        let Some(path) = &self.cfg.status_file else {
            return;
        };
        let tmp = path.with_extension("json.tmp");
        let result =
            std::fs::write(&tmp, body.to_string()).and_then(|_| std::fs::rename(&tmp, path));
        if let Err(e) = result {
            tracing::warn!("could not write status file {}: {e}", path.display());
        }
    }

    /// The service: startup gap fill, then RID polls on the hour (+offset)
    /// with retries while RID is behind, and the hourly sources once per
    /// clock hour. Returns when a shutdown signal arrives.
    pub async fn run(&self) -> Result<()> {
        db::ensure_schema(&self.pool).await?;
        let shutdown = shutdown_signal();
        tokio::pin!(shutdown);
        // the website, when WEB_LISTEN is set: stopped with the collector
        let (stop_web, web_stopped) = tokio::sync::watch::channel(false);
        let web = self.cfg.web.listen.clone().map(|listen| {
            let state = self.web_state(true);
            let mut rx = web_stopped.clone();
            tokio::spawn(async move {
                let stop = async move {
                    let _ = rx.wait_for(|stop| *stop).await;
                };
                if let Err(e) = crate::web::serve(state, listen, stop).await {
                    tracing::error!("website stopped: {e:#}");
                }
            })
        });
        let result = self.collect_loop(&mut shutdown).await;
        let _ = stop_web.send(true);
        if let Some(handle) = web {
            let _ = tokio::time::timeout(std::time::Duration::from_secs(10), handle).await;
        }
        result
    }

    async fn collect_loop(
        &self,
        shutdown: &mut std::pin::Pin<&mut impl std::future::Future<Output = ()>>,
    ) -> Result<()> {
        if self.cfg.gap_fill_days_on_start > 0 {
            let days = Some(self.cfg.gap_fill_days_on_start);
            let stations = self.stations();
            tokio::select! {
                r = rid::fill_gaps(&self.rid, &self.pool, &stations, days, bangkok_now()) => {
                    if let Err(e) = r { tracing::error!("startup gap fill failed: {e:#}") }
                }
                _ = &mut *shutdown => return Ok(()),
            }
        }

        let mut stats = Stats::default();
        let mut last_hourly_slot: Option<NaiveDateTime> = None;
        let mut forecaster = self.cfg.forecast_enabled.then(|| {
            tracing::info!(
                "FORECAST_IN_COLLECTOR: forecasting from {} and notifying here",
                self.cfg.models_dir.display()
            );
            Forecaster::new(&self.cfg, self.http.clone())
        });
        loop {
            let start = bangkok_now();
            self.write_status(&stats, true);
            let outcome: CycleOutcome =
                rid::run_cycle(&self.rid, &self.pool, &self.stations(), start).await;
            stats.total_runs += 1;
            stats.last_run = Some(start);
            if outcome.current {
                stats.successful_runs += 1;
            } else {
                stats.failed_runs += 1;
            }
            stats.feed_current = Some(outcome.current);
            stats.readings_saved = outcome.saved;
            if outcome.latest.is_some() {
                stats.latest_reading = outcome.latest.max(stats.latest_reading);
            }

            let slot = floor_hour(bangkok_now());
            if last_hourly_slot != Some(slot) {
                last_hourly_slot = Some(slot);
                self.collect_hourly(bangkok_now()).await;
            }
            if let Some(f) = forecaster.as_mut()
                && let Some(info) = f.maybe_run(&self.pool, &self.rain, bangkok_now()).await
            {
                stats.forecast = Some(info);
            }

            let now = bangkok_now();
            let next = rid::next_poll_time(
                now,
                outcome.current,
                self.cfg.poll_offset_minutes,
                self.cfg.retry_minutes,
            );
            stats.next_run = Some(next);
            self.write_status(&stats, false);
            tracing::info!(
                "next RID poll at {}{}",
                next.format("%H:%M"),
                if outcome.current {
                    ""
                } else {
                    " (current hour not published yet)"
                }
            );
            let wait = (next - now)
                .to_std()
                .unwrap_or_default()
                .max(std::time::Duration::from_secs(30));
            tokio::select! {
                _ = tokio::time::sleep(wait) => {}
                _ = self.trigger.notified() => tracing::info!("collection triggered via the API"),
                _ = &mut *shutdown => {
                    tracing::info!("shutdown signal received");
                    return Ok(());
                }
            }
        }
    }

    /// One cycle of every source; with `dry_run` nothing is written and the
    /// RID fetch is compared against what the database already holds.
    pub async fn once(&self, dry_run: bool) -> Result<()> {
        let now = bangkok_now();
        if !dry_run {
            db::ensure_schema(&self.pool).await?;
            let o = rid::run_cycle(&self.rid, &self.pool, &self.stations(), now).await;
            println!(
                "RID: reports {:?}, {} fetched, {} saved, latest {:?}, current hour in: {}{}",
                o.dates,
                o.fetched,
                o.saved,
                o.latest,
                o.current,
                o.error.map(|e| format!(", error: {e}")).unwrap_or_default()
            );
            self.collect_hourly(now).await;
            return Ok(());
        }

        let dates = rid::cycle_dates(&self.pool, now).await;
        let (readings, error) = rid::fetch_dates(&self.rid, &dates).await;
        let stored = db::count_stored_rid(&self.pool, &readings).await?;
        let latest = readings.iter().map(|r| r.timestamp).max();
        let db_latest = db::latest_rid_timestamp(&self.pool).await?;
        println!(
            "RID (dry run): reports {dates:?}: {} readings fetched, {stored} already stored, {} would be added; latest fetched {latest:?}, latest stored {db_latest:?}{}",
            readings.len(),
            readings.len() as i64 - stored,
            error.map(|e| format!(", error: {e}")).unwrap_or_default()
        );
        if self.cfg.hii_enabled {
            match self.hii.fetch_rain(self.cfg.hii_basin_code).await {
                Ok(r) => println!("HII rainfall (dry run): {} Ping-basin stations", r.len()),
                Err(e) => println!("HII rainfall (dry run): {e:#}"),
            }
            match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await {
                Ok(r) => println!("HII water level (dry run): {} Ping-basin stations", r.len()),
                Err(e) => println!("HII water level (dry run): {e:#}"),
            }
        }
        if self.cfg.dams_enabled {
            match self.dams.fetch_day(None).await {
                Ok(r) => println!("RID dams (dry run): {} dams today", r.len()),
                Err(e) => println!("RID dams (dry run): {e:#}"),
            }
        }
        if self.cfg.rain_enabled {
            match self.rain.fetch_forecast(5, 2).await {
                Ok(r) => println!("Open-Meteo (dry run): {} hourly rows", r.len()),
                Err(e) => println!("Open-Meteo (dry run): {e:#}"),
            }
        }
        Ok(())
    }
}

impl Collector {
    /// `ping-collector forecast`: compute from the database now (or at
    /// `now`), print a table or the /forecast JSON, store only with `save`.
    pub async fn forecast(
        &self,
        json_out: bool,
        save: bool,
        now: Option<NaiveDateTime>,
        rain: Option<RainSeries>,
    ) -> Result<()> {
        let mut store = ModelStore::new(&self.cfg.models_dir);
        let models = store.refresh()?;
        if models.is_empty() {
            bail!(
                "no exported models in {} (run `python -m src.ml.export`)",
                self.cfg.models_dir.display()
            );
        }
        let now = now.unwrap_or_else(bangkok_now);
        let snap = forecasting::compute(&self.pool, &self.rain, &models, now, rain).await?;
        if json_out {
            println!("{}", serde_json::to_string(&snap.forecasts)?);
        } else {
            println!(
                "{:<6} {:>3}  {:<9} {:<16} {:>7} {:>7} {:>9} {:>8}",
                "code", "h", "source", "as_of", "level", "peak", "p_warning", "p_danger"
            );
            for f in &snap.forecasts {
                println!(
                    "{:<6} {:>3}  {:<9} {:<16} {:>7.2} {:>7.2} {:>9.3} {:>8.3}",
                    f.station_code,
                    f.horizon_hours,
                    f.source,
                    f.as_of.format("%Y-%m-%d %H:%M"),
                    f.current_level,
                    f.predicted_max_level,
                    f.p_warning,
                    f.p_danger
                );
            }
            eprintln!(
                "{} models from {}; {}",
                models.len(),
                self.cfg.models_dir.display(),
                forecasting::summarize(&snap)
            );
        }
        if save {
            db::ensure_schema(&self.pool).await?;
            let n = db::save_forecast_history(&self.pool, &snap.forecasts).await?;
            db::save_forecast_current(&self.pool, now, &snap.forecasts).await?;
            eprintln!("stored: {n} forecast_history rows, forecast_current");
        }
        Ok(())
    }

    /// `ping-collector notify`: what would be (or, without `dry_run`, is)
    /// published right now against the stored notification state.
    pub async fn notify(&self, dry_run: bool) -> Result<()> {
        if !dry_run && !self.cfg.forecast_enabled {
            bail!(
                "the Python app notifies while FORECAST_IN_COLLECTOR is not true; \
                 use --dry-run"
            );
        }
        let Some(notifier) = forecasting::notifier(&self.cfg, self.http.clone()) else {
            bail!(
                "notifications are off: neither NTFY_SERVER nor MATRIX_ACCESS_TOKEN/MATRIX_ROOM_ID is set"
            );
        };
        let mut store = ModelStore::new(&self.cfg.models_dir);
        let models = store.refresh()?;
        let now = bangkok_now();
        let snap = forecasting::compute(&self.pool, &self.rain, &models, now, None).await?;
        let readings =
            notify::readings_with_hii(&db::latest_rid_readings(&self.pool).await?, &snap.fills);
        for r in &readings {
            println!(
                "{:<6} {} {:>7} {:?}{}",
                r.station_code,
                r.timestamp.format("%Y-%m-%d %H:%M"),
                r.water_level.map_or("-".into(), |l| format!("{l:.2}")),
                r.source,
                r.discharge_percent
                    .map_or(String::new(), |p| format!(" {p:.0}%"))
            );
        }
        db::ensure_schema(&self.pool).await?;
        if dry_run {
            let states = db::notification_states(&self.pool).await?;
            let mut plan = vec![];
            if notifier.ntfy.is_some() {
                plan.extend(
                    notify::evaluate(&readings, &snap.forecasts, &states, &notifier.settings, now)
                        .into_iter()
                        .map(|p| (p.notification.topic.clone(), p)),
                );
            }
            if notifier.matrix.is_some() {
                let mstates = notifier
                    .matrix_states(&self.pool, &states, now, false)
                    .await?;
                plan.extend(
                    notifier
                        .matrix_plan(&readings, &snap.forecasts, &mstates, now)
                        .into_iter()
                        .map(|p| ("matrix".to_string(), p)),
                );
            }
            if plan.is_empty() {
                println!("no transitions: nothing would be published");
            }
            for (to, p) in plan {
                println!(
                    "would publish {to}: {} -- {}",
                    p.notification.title, p.notification.message
                );
            }
        } else {
            for n in notifier
                .run(&self.pool, &readings, &snap.forecasts, now)
                .await?
            {
                println!("published {}: {}", n.topic, n.title);
            }
        }
        Ok(())
    }
}

async fn shutdown_signal() {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{SignalKind, signal};
        let mut term = signal(SignalKind::terminate()).expect("SIGTERM handler");
        tokio::select! {
            _ = tokio::signal::ctrl_c() => {}
            _ = term.recv() => {}
        }
    }
    #[cfg(not(unix))]
    {
        let _ = tokio::signal::ctrl_c().await;
    }
}