//! PostgreSQL access: the same tables and upsert semantics as the Python
//! collectors (src/database_adapters.py, hii_collector.py, rid_reservoir.py,
//! ml/rain.py), written as one set-based statement per table (UNNEST) instead
//! of a statement per row.

use crate::dams::DamRecord;
use crate::hii::{RainRecord, WaterLevelRecord};
use crate::rain::RainRow;
use crate::rid::Reading;
use crate::stations::Stations;
use anyhow::{Context, Result};
use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
use sqlx::PgPool;
use sqlx::postgres::PgPoolOptions;
use std::collections::{BTreeMap, HashMap, HashSet};

pub async fn connect(url: &str, max_connections: u32) -> Result<PgPool> {
    use sqlx::ConnectOptions;
    use std::str::FromStr;
    // sqlx warns about any statement over 1 s; the whole-history gap scan
    // legitimately takes several seconds on ~1M rows.
    let options = sqlx::postgres::PgConnectOptions::from_str(url)
        .context("invalid database URL")?
        .log_slow_statements(log::LevelFilter::Warn, std::time::Duration::from_secs(30));
    PgPoolOptions::new()
        .max_connections(max_connections)
        .acquire_timeout(std::time::Duration::from_secs(30))
        .connect_with(options)
        .await
        .context("connecting to PostgreSQL")
}

/// Idempotent DDL, identical to what the Python app created and the Python
/// backfill stores (src/hii_collector.py, src/rid_reservoir.py) still create,
/// so either can start first against an empty database.
const SCHEMA: &[&str] = &[
    "CREATE TABLE IF NOT EXISTS stations (
        id SERIAL PRIMARY KEY,
        station_code VARCHAR(10) UNIQUE NOT NULL,
        thai_name VARCHAR(255) NOT NULL,
        english_name VARCHAR(255) NOT NULL,
        latitude DECIMAL(10,8),
        longitude DECIMAL(11,8),
        geohash VARCHAR(20),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
    "CREATE TABLE IF NOT EXISTS water_measurements (
        id BIGSERIAL PRIMARY KEY,
        timestamp TIMESTAMP NOT NULL,
        station_id INTEGER NOT NULL,
        water_level NUMERIC(10,3),
        discharge NUMERIC(10,2),
        discharge_percent NUMERIC(5,2),
        status VARCHAR(20) DEFAULT 'active',
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (station_id) REFERENCES stations(id),
        UNIQUE(timestamp, station_id))",
    "CREATE INDEX IF NOT EXISTS idx_timestamp ON water_measurements(timestamp)",
    "CREATE INDEX IF NOT EXISTS idx_station_timestamp ON water_measurements(station_id, timestamp DESC)",
    "CREATE TABLE IF NOT EXISTS hii_rain_stations (
        id INTEGER PRIMARY KEY,
        oldcode VARCHAR(60),
        name_th VARCHAR(255),
        name_en VARCHAR(255),
        latitude NUMERIC(10,6),
        longitude NUMERIC(10,6),
        sub_basin_id VARCHAR(10),
        agency VARCHAR(40),
        updated_at TIMESTAMP)",
    "CREATE TABLE IF NOT EXISTS hii_rainfall (
        station_id INTEGER NOT NULL,
        timestamp TIMESTAMP NOT NULL,
        rain_1h NUMERIC(7,2),
        rain_24h NUMERIC(8,2),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (station_id, timestamp))",
    "CREATE TABLE IF NOT EXISTS hii_wl_stations (
        id INTEGER PRIMARY KEY,
        oldcode VARCHAR(60),
        rid_code VARCHAR(10),
        name_th VARCHAR(255),
        name_en VARCHAR(255),
        latitude NUMERIC(10,6),
        longitude NUMERIC(10,6),
        agency VARCHAR(40),
        river_name VARCHAR(255),
        offset_msl NUMERIC(8,3),
        ground_level_msl NUMERIC(8,3),
        min_bank_msl NUMERIC(8,3),
        critical_level_msl NUMERIC(8,3),
        critical_level_m NUMERIC(8,3),
        qmax NUMERIC(10,2),
        is_key_station BOOLEAN,
        updated_at TIMESTAMP)",
    "CREATE TABLE IF NOT EXISTS hii_waterlevel (
        station_id INTEGER NOT NULL,
        timestamp TIMESTAMP NOT NULL,
        wl_msl NUMERIC(8,3),
        wl_m NUMERIC(8,3),
        discharge NUMERIC(10,2),
        flow_rate NUMERIC(10,2),
        storage_percent NUMERIC(6,2),
        situation_level INTEGER,
        diff_wl_bank NUMERIC(8,3),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (station_id, timestamp))",
    "CREATE INDEX IF NOT EXISTS idx_hii_rainfall_ts ON hii_rainfall(timestamp)",
    "CREATE INDEX IF NOT EXISTS idx_hii_waterlevel_ts ON hii_waterlevel(timestamp)",
    "CREATE TABLE IF NOT EXISTS rid_dams (
        dam_id VARCHAR(10) PRIMARY KEY,
        region VARCHAR(40),
        name_th VARCHAR(255),
        latitude NUMERIC(10,6),
        longitude NUMERIC(10,6),
        capacity_max_mcm NUMERIC(10,2),
        capacity_normal_mcm NUMERIC(10,2),
        updated_at TIMESTAMP)",
    "CREATE TABLE IF NOT EXISTS rid_reservoir_daily (
        dam_id VARCHAR(10) NOT NULL,
        date DATE NOT NULL,
        storage_mcm NUMERIC(10,2),
        storage_pct NUMERIC(8,2),
        inflow_mcm NUMERIC(10,2),
        outflow_mcm NUMERIC(10,2),
        level_msl NUMERIC(8,2),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (dam_id, date))",
    "CREATE INDEX IF NOT EXISTS idx_rid_reservoir_date ON rid_reservoir_daily(date)",
    "CREATE TABLE IF NOT EXISTS openmeteo_rain (
        timestamp TIMESTAMP PRIMARY KEY,
        chiang_dao NUMERIC(6,2), mae_taeng NUMERIC(6,2), mae_ngat NUMERIC(6,2),
        mae_rim NUMERIC(6,2), chiang_mai NUMERIC(6,2),
        catchment_mean NUMERIC(6,2),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
    // src/forecast_history.py
    "CREATE TABLE IF NOT EXISTS forecast_history (
        as_of TIMESTAMP NOT NULL,
        station_code VARCHAR(10) NOT NULL,
        horizon_hours INTEGER NOT NULL,
        predicted_max_level NUMERIC(8,3),
        p_warning NUMERIC(7,5),
        p_danger NUMERIC(7,5),
        current_level NUMERIC(8,3),
        model_version VARCHAR(64),
        source VARCHAR(16),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (as_of, station_code, horizon_hours))",
    "CREATE INDEX IF NOT EXISTS idx_forecast_history_station ON forecast_history(station_code, as_of)",
    // notify.rs
    "CREATE TABLE IF NOT EXISTS notification_state (
        key VARCHAR(64) PRIMARY KEY, state VARCHAR(16) NOT NULL,
        value NUMERIC(8,3), updated_at TIMESTAMP NOT NULL)",
    // The current /forecast payload (one row, id 'all'), served by
    // `ping-collector serve` when the website runs as its own process.
    "CREATE TABLE IF NOT EXISTS forecast_current (
        id VARCHAR(16) PRIMARY KEY,
        computed_at TIMESTAMP NOT NULL,
        model_versions TEXT,
        payload TEXT NOT NULL)",
];

