prop/rust/prop_grid_rs/src/hrrr_points.rs

//! Per-QSO HRRR point extraction pipeline.
//!
//! Phase 3 Stream C Rust side. Reads batches from `hrrr_fetch_tasks`
//! (one row per valid_time with a JSONB array of `{lat, lon}` points),
//! fetches the HRRR cycle once, extracts each requested point, and
//! bulk-inserts into `hrrr_profiles`.
//!
//! The Elixir side already snaps lat/lon to the 3 km HRRR grid and
//! skips points that already have a profile, so duplicates here are
//! rare and handled via `ON CONFLICT DO UPDATE` on (lat, lon,
//! valid_time).

use std::sync::OnceLock;

use chrono::{DateTime, Timelike, Utc};
use sqlx::PgPool;
use uuid::Uuid;

use crate::decoder::{self, CellValues};
use crate::fetcher::{self, HrrrClient, Product};
use crate::grid::{wgrib2_grid_spec, STEP};

// wgrib2 filter patterns are derived from static SURFACE_MESSAGES /
// GRID_PRESSURE_LEVELS tables and never change at runtime. Building
// them per batch was wasted allocation on every hrrr_fetch_tasks tick.
fn sfc_pattern() -> &'static str {
    static P: OnceLock<String> = OnceLock::new();
    P.get_or_init(|| build_pattern(&fetcher::surface_messages_owned()))
}

fn prs_pattern() -> &'static str {
    static P: OnceLock<String> = OnceLock::new();
    P.get_or_init(|| build_pattern(&fetcher::pressure_messages_grid()))
}

fn build_pattern(msgs: &[(String, String)]) -> String {
    let vars: Vec<&str> = msgs.iter().map(|(v, _)| v.as_str()).collect();
    format!(":({}):", vars.join("|"))
}

