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::path::Path;

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;

#[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, _tmp_dir),
    fields(
        valid_time = %valid_time,
        points = points.len(),
    )
)]
pub async fn process_batch(
    client: &HrrrClient,
    pool: &PgPool,
    valid_time: DateTime<Utc>,
    points: &[(f64, f64)],
    _tmp_dir: &Path,
) -> 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_pattern = format!(
        ":({}):",
        sfc_wanted
            .iter()
            .map(|(v, _)| v.as_str())
            .collect::<Vec<_>>()
            .join("|")
    );
    let prs_pattern = format!(
        ":({}):",
        prs_wanted
            .iter()
            .map(|(v, _)| v.as_str())
            .collect::<Vec<_>>()
            .join("|")
    );

    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_pattern, grid_spec)
    })
    .await
    .expect("blocking join")?;

    let prs_grid = tokio::task::spawn_blocking(move || {
        decoder::extract_grid(&prs_blob, &prs_pattern, 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. Skip the point silently if extract_grid
    // produced nothing for its snapped key (e.g. off-CONUS).
    let mut inserted = 0u32;
    for &(lat, lon) in points {
        let key = ((lat * 1000.0).round() as i32, (lon * 1000.0).round() as i32);
        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_PRESSURE_LEVELS
        .iter()
        .filter_map(|&p| {
            let pres_mb = p as f64;
            let t = cell.get(&format!("TMP:{p} mb")).copied()?;
            let h = cell.get(&format!("HGT:{p} mb")).copied()?;
            let d = cell.get(&format!("DPT:{p} mb")).copied();
            let mut m = serde_json::Map::new();
            m.insert("pres_mb".into(), serde_json::json!(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(())
}

// Integration tests for process_batch + upsert_profile live in
// tests/hrrr_points_live.rs (network + DB gated). No unit tests here
// — the two public functions are thin glue over fetcher/decoder which
// carry their own test coverage.