pub async fn ensure_schema(pool: &PgPool) -> Result<()> {
    for ddl in SCHEMA {
        sqlx::query(ddl)
            .execute(pool)
            .await
            .with_context(|| format!("DDL failed: {ddl}"))?;
    }
    Ok(())
}

/// Keep the last record per key (a later report revision wins), preserving
/// first-seen order; Postgres refuses to upsert the same key twice in one
/// statement.
fn dedupe_last<T, K: std::hash::Hash + Eq + Clone>(items: &[T], key: impl Fn(&T) -> K) -> Vec<&T> {
    let mut index: HashMap<K, usize> = HashMap::new();
    let mut out: Vec<&T> = Vec::new();
    for item in items {
        let k = key(item);
        match index.get(&k) {
            Some(&i) => out[i] = item,
            None => {
                index.insert(k, out.len());
                out.push(item);
            }
        }
    }
    out
}

pub async fn save_rid(pool: &PgPool, readings: &[Reading], stations: &Stations) -> Result<usize> {
    let rows = dedupe_last(readings, |r| (r.timestamp, r.station_num));
    let station_nums: Vec<u32> = rows
        .iter()
        .map(|r| r.station_num)
        .collect::<std::collections::BTreeSet<_>>()
        .into_iter()
        .collect();
    let infos: Vec<_> = station_nums.iter().map(|&n| stations.get(n)).collect();

    let mut tx = pool.begin().await?;
    sqlx::query(
        "INSERT INTO stations (id, station_code, thai_name, english_name, latitude, longitude, geohash, updated_at)
         SELECT *, NOW() FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::float8[], $6::float8[], $7::text[])
         ON CONFLICT (id) DO UPDATE SET
            thai_name = EXCLUDED.thai_name, english_name = EXCLUDED.english_name,
            latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
            geohash = EXCLUDED.geohash, updated_at = NOW()",
    )
    .bind(station_nums.iter().map(|&n| n as i32).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.code.clone()).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.thai_name.clone()).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.english_name.clone()).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.latitude).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.longitude).collect::<Vec<_>>())
    .bind(infos.iter().map(|s| s.geohash.clone()).collect::<Vec<_>>())
    .execute(&mut *tx)
    .await
    .context("upserting stations")?;

    sqlx::query(
        "INSERT INTO water_measurements (timestamp, station_id, water_level, discharge, discharge_percent, status)
         SELECT *, 'active' FROM UNNEST($1::timestamp[], $2::int4[], $3::float8[], $4::float8[], $5::float8[])
         ON CONFLICT (timestamp, station_id) DO UPDATE SET
            water_level = EXCLUDED.water_level, discharge = EXCLUDED.discharge,
            discharge_percent = EXCLUDED.discharge_percent, status = EXCLUDED.status",
    )
    .bind(rows.iter().map(|r| r.timestamp).collect::<Vec<_>>())
    .bind(rows.iter().map(|r| r.station_num as i32).collect::<Vec<_>>())
    .bind(rows.iter().map(|r| r.water_level).collect::<Vec<_>>())
    .bind(rows.iter().map(|r| r.discharge).collect::<Vec<_>>())
    .bind(rows.iter().map(|r| r.discharge_percent).collect::<Vec<_>>())
    .execute(&mut *tx)
    .await
    .context("upserting water_measurements")?;
    tx.commit().await?;
    Ok(rows.len())
}