#[derive(Debug, thiserror::Error)]
pub enum PointFetchError {
    #[error("fetch: {0}")]
    Fetch(#[from] fetcher::FetchError),
    #[error("decode: {0}")]
    Decode(#[from] decoder::DecodeError),
    #[error("db: {0}")]
    Db(#[from] sqlx::Error),
}

#[derive(Debug, Clone)]
pub struct PointFetchStats {
    pub points_requested: u32,
    pub profiles_inserted: u32,
}

/// Fetch + decode + persist for one `hrrr_fetch_tasks` batch.
#[tracing::instrument(
    name = "hrrr_points.process_batch",
    skip(client, pool),
    fields(
        valid_time = %valid_time,
        points = points.len(),
    )
)]
pub async fn process_batch(
    client: &HrrrClient,
    pool: &PgPool,
    valid_time: DateTime<Utc>,
    points: &[(f64, f64)],
) -> Result<PointFetchStats, PointFetchError> {
    let date = valid_time.date_naive();
    let hour = valid_time.hour() as u8;
    let grid_spec = wgrib2_grid_spec();

    let sfc_wanted = fetcher::surface_messages_owned();
    let prs_wanted = fetcher::pressure_messages_grid();
    let sfc_pat = sfc_pattern();
    let prs_pat = prs_pattern();

    let sfc_fut = client.fetch_product_blob(date, hour, Product::Surface, 0, &sfc_wanted);
    let prs_fut = client.fetch_product_blob(date, hour, Product::Pressure, 0, &prs_wanted);
    let (sfc_blob, prs_blob) = tokio::try_join!(sfc_fut, prs_fut)?;

    let sfc_grid =
        tokio::task::spawn_blocking(move || decoder::extract_grid(&sfc_blob, sfc_pat, grid_spec))
            .await
            .expect("blocking join")?;

    let prs_grid =
        tokio::task::spawn_blocking(move || decoder::extract_grid(&prs_blob, prs_pat, grid_spec))
            .await
            .expect("blocking join")?;

    // Empty grids almost always mean the HRRR archive doesn't have this
    // cycle (common for backfills older than the rolling retention
    // window). Log so "0 inserted" doesn't look like a silent success.
    if sfc_grid.is_empty() && prs_grid.is_empty() {
        tracing::warn!(
            valid_time = %valid_time,
            points = points.len(),
            "hrrr fetch returned empty surface and pressure grids — upstream archive likely missing this cycle"
        );
    }

    // For each requested point, pull the merged cell and build an
    // hrrr_profiles row. Elixir snaps points to 0.01° (see
    // `Weather.round_to_hrrr_grid/2`), but wgrib2 resamples to the 0.125°
    // CONUS prop grid — so we must re-snap before keying into `sfc_grid`
    // or the lookup misses and the task completes with zero inserts.
    // Store the profile at the original requested (lat, lon) so
    // `Weather.has_hrrr_profile?/3` (0.07° window) finds it on the
    // next backfill tick and flips the contact to `:complete`.
    let mut inserted = 0u32;
    for &(lat, lon) in points {
        let key = snap_key(lat, lon);
        let mut cell: CellValues = sfc_grid.get(&key).cloned().unwrap_or_default();
        if let Some(prs) = prs_grid.get(&key) {
            cell.extend(prs.iter().map(|(k, v)| (k.clone(), *v)));
        }
        if cell.is_empty() {
            continue;
        }

        upsert_profile(pool, lat, lon, valid_time, &cell).await?;
        inserted += 1;
    }

    Ok(PointFetchStats {
        points_requested: points.len() as u32,
        profiles_inserted: inserted,
    })
}

async fn upsert_profile(
    pool: &PgPool,
    lat: f64,
    lon: f64,
    valid_time: DateTime<Utc>,
    cell: &CellValues,
) -> Result<(), sqlx::Error> {
    let surface_temp_c = cell
        .get("TMP:2 m above ground")
        .copied()
        .map(|v| (v as f64) - 273.15);
    let surface_dewpoint_c = cell
        .get("DPT:2 m above ground")
        .copied()
        .map(|v| (v as f64) - 273.15);
    let surface_pressure_mb = cell
        .get("PRES:surface")
        .copied()
        .map(|v| (v as f64) / 100.0);
    let hpbl_m = cell.get("HPBL:surface").copied().map(|v| v as f64);
    let pwat_mm = cell
        .get("PWAT:entire atmosphere (considered as a single layer)")
        .copied()
        .map(|v| v as f64);

    // Pressure-level profile list. Postgres column `profile` is jsonb[]
    // (ARRAY of jsonb), one element per level, matching Elixir's
    // HrrrClient.build_profile output shape:
    //   [{"pres_mb", "hght_m", "tmpc", "dwpc"}, …]
    // We build a Vec<Json<Value>> so sqlx encodes it as jsonb[] rather
    // than wrapping the whole list in a single jsonb (which yields the
    // "expression is of type jsonb" type-mismatch error at runtime).
    let levels: Vec<sqlx::types::Json<serde_json::Value>> = fetcher::grid_level_keys()
        .iter()
        .filter_map(|k| {
            let t = cell.get(k.tmp.as_str()).copied()?;
            let h = cell.get(k.hgt.as_str()).copied()?;
            let d = cell.get(k.dpt.as_str()).copied();
            let mut m = serde_json::Map::new();
            m.insert("pres_mb".into(), serde_json::json!(k.pres_mb));
            m.insert("hght_m".into(), serde_json::json!(h as f64));
            m.insert("tmpc".into(), serde_json::json!((t as f64) - 273.15));
            if let Some(dv) = d {
                m.insert("dwpc".into(), serde_json::json!((dv as f64) - 273.15));
            }
            Some(sqlx::types::Json(serde_json::Value::Object(m)))
        })
        .collect();

    let id = Uuid::new_v4();
    sqlx::query(
        r#"
        INSERT INTO hrrr_profiles
          (id, valid_time, lat, lon, run_time, profile, hpbl_m, pwat_mm,
           surface_temp_c, surface_dewpoint_c, surface_pressure_mb,
           is_grid_point, inserted_at, updated_at)
        VALUES
          ($1, $2, $3, $4, $2, $5, $6, $7, $8, $9, $10, false, NOW(), NOW())
        ON CONFLICT (lat, lon, valid_time) DO UPDATE
        SET profile = EXCLUDED.profile,
            hpbl_m = EXCLUDED.hpbl_m,
            pwat_mm = EXCLUDED.pwat_mm,
            surface_temp_c = EXCLUDED.surface_temp_c,
            surface_dewpoint_c = EXCLUDED.surface_dewpoint_c,
            surface_pressure_mb = EXCLUDED.surface_pressure_mb,
            updated_at = NOW()
        "#,
    )
    .bind(id)
    .bind(valid_time.naive_utc())
    .bind(lat)
    .bind(lon)
    .bind(&levels)
    .bind(hpbl_m)
    .bind(pwat_mm)
    .bind(surface_temp_c)
    .bind(surface_dewpoint_c)
    .bind(surface_pressure_mb)
    .execute(pool)
    .await?;
    Ok(())
}

/// Snap a requested (lat, lon) to the nearest 0.125° grid cell and
/// return its integer-millidegree key. The wgrib2 subprocess resamples
/// HRRR to that grid (see `wgrib2_grid_spec()`), so a request on the
/// HRRR native ~3 km grid (e.g. 44.94, -91.63 after
/// `Weather.round_to_hrrr_grid/2`, which rounds to 0.01°) never matches
/// a 0.125° cell key exactly and the point is silently dropped —
/// producing the "task done with 0 profiles inserted" pattern that
/// wedged 13k+ contacts at `hrrr_status = :queued`.
fn snap_key(lat: f64, lon: f64) -> (i32, i32) {
    let snap_lat = (lat / STEP).round() * STEP;
    let snap_lon = (lon / STEP).round() * STEP;
    (
        (snap_lat * 1000.0).round() as i32,
        (snap_lon * 1000.0).round() as i32,
    )
}

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

    #[test]
    fn snap_key_on_grid_is_identity() {
        assert_eq!(snap_key(45.0, -91.625), (45_000, -91_625));
        assert_eq!(snap_key(25.0, -125.0), (25_000, -125_000));
    }

    #[test]
    fn snap_key_rounds_to_nearest_0_125_cell() {
        // 44.94 → 45.0 (0.06 away vs 44.875 which is 0.065 away)
        assert_eq!(snap_key(44.94, -91.63), (45_000, -91_625));
        // 44.83 → 44.875 (0.045 vs 44.75 which is 0.08)
        assert_eq!(snap_key(44.83, -91.5), (44_875, -91_500));
        // 44.73 → 44.75 (0.02 vs 44.625 which is 0.105)
        assert_eq!(snap_key(44.73, -91.38), (44_750, -91_375));
    }
}

// Integration tests for process_batch + upsert_profile live in
// tests/hrrr_points_live.rs (network + DB gated).