/// Stored RID levels of report columns `nums` in [start, end], for the
/// screening (rid::screen).
pub async fn rid_readings_between(
    pool: &PgPool,
    nums: &[u32],
    start: NaiveDateTime,
    end: NaiveDateTime,
) -> Result<Vec<crate::rid::StoredReading>> {
    let rows: Vec<(NaiveDateTime, i32, f64, Option<f64>)> = sqlx::query_as(
        "SELECT timestamp, station_id, water_level::float8, discharge::float8
         FROM water_measurements
         WHERE station_id = ANY($1) AND timestamp >= $2 AND timestamp <= $3
           AND water_level IS NOT NULL",
    )
    .bind(nums.iter().map(|&n| n as i32).collect::<Vec<_>>())
    .bind(start)
    .bind(end)
    .fetch_all(pool)
    .await
    .context("reading stored RID levels")?;
    Ok(rows
        .into_iter()
        .map(|(t, n, l, q)| (t, n as u32, l, q))
        .collect())
}

/// Delete stored RID readings (timestamp, report column).
pub async fn delete_rid_readings(pool: &PgPool, keys: &[(NaiveDateTime, u32)]) -> Result<usize> {
    if keys.is_empty() {
        return Ok(0);
    }
    let r = sqlx::query(
        "DELETE FROM water_measurements w
         USING UNNEST($1::timestamp[], $2::int4[]) AS k(ts, sid)
         WHERE w.timestamp = k.ts AND w.station_id = k.sid",
    )
    .bind(keys.iter().map(|(t, _)| *t).collect::<Vec<_>>())
    .bind(keys.iter().map(|(_, n)| *n as i32).collect::<Vec<_>>())
    .execute(pool)
    .await
    .context("deleting RID spikes")?;
    Ok(r.rows_affected() as usize)
}

/// Delete the model rows of forecast_history issued at or after `from`, so
/// a catch-up recomputes them (after the readings they were computed from
/// changed).
pub async fn delete_model_forecasts_since(pool: &PgPool, from: NaiveDateTime) -> Result<usize> {
    let r = sqlx::query("DELETE FROM forecast_history WHERE as_of >= $1 AND source = 'model'")
        .bind(from)
        .execute(pool)
        .await
        .context("deleting forecast_history rows")?;
    Ok(r.rows_affected() as usize)
}

/// Day -> hours (0-23) that have at least one RID reading, for [start, end].
pub async fn recorded_hours_by_day(
    pool: &PgPool,
    start: NaiveDate,
    end: NaiveDate,
) -> Result<HashMap<NaiveDate, HashSet<u32>>> {
    let rows: Vec<(NaiveDate, i32)> = sqlx::query_as(
        "SELECT CAST(timestamp AS DATE) AS day, CAST(EXTRACT(HOUR FROM timestamp) AS INTEGER) AS hour
         FROM water_measurements WHERE timestamp >= $1 AND timestamp < $2
         GROUP BY 1, 2",
    )
    .bind(start.and_time(NaiveTime::MIN))
    .bind((end + chrono::Duration::days(1)).and_time(NaiveTime::MIN))
    .fetch_all(pool)
    .await?;
    let mut out: HashMap<NaiveDate, HashSet<u32>> = HashMap::new();
    for (day, hour) in rows {
        out.entry(day).or_default().insert(hour as u32);
    }
    Ok(out)
}

pub async fn earliest_rid_timestamp(pool: &PgPool) -> Result<Option<NaiveDateTime>> {
    Ok(
        sqlx::query_scalar("SELECT MIN(timestamp) FROM water_measurements")
            .fetch_one(pool)
            .await?,
    )
}

/// (hour, station) pairs in `[start, end)` where one of `codes` has a RID
/// level but no model forecast: hours that arrived together with a later
/// one (RID's overnight batch, a late report), so no run forecast from them.
pub async fn hours_without_model_forecast(
    pool: &PgPool,
    codes: &[String],
    start: NaiveDateTime,
    end: NaiveDateTime,
) -> Result<Vec<(NaiveDateTime, String)>> {
    sqlx::query_as(
        "SELECT DISTINCT date_trunc('hour', m.timestamp) AS h, s.station_code::text
         FROM water_measurements m JOIN stations s ON s.id = m.station_id
         WHERE m.timestamp >= $1 AND m.timestamp < $2 AND m.water_level IS NOT NULL
           AND s.station_code = ANY($3)
           AND NOT EXISTS (
               SELECT 1 FROM forecast_history f
                WHERE f.station_code = s.station_code
                  AND f.as_of = date_trunc('hour', m.timestamp) AND f.source = 'model')
         ORDER BY 1, 2",
    )
    .bind(start)
    .bind(end)
    .bind(codes)
    .fetch_all(pool)
    .await
    .context("finding hours without a model forecast")
}

pub async fn latest_rid_timestamp(pool: &PgPool) -> Result<Option<NaiveDateTime>> {
    Ok(
        sqlx::query_scalar("SELECT MAX(timestamp) FROM water_measurements")
            .fetch_one(pool)
            .await?,
    )
}

/// How many of `readings` are already stored (dry-run comparison).
pub async fn count_stored_rid(pool: &PgPool, readings: &[Reading]) -> Result<i64> {
    Ok(sqlx::query_scalar(
        "SELECT COUNT(*) FROM water_measurements w
         JOIN UNNEST($1::timestamp[], $2::int4[]) AS k(ts, sid) ON w.timestamp = k.ts AND w.station_id = k.sid",
    )
    .bind(readings.iter().map(|r| r.timestamp).collect::<Vec<_>>())
    .bind(readings.iter().map(|r| r.station_num as i32).collect::<Vec<_>>())
    .fetch_one(pool)
    .await?)
}

pub async fn save_hii_rain(
    pool: &PgPool,
    records: &[RainRecord],
    now: NaiveDateTime,
) -> Result<usize> {
    if records.is_empty() {
        return Ok(0);
    }
    let st = dedupe_last(records, |r| r.station_id);
    let ms = dedupe_last(records, |r| (r.station_id, r.timestamp));
    let mut tx = pool.begin().await?;
    sqlx::query(
        "INSERT INTO hii_rain_stations (id, oldcode, name_th, name_en, latitude, longitude, sub_basin_id, agency, updated_at)
         SELECT * FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::float8[], $6::float8[], $7::text[], $8::text[], $9::timestamp[])
         ON CONFLICT (id) DO UPDATE SET
            oldcode = EXCLUDED.oldcode, name_th = EXCLUDED.name_th, name_en = EXCLUDED.name_en,
            latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
            sub_basin_id = EXCLUDED.sub_basin_id, agency = EXCLUDED.agency, updated_at = EXCLUDED.updated_at",
    )
    .bind(st.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.oldcode.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.name_en.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.latitude).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.longitude).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.sub_basin_id.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.agency.clone()).collect::<Vec<_>>())
    .bind(vec![now; st.len()])
    .execute(&mut *tx)
    .await
    .context("upserting hii_rain_stations")?;
    sqlx::query(
        "INSERT INTO hii_rainfall (station_id, timestamp, rain_1h, rain_24h)
         SELECT * FROM UNNEST($1::int4[], $2::timestamp[], $3::float8[], $4::float8[])
         ON CONFLICT (station_id, timestamp) DO UPDATE SET
            rain_1h = EXCLUDED.rain_1h, rain_24h = EXCLUDED.rain_24h",
    )
    .bind(ms.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.timestamp).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.rain_1h).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.rain_24h).collect::<Vec<_>>())
    .execute(&mut *tx)
    .await
    .context("upserting hii_rainfall")?;
    tx.commit().await?;
    Ok(ms.len())
}

pub async fn save_hii_waterlevel(
    pool: &PgPool,
    records: &[WaterLevelRecord],
    now: NaiveDateTime,
) -> Result<usize> {
    if records.is_empty() {
        return Ok(0);
    }
    let st = dedupe_last(records, |r| r.station_id);
    let ms = dedupe_last(records, |r| (r.station_id, r.timestamp));
    let mut tx = pool.begin().await?;
    sqlx::query(
        "INSERT INTO hii_wl_stations (id, oldcode, rid_code, name_th, name_en, latitude, longitude, agency,
            river_name, offset_msl, ground_level_msl, min_bank_msl, critical_level_msl, critical_level_m,
            qmax, is_key_station, updated_at)
         SELECT * FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::text[], $6::float8[],
            $7::float8[], $8::text[], $9::text[], $10::float8[], $11::float8[], $12::float8[], $13::float8[],
            $14::float8[], $15::float8[], $16::bool[], $17::timestamp[])
         ON CONFLICT (id) DO UPDATE SET
            oldcode = EXCLUDED.oldcode, rid_code = EXCLUDED.rid_code, name_th = EXCLUDED.name_th,
            name_en = EXCLUDED.name_en, latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
            agency = EXCLUDED.agency, river_name = EXCLUDED.river_name, offset_msl = EXCLUDED.offset_msl,
            ground_level_msl = EXCLUDED.ground_level_msl, min_bank_msl = EXCLUDED.min_bank_msl,
            critical_level_msl = EXCLUDED.critical_level_msl, critical_level_m = EXCLUDED.critical_level_m,
            qmax = EXCLUDED.qmax, is_key_station = EXCLUDED.is_key_station, updated_at = EXCLUDED.updated_at",
    )
    .bind(st.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.oldcode.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.rid_code.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.name_en.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.latitude).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.longitude).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.agency.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.river_name.clone()).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.offset_msl).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.ground_level_msl).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.min_bank_msl).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.critical_level_msl).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.critical_level_m).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.qmax).collect::<Vec<_>>())
    .bind(st.iter().map(|r| r.is_key_station).collect::<Vec<_>>())
    .bind(vec![now; st.len()])
    .execute(&mut *tx)
    .await
    .context("upserting hii_wl_stations")?;
    sqlx::query(
        "INSERT INTO hii_waterlevel (station_id, timestamp, wl_msl, wl_m, discharge, flow_rate,
            storage_percent, situation_level, diff_wl_bank)
         SELECT * FROM UNNEST($1::int4[], $2::timestamp[], $3::float8[], $4::float8[], $5::float8[],
            $6::float8[], $7::float8[], $8::int4[], $9::float8[])
         ON CONFLICT (station_id, timestamp) DO UPDATE SET
            wl_msl = EXCLUDED.wl_msl, wl_m = EXCLUDED.wl_m, discharge = EXCLUDED.discharge,
            flow_rate = EXCLUDED.flow_rate, storage_percent = EXCLUDED.storage_percent,
            situation_level = EXCLUDED.situation_level, diff_wl_bank = EXCLUDED.diff_wl_bank",
    )
    .bind(ms.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.timestamp).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.wl_msl).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.wl_m).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.discharge).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.flow_rate).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.storage_percent).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.situation_level).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.diff_wl_bank).collect::<Vec<_>>())
    .execute(&mut *tx)
    .await
    .context("upserting hii_waterlevel")?;
    tx.commit().await?;
    Ok(ms.len())
}

/// rid_dams metadata never blanks a stored value (COALESCE), and the daily
/// `level_msl` is preserved too: only the range backfill carries a level, the
/// daily snapshot would NULL it on every cycle.
pub async fn save_dams(pool: &PgPool, records: &[DamRecord], now: NaiveDateTime) -> Result<usize> {
    if records.is_empty() {
        return Ok(0);
    }
    let dams = dedupe_last(records, |r| r.dam_id.clone());
    let ms = dedupe_last(records, |r| (r.dam_id.clone(), r.date));
    let mut tx = pool.begin().await?;
    sqlx::query(
        "INSERT INTO rid_dams (dam_id, region, name_th, latitude, longitude, capacity_max_mcm,
            capacity_normal_mcm, updated_at)
         SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::float8[],
            $6::float8[], $7::float8[], $8::timestamp[])
         ON CONFLICT (dam_id) DO UPDATE SET
            region = COALESCE(EXCLUDED.region, rid_dams.region),
            name_th = COALESCE(EXCLUDED.name_th, rid_dams.name_th),
            latitude = COALESCE(EXCLUDED.latitude, rid_dams.latitude),
            longitude = COALESCE(EXCLUDED.longitude, rid_dams.longitude),
            capacity_max_mcm = COALESCE(EXCLUDED.capacity_max_mcm, rid_dams.capacity_max_mcm),
            capacity_normal_mcm = COALESCE(EXCLUDED.capacity_normal_mcm, rid_dams.capacity_normal_mcm),
            updated_at = EXCLUDED.updated_at",
    )
    .bind(dams.iter().map(|r| r.dam_id.clone()).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.region.clone()).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.latitude).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.longitude).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.capacity_max_mcm).collect::<Vec<_>>())
    .bind(dams.iter().map(|r| r.capacity_normal_mcm).collect::<Vec<_>>())
    .bind(vec![now; dams.len()])
    .execute(&mut *tx)
    .await
    .context("upserting rid_dams")?;
    sqlx::query(
        "INSERT INTO rid_reservoir_daily (dam_id, date, storage_mcm, storage_pct, inflow_mcm, outflow_mcm, level_msl)
         SELECT * FROM UNNEST($1::text[], $2::date[], $3::float8[], $4::float8[], $5::float8[], $6::float8[], $7::float8[])
         ON CONFLICT (dam_id, date) DO UPDATE SET
            storage_mcm = EXCLUDED.storage_mcm, storage_pct = EXCLUDED.storage_pct,
            inflow_mcm = EXCLUDED.inflow_mcm, outflow_mcm = EXCLUDED.outflow_mcm,
            level_msl = COALESCE(EXCLUDED.level_msl, rid_reservoir_daily.level_msl)",
    )
    .bind(ms.iter().map(|r| r.dam_id.clone()).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.date).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.storage_mcm).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.storage_pct).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.inflow_mcm).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.outflow_mcm).collect::<Vec<_>>())
    .bind(ms.iter().map(|r| r.level_msl).collect::<Vec<_>>())
    .execute(&mut *tx)
    .await
    .context("upserting rid_reservoir_daily")?;
    tx.commit().await?;
    Ok(ms.len())
}

pub async fn save_rain(pool: &PgPool, rows: &[RainRow]) -> Result<usize> {
    if rows.is_empty() {
        return Ok(0);
    }
    let rows = dedupe_last(rows, |r| r.timestamp);
    let col = |i: usize| rows.iter().map(|r| r.points[i]).collect::<Vec<_>>();
    sqlx::query(
        "INSERT INTO openmeteo_rain (timestamp, chiang_dao, mae_taeng, mae_ngat, mae_rim, chiang_mai, catchment_mean)
         SELECT * FROM UNNEST($1::timestamp[], $2::float8[], $3::float8[], $4::float8[], $5::float8[], $6::float8[], $7::float8[])
         ON CONFLICT (timestamp) DO UPDATE SET
            chiang_dao = EXCLUDED.chiang_dao, mae_taeng = EXCLUDED.mae_taeng, mae_ngat = EXCLUDED.mae_ngat,
            mae_rim = EXCLUDED.mae_rim, chiang_mai = EXCLUDED.chiang_mai, catchment_mean = EXCLUDED.catchment_mean",
    )
    .bind(rows.iter().map(|r| r.timestamp).collect::<Vec<_>>())
    .bind(col(0))
    .bind(col(1))
    .bind(col(2))
    .bind(col(3))
    .bind(col(4))
    .bind(rows.iter().map(|r| r.catchment_mean).collect::<Vec<_>>())
    .execute(pool)
    .await
    .context("upserting openmeteo_rain")?;
    Ok(rows.len())
}

/// Row counts per table, for `check-db`.
pub async fn table_counts(pool: &PgPool) -> Result<BTreeMap<&'static str, i64>> {
    let mut out = BTreeMap::new();
    for table in [
        "water_measurements",
        "hii_rainfall",
        "hii_waterlevel",
        "rid_reservoir_daily",
        "openmeteo_rain",
    ] {
        let n: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table}"))
            .fetch_one(pool)
            .await?;
        out.insert(table, n);
    }
    Ok(out)
}

// ---------------------------------------------------------------------------
// Forecast inputs and outputs (FORECAST_IN_COLLECTOR)
//
// NUMERIC columns are read as text and parsed: Python's `float(Decimal)` is a
// correctly rounded parse of the same digits, so both sides get the same f64.
// Forecast values are written as text (shortest round-trip digits, as Python's
// repr) so PostgreSQL rounds the same decimal as it does for psycopg2.
// ---------------------------------------------------------------------------

/// A row of (key, key, NUMERIC as text, NUMERIC as text).
type TextRow<A, B> = (A, B, Option<String>, Option<String>);
/// (date, storage_pct, inflow_mcm, outflow_mcm) with NUMERIC as text.
type DamRow = (NaiveDate, Option<String>, Option<String>, Option<String>);

fn parse_num(v: Option<String>) -> f64 {
    v.and_then(|s| s.parse().ok()).unwrap_or(f64::NAN)
}

fn num_text(v: f64) -> Option<String> {
    v.is_finite().then(|| format!("{v}"))
}

/// RID readings in [start, end], as `data._fetch_from_db`.
pub async fn load_rid_window(
    pool: &PgPool,
    start: NaiveDateTime,
    end: NaiveDateTime,
) -> Result<Vec<ping_forecast::Row>> {
    let rows: Vec<TextRow<NaiveDateTime, String>> = sqlx::query_as(
        "SELECT m.timestamp, s.station_code, m.water_level::text, m.discharge::text
         FROM water_measurements m JOIN stations s ON m.station_id = s.id
         WHERE m.timestamp >= $1 AND m.timestamp <= $2 ORDER BY m.timestamp",
    )
    .bind(start)
    .bind(end)
    .fetch_all(pool)
    .await
    .context("loading the RID window")?;
    Ok(rows
        .into_iter()
        .map(|(t, code, level, q)| ping_forecast::Row {
            timestamp: ping_forecast::grid::floor_hour(t),
            station_code: code,
            water_level: parse_num(level),
            discharge: parse_num(q),
        })
        .collect())
}

/// HII mirror readings for `codes` in [start, end], as `data._fetch_hii_levels`.
pub async fn load_hii_window(
    pool: &PgPool,
    codes: &[String],
    start: NaiveDateTime,
    end: NaiveDateTime,
) -> Result<Vec<ping_forecast::hii::HiiRow>> {
    let rows: Vec<TextRow<NaiveDateTime, String>> = sqlx::query_as(
        "SELECT m.timestamp, s.rid_code, m.wl_msl::text, m.discharge::text
         FROM hii_waterlevel m JOIN hii_wl_stations s ON s.id = m.station_id
         WHERE s.rid_code IS NOT NULL AND m.timestamp >= $1 AND m.timestamp <= $2
           AND s.rid_code = ANY($3)
         ORDER BY m.timestamp, m.station_id",
    )
    .bind(start)
    .bind(end)
    .bind(codes)
    .fetch_all(pool)
    .await
    .context("loading HII mirror levels")?;
    Ok(rows
        .into_iter()
        .map(|(t, code, wl, q)| ping_forecast::hii::HiiRow {
            timestamp: t,
            station_code: code,
            wl_msl: parse_num(wl),
            discharge: parse_num(q),
        })
        .collect())
}

/// Newest HII mirror reading for `codes` (a forecast trigger).
pub async fn latest_hii_timestamp(
    pool: &PgPool,
    codes: &[String],
) -> Result<Option<NaiveDateTime>> {
    Ok(sqlx::query_scalar(
        "SELECT MAX(m.timestamp) FROM hii_waterlevel m JOIN hii_wl_stations s ON s.id = m.station_id
         WHERE s.rid_code = ANY($1) AND m.wl_msl IS NOT NULL",
    )
    .bind(codes)
    .fetch_one(pool)
    .await?)
}

/// Daily reservoir rows from `start`, as `dam.load_daily`.
pub async fn load_dam_days(
    pool: &PgPool,
    dam_id: &str,
    start: NaiveDate,
) -> Result<Vec<ping_forecast::dam::DamDay>> {
    let rows: Vec<DamRow> = sqlx::query_as(
        "SELECT date, storage_pct::text, inflow_mcm::text, outflow_mcm::text
         FROM rid_reservoir_daily WHERE dam_id = $1 AND date >= $2 ORDER BY date",
    )
    .bind(dam_id)
    .bind(start)
    .fetch_all(pool)
    .await
    .context("loading the dam series")?;
    Ok(rows
        .into_iter()
        .map(|(date, s, i, o)| ping_forecast::dam::DamDay {
            date,
            storage_pct: parse_num(s),
            inflow_mcm: parse_num(i),
            outflow_mcm: parse_num(o),
        })
        .collect())
}

/// The newest RID row per station (`get_latest_measurements(200)`).
#[derive(Clone, Debug, PartialEq)]
pub struct LatestReading {
    pub timestamp: NaiveDateTime,
    pub station_code: String,
    pub water_level: Option<f64>,
    pub discharge_percent: Option<f64>,
}

pub async fn latest_rid_readings(pool: &PgPool) -> Result<Vec<LatestReading>> {
    let rows: Vec<TextRow<NaiveDateTime, String>> = sqlx::query_as(
        "SELECT m.timestamp, s.station_code, m.water_level::text, m.discharge_percent::text
         FROM water_measurements m
         JOIN stations s ON m.station_id = s.id
         INNER JOIN (
             SELECT station_id, MAX(timestamp) AS max_timestamp
             FROM water_measurements GROUP BY station_id
         ) latest ON m.station_id = latest.station_id AND m.timestamp = latest.max_timestamp
         ORDER BY s.station_code
         LIMIT 200",
    )
    .fetch_all(pool)
    .await
    .context("loading the latest readings")?;
    let num = |v: Option<String>| v.and_then(|s| s.parse().ok());
    Ok(rows
        .into_iter()
        .map(|(t, code, level, pct)| LatestReading {
            timestamp: t,
            station_code: code,
            water_level: num(level),
            discharge_percent: num(pct),
        })
        .collect())
}

/// `ForecastHistoryStore.save_rows`: one row per (as_of, station, horizon).
/// Store issued forecasts; a later run's row for the same hour replaces the
/// earlier one, except that a heuristic row never replaces a model row. The
/// hourly run re-issues the newest reading's hour while RID is silent
/// overnight, and after 6 h (`stale_after_h`) it does so as the heuristic:
/// without the exception the 21:00 model forecast was overwritten each night.
pub async fn save_forecast_history(
    pool: &PgPool,
    rows: &[ping_forecast::Forecast],
) -> Result<usize> {
    let rows = dedupe_last(rows, |r| (r.as_of, r.station_code.clone(), r.horizon_hours));
    if rows.is_empty() {
        return Ok(0);
    }
    sqlx::query(
        "INSERT INTO forecast_history (as_of, station_code, horizon_hours, predicted_max_level,
            p_warning, p_danger, current_level, model_version, source)
         SELECT a, c, h, pm::numeric, pw::numeric, pd::numeric, cl::numeric, mv, src
         FROM UNNEST($1::timestamp[], $2::text[], $3::int4[], $4::text[], $5::text[], $6::text[],
                     $7::text[], $8::text[], $9::text[])
              AS t(a, c, h, pm, pw, pd, cl, mv, src)
         ON CONFLICT (as_of, station_code, horizon_hours) DO UPDATE SET
            predicted_max_level = EXCLUDED.predicted_max_level, p_warning = EXCLUDED.p_warning,
            p_danger = EXCLUDED.p_danger, current_level = EXCLUDED.current_level,
            model_version = EXCLUDED.model_version, source = EXCLUDED.source
         WHERE forecast_history.source IS DISTINCT FROM 'model' OR EXCLUDED.source = 'model'",
    )
    .bind(rows.iter().map(|r| r.as_of).collect::<Vec<_>>())
    .bind(
        rows.iter()
            .map(|r| r.station_code.clone())
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| r.horizon_hours as i32)
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| num_text(r.predicted_max_level))
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| num_text(r.p_warning))
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| num_text(r.p_danger))
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| num_text(r.current_level))
            .collect::<Vec<_>>(),
    )
    .bind(
        rows.iter()
            .map(|r| r.model_version.clone())
            .collect::<Vec<_>>(),
    )
    .bind(rows.iter().map(|r| r.source.clone()).collect::<Vec<_>>())
    .execute(pool)
    .await
    .context("upserting forecast_history")?;
    Ok(rows.len())
}

/// The payload /forecast serves (see `forecast_current` above).
pub async fn save_forecast_current(
    pool: &PgPool,
    computed_at: NaiveDateTime,
    rows: &[ping_forecast::Forecast],
) -> Result<()> {
    let versions: std::collections::BTreeSet<&str> =
        rows.iter().map(|r| r.model_version.as_str()).collect();
    sqlx::query(
        "INSERT INTO forecast_current (id, computed_at, model_versions, payload)
         VALUES ('all', $1, $2, $3)
         ON CONFLICT (id) DO UPDATE SET computed_at = EXCLUDED.computed_at,
            model_versions = EXCLUDED.model_versions, payload = EXCLUDED.payload",
    )
    .bind(computed_at)
    .bind(versions.into_iter().collect::<Vec<_>>().join(","))
    .bind(serde_json::to_string(rows)?)
    .execute(pool)
    .await
    .context("writing forecast_current")?;
    Ok(())
}

/// notification_state as a key -> state map.
pub async fn notification_states(pool: &PgPool) -> Result<HashMap<String, String>> {
    let rows: Vec<(String, String)> = sqlx::query_as("SELECT key, state FROM notification_state")
        .fetch_all(pool)
        .await
        .context("reading notification_state")?;
    Ok(rows.into_iter().collect())
}

pub async fn set_notification_state(
    pool: &PgPool,
    key: &str,
    state: &str,
    value: Option<f64>,
    now: NaiveDateTime,
) -> Result<()> {
    sqlx::query(
        "INSERT INTO notification_state (key, state, value, updated_at)
         VALUES ($1, $2, $3::numeric, $4) ON CONFLICT (key) DO UPDATE SET
            state = EXCLUDED.state, value = EXCLUDED.value, updated_at = EXCLUDED.updated_at",
    )
    .bind(key)
    .bind(state)
    .bind(value.and_then(num_text))
    .bind(now)
    .execute(pool)
    .await
    .with_context(|| format!("saving notification state {key}"))?;
    Ok(())
}