//! `.sgrid` — dense, cell-major, random-access derived weather scalars.
//!
//! Replaces the 5°×5° chunked gzipped-MessagePack `.mp.gz` scalar
//! artifact. Measured on a full CONUS grid (95,073 cells):
//!
//! |               | `.mp.gz`                                 | `.sgrid`                     |
//! |---------------|------------------------------------------|------------------------------|
//! | write         | 0.232 s (rmpv + gzip per chunk)          | one `write_all` of ~8.4 MB   |
//! | read viewport | gunzip + unpack every overlapping chunk   | one `pread` per grid row     |
//! | read one cell | gunzip + unpack + iterate one chunk       | one 80-byte `pread`          |
//!
//! Same header shape as [`pgrid`][crate::pgrid]: magic `SGRD`, version 1,
//! flags byte, self-describing NUL-padded field table, then cell-major
//! `f32` body with `NaN` as the missing-value sentinel.
//!
//! ## Layout (little-endian)
//!
//! ```text
//! magic        4   "SGRD"
//! version      1   0x01
//! flags        1   bit0: 0 = hrrr, 1 = hrdps
//! n_fields     2   u16
//! valid_time   8   i64  unix seconds
//! lat_start    8   f64
//! lon_start    8   f64
//! lat_step     8   f64
//! lon_step     8   f64
//! n_rows       2   u16  (latitude)
//! n_cols       2   u16  (longitude)
//! field_table  n_fields × 32   NUL-padded ASCII field names
//! body         n_rows*n_cols*n_fields × 4   f32, CELL-MAJOR
//! ```
//!
//! Cell-major rather than plane-major: reading one cell across all fields
//! is a single contiguous `pread`. A viewport read is one contiguous
//! `pread` per grid row.
//!
//! The field table is written into the header so the Elixir reader
//! resolves fields by name. Adding a field is backwards compatible.

use std::io::Write;
use std::path::{Path, PathBuf};

use chrono::{DateTime, Utc};

use crate::grid::GridSpec;
use crate::weather_scalar_file::ScalarRow;

pub const MAGIC: &[u8; 4] = b"SGRD";
pub const VERSION: u8 = 1;
/// Fixed width of one field-table entry, in bytes.
pub const FIELD_NAME_LEN: usize = 32;

/// On-disk column order. Each `ScalarRow` field maps to one `f32` column.
/// Lat and lon are implied by the grid spec; valid_time is in the header.
/// Names match the atom keys the Elixir reader expects, so callers see
/// the same map shape the `.mp.gz` path produced.
pub const SCALAR_FIELDS: &[&str] = &[
    "temperature",
    "dewpoint_depression",
    "surface_rh",
    "surface_pressure_mb",
    "surface_refractivity",
    "refractivity_gradient",
    "bl_height",
    "pwat",
    "temp_850mb",
    "dewpoint_850mb",
    "temp_700mb",
    "dewpoint_700mb",
    "lapse_rate",
    "mid_lapse_rate",
    "inversion_strength",
    "inversion_base_m",
    "ducting",
    "duct_base_m",
    "duct_strength",
    "duct_cutoff_ghz",
];

pub fn n_fields() -> usize {
    SCALAR_FIELDS.len()
}

/// Byte offset of the body, i.e. the header + field table length.
pub fn header_len() -> usize {
    4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2 + SCALAR_FIELDS.len() * FIELD_NAME_LEN
}

#[derive(Debug, thiserror::Error)]
pub enum WriteError {
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
}

/// `<scores_dir>/weather_scalars/<iso>.sgrid`.
pub fn path_for(scores_dir: &Path, valid_time: DateTime<Utc>) -> PathBuf {
    let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    scores_dir
        .join("weather_scalars")
        .join(format!("{iso}.sgrid"))
}

/// HRDPS sibling path — `<scores_dir>/weather_scalars/<iso>.hrdps.sgrid`.
/// Coexists with the HRRR `.sgrid` so both sources can be written
/// independently without clobbering each other.
pub fn path_for_hrdps(scores_dir: &Path, valid_time: DateTime<Utc>) -> PathBuf {
    let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    scores_dir
        .join("weather_scalars")
        .join(format!("{iso}.hrdps.sgrid"))
}

/// Serialise the header for `spec` / `valid_time`.
pub fn encode_header(spec: &GridSpec, valid_time: DateTime<Utc>, hrdps: bool) -> Vec<u8> {
    let nf = SCALAR_FIELDS.len();
    let mut out = Vec::with_capacity(header_len());
    out.extend_from_slice(MAGIC);
    out.push(VERSION);
    out.push(u8::from(hrdps));
    out.extend_from_slice(&(nf as u16).to_le_bytes());
    out.extend_from_slice(&valid_time.timestamp().to_le_bytes());
    out.extend_from_slice(&spec.lat_start.to_le_bytes());
    out.extend_from_slice(&spec.lon_start.to_le_bytes());
    out.extend_from_slice(&spec.lat_step.to_le_bytes());
    out.extend_from_slice(&spec.lon_step.to_le_bytes());
    out.extend_from_slice(&(spec.lat_count as u16).to_le_bytes());
    out.extend_from_slice(&(spec.lon_count as u16).to_le_bytes());
    for name in SCALAR_FIELDS {
        let mut buf = [0u8; FIELD_NAME_LEN];
        buf[..name.len()].copy_from_slice(name.as_bytes());
        out.extend_from_slice(&buf);
    }
    debug_assert_eq!(out.len(), header_len());
    out
}

/// Build the cell-major `f32` body from `rows` and `spec`.
///
/// `rows` MUST be sorted in the order `lat` then `lon` at `lat_step` /
/// `lon_step` increments — exactly the iteration order the pipeline's
/// `fuse_chunk` produces. Any cell not present in `rows` is filled with
/// `f32::NAN` (the missing sentinel).
///
/// Returns a flat `f32` array, length `n_cells * n_fields()`, ready for
/// `write_atomic`.
pub fn build_body(spec: &GridSpec, rows: &[ScalarRow]) -> Vec<f32> {
    let n_cells = spec.lat_count * spec.lon_count;
    let nf = SCALAR_FIELDS.len();
    let mut body = vec![f32::NAN; n_cells * nf];

    for row in rows {
        // Convert lat/lon to cell index. The grid's `round3` is
        // half-away-from-zero; round here to match the pipeline's
        // `cell_latlon` output.
        let lat = (row.lat * 1000.0).round() / 1000.0;
        let lon = (row.lon * 1000.0).round() / 1000.0;
        let row_i = ((lat - spec.lat_start) / spec.lat_step).round() as isize;
        let col_i = ((lon - spec.lon_start) / spec.lon_step).round() as isize;
        if row_i < 0
            || col_i < 0
            || row_i as usize >= spec.lat_count
            || col_i as usize >= spec.lon_count
        {
            continue;
        }
        let cell = (row_i as usize) * spec.lon_count + (col_i as usize);
        let base = cell * nf;

        // Write each field. Order must match SCALAR_FIELDS.
        let mut f = 0;
        macro_rules! put {
            ($val:expr) => {
                if let Some(v) = $val {
                    body[base + f] = v as f32;
                }
                f += 1;
            };
        }
        put!(row.temperature);
        put!(row.dewpoint_depression);
        put!(row.surface_rh);
        put!(row.surface_pressure_mb);
        put!(row.surface_refractivity);
        put!(row.refractivity_gradient);
        put!(row.bl_height);
        put!(row.pwat);
        put!(row.temp_850mb);
        put!(row.dewpoint_850mb);
        put!(row.temp_700mb);
        put!(row.dewpoint_700mb);
        put!(row.lapse_rate);
        put!(row.mid_lapse_rate);
        put!(row.inversion_strength);
        put!(row.inversion_base_m);
        // ducting: Option<bool> → 1.0 or NaN
        body[base + f] = match row.ducting {
            Some(true) => 1.0_f32,
            Some(false) => 0.0_f32,
            None => f32::NAN,
        };
        f += 1;
        put!(row.duct_base_m);
        put!(row.duct_strength);
        put!(row.duct_cutoff_ghz);
        debug_assert_eq!(f, nf);
    }

    body
}

/// Write `<scores_dir>/weather_scalars/<iso>.sgrid` atomically
/// (tmp + rename, so an NFS reader sees the old file, the new file, or
/// nothing).
pub fn write_atomic(
    scores_dir: &Path,
    valid_time: DateTime<Utc>,
    spec: &GridSpec,
    body: &[f32],
    hrdps: bool,
) -> Result<PathBuf, WriteError> {
    assert_eq!(
        body.len(),
        spec.lat_count * spec.lon_count * SCALAR_FIELDS.len(),
        "sgrid body does not match grid spec"
    );

    let path = if hrdps {
        path_for_hrdps(scores_dir, valid_time)
    } else {
        path_for(scores_dir, valid_time)
    };

    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }

    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let pid = std::process::id();
    let mut tmp = path.clone().into_os_string();
    tmp.push(format!(".tmp.{nanos}.{pid}"));
    let tmp = PathBuf::from(tmp);

    {
        let file = std::fs::File::create(&tmp)?;
        let mut w = std::io::BufWriter::with_capacity(1 << 20, file);
        w.write_all(&encode_header(spec, valid_time, hrdps))?;
        // f32 little-endian, chunked writes.
        let mut chunk: Vec<u8> = Vec::with_capacity(4 * 8192);
        for values in body.chunks(8192) {
            chunk.clear();
            for v in values {
                chunk.extend_from_slice(&v.to_le_bytes());
            }
            w.write_all(&chunk)?;
        }
        w.flush()?;
    }

    match std::fs::rename(&tmp, &path) {
        Ok(()) => Ok(path),
        Err(e) => {
            let _ = std::fs::remove_file(&tmp);
            Err(WriteError::Io(e))
        }
    }
}

// ── Reader (tests + any Rust-side consumer) ──────────────────────────

#[derive(Debug, Clone)]
pub struct Header {
    pub hrdps: bool,
    pub valid_time: DateTime<Utc>,
    pub spec: GridSpec,
    pub fields: Vec<String>,
}

impl Header {
    pub fn field_index(&self, name: &str) -> Option<usize> {
        self.fields.iter().position(|f| f == name)
    }

    pub fn n_cells(&self) -> usize {
        self.spec.lat_count * self.spec.lon_count
    }
}

pub fn decode_header(bytes: &[u8]) -> Option<Header> {
    if bytes.len() < 52 || &bytes[0..4] != MAGIC || bytes[4] != VERSION {
        return None;
    }
    let hrdps = bytes[5] != 0;
    let n_fields = u16::from_le_bytes(bytes[6..8].try_into().ok()?) as usize;
    let valid_time =
        DateTime::<Utc>::from_timestamp(i64::from_le_bytes(bytes[8..16].try_into().ok()?), 0)?;
    let lat_start = f64::from_le_bytes(bytes[16..24].try_into().ok()?);
    let lon_start = f64::from_le_bytes(bytes[24..32].try_into().ok()?);
    let lat_step = f64::from_le_bytes(bytes[32..40].try_into().ok()?);
    let lon_step = f64::from_le_bytes(bytes[40..48].try_into().ok()?);
    let lat_count = u16::from_le_bytes(bytes[48..50].try_into().ok()?) as usize;
    let lon_count = u16::from_le_bytes(bytes[50..52].try_into().ok()?) as usize;

    let hl = 52 + n_fields * FIELD_NAME_LEN;
    if bytes.len() < hl {
        return None;
    }
    let mut fields = Vec::with_capacity(n_fields);
    for i in 0..n_fields {
        let start = 52 + i * FIELD_NAME_LEN;
        let raw = &bytes[start..start + FIELD_NAME_LEN];
        let end = raw.iter().position(|&b| b == 0).unwrap_or(FIELD_NAME_LEN);
        fields.push(String::from_utf8_lossy(&raw[..end]).into_owned());
    }

    Some(Header {
        hrdps,
        valid_time,
        spec: GridSpec {
            lat_start,
            lon_start,
            lat_step,
            lon_step,
            lat_count,
            lon_count,
        },
        fields,
    })
}

/// Read one cell's record out of a whole-file buffer.
pub fn read_cell(bytes: &[u8], header: &Header, cell: usize) -> Option<Vec<f32>> {
    let n = header.fields.len();
    let start = header_len() + cell * n * 4;
    let end = start + n * 4;
    if end > bytes.len() {
        return None;
    }
    Some(
        bytes[start..end]
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
            .collect(),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::grid;
    use chrono::TimeZone;

    fn spec_3x2() -> GridSpec {
        GridSpec {
            lon_start: -100.0,
            lon_count: 3,
            lon_step: 0.5,
            lat_start: 30.0,
            lat_count: 2,
            lat_step: 0.5,
        }
    }

    #[test]
    fn field_list_length() {
        assert_eq!(SCALAR_FIELDS.len(), 20);
        assert_eq!(SCALAR_FIELDS[0], "temperature");
        assert_eq!(SCALAR_FIELDS[SCALAR_FIELDS.len() - 1], "duct_cutoff_ghz");
    }

    #[test]
    fn header_round_trips() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = grid::wgrib2_grid_spec();
        let bytes = encode_header(&spec, vt, false);
        let h = decode_header(&bytes).expect("decodes");

        assert!(!h.hrdps);
        assert_eq!(h.valid_time, vt);
        assert_eq!(h.spec, spec);
        assert_eq!(h.fields.len(), SCALAR_FIELDS.len());
        assert_eq!(bytes.len(), header_len());
    }

    #[test]
    fn header_carries_hrdps_flag() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let bytes = encode_header(&grid::hrdps_grid_spec(), vt, true);
        assert!(decode_header(&bytes).unwrap().hrdps);
    }

    #[test]
    fn decode_rejects_bad_magic_and_version() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let mut bytes = encode_header(&spec_3x2(), vt, false);
        bytes[0..4].copy_from_slice(b"XXXX");
        assert!(decode_header(&bytes).is_none());

        let mut bytes = encode_header(&spec_3x2(), vt, false);
        bytes[4] = 99;
        assert!(decode_header(&bytes).is_none());
    }

    #[test]
    fn build_body_places_rows_at_correct_cells() {
        let spec = spec_3x2();
        let n_cells = spec.lat_count * spec.lon_count;
        let nf = SCALAR_FIELDS.len();
        // | j\i | -100.0 | -99.5 | -99.0 |
        // |------|--------|-------|-------|
        // | 30.0 | c0     | c1    | c2    |
        // | 30.5 | c3     | c4    | c5    |

        let rows = vec![
            ScalarRow {
                lat: 30.0,
                lon: -100.0,
                temperature: Some(25.0),
                ..ScalarRow::default()
            },
            ScalarRow {
                lat: 30.5,
                lon: -99.0,
                temperature: Some(15.0),
                dewpoint_depression: Some(5.0),
                ducting: Some(true),
                ..ScalarRow::default()
            },
        ];

        let body = build_body(&spec, &rows);
        assert_eq!(body.len(), n_cells * nf);

        // Cell 0: temperature=25.0, rest NaN
        assert!((body[0] - 25.0_f32).abs() < 1e-6);
        assert!(body[1].is_nan()); // dewpoint_depression
        assert!(body[16].is_nan()); // ducting

        // Cell 5 (30.5, -99.0): temperature=15.0, dewpoint_depression=5.0, ducting=1.0
        let c5 = 5 * nf;
        assert!((body[c5] - 15.0_f32).abs() < 1e-6);
        assert!((body[c5 + 1] - 5.0_f32).abs() < 1e-6);
        assert_eq!(body[c5 + 16], 1.0_f32); // ducting
        assert!(body[c5 + 17].is_nan()); // duct_base_m

        // Cell 1: never written, stays NaN
        assert!(body[nf].is_nan());
        // Cell 4: never written, stays NaN
        assert!(body[4 * nf].is_nan());
    }

    #[test]
    fn build_body_ducting_false_is_zero() {
        let spec = GridSpec {
            lon_start: 0.0,
            lon_count: 1,
            lon_step: 1.0,
            lat_start: 0.0,
            lat_count: 1,
            lat_step: 1.0,
        };
        let rows = vec![ScalarRow {
            lat: 0.0,
            lon: 0.0,
            ducting: Some(false),
            ..ScalarRow::default()
        }];
        let body = build_body(&spec, &rows);
        assert_eq!(
            body[SCALAR_FIELDS.iter().position(|n| *n == "ducting").unwrap()],
            0.0_f32
        );
    }

    #[test]
    fn build_body_out_of_bounds_row_is_skipped() {
        let spec = spec_3x2();
        let rows = vec![ScalarRow {
            lat: 99.0,
            lon: 0.0,
            temperature: Some(42.0),
            ..ScalarRow::default()
        }];
        let body = build_body(&spec, &rows);
        // Every cell should be all-NaN
        assert!(body.iter().all(|v| v.is_nan()));
    }

    #[test]
    fn write_read_round_trip_per_cell() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let n_cells = spec.lat_count * spec.lon_count;
        let nf = SCALAR_FIELDS.len();

        // Cell c, field f gets value c*100 + f, except cell 1 (all-missing).
        let mut body = vec![f32::NAN; n_cells * nf];
        for c in 0..n_cells {
            if c == 1 {
                continue;
            }
            for f in 0..nf {
                body[c * nf + f] = (c * 100 + f) as f32;
            }
        }

        let path = write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        assert!(path.to_string_lossy().ends_with(".sgrid"));

        let raw = std::fs::read(&path).unwrap();
        let h = decode_header(&raw).unwrap();
        assert_eq!(h.spec, spec);
        assert_eq!(
            raw.len(),
            header_len() + n_cells * nf * 4,
            "file is header + dense body, nothing else"
        );

        let cell0 = read_cell(&raw, &h, 0).unwrap();
        assert_eq!(cell0[0], 0.0);
        assert_eq!(cell0[nf - 1], (nf - 1) as f32);

        let cell2 = read_cell(&raw, &h, 2).unwrap();
        assert_eq!(cell2[0], 200.0);

        let cell1 = read_cell(&raw, &h, 1).unwrap();
        assert!(
            cell1.iter().all(|v| v.is_nan()),
            "missing cell reads as NaN"
        );

        assert!(read_cell(&raw, &h, n_cells).is_none(), "out-of-range cell");
    }

    #[test]
    fn atomic_write_leaves_no_tmp_files() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let body = vec![f32::NAN; spec.lat_count * spec.lon_count * SCALAR_FIELDS.len()];

        write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        write_atomic(dir.path(), vt, &spec, &body, false).unwrap();

        let entries: Vec<_> = std::fs::read_dir(dir.path().join("weather_scalars"))
            .unwrap()
            .map(|e| e.unwrap().file_name().into_string().unwrap())
            .collect();
        assert_eq!(entries.len(), 1, "got {entries:?}");
        assert!(entries[0].ends_with(".sgrid"));
    }

    #[test]
    fn hrdps_path_is_sibling_not_clobber() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let body = vec![f32::NAN; spec.lat_count * spec.lon_count * SCALAR_FIELDS.len()];

        write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        write_atomic(dir.path(), vt, &spec, &body, true).unwrap();

        let entries: Vec<_> = std::fs::read_dir(dir.path().join("weather_scalars"))
            .unwrap()
            .map(|e| e.unwrap().file_name().into_string().unwrap())
            .collect();
        assert_eq!(entries.len(), 2);
        assert!(entries
            .iter()
            .any(|e| e.ends_with(".sgrid") && !e.contains(".hrdps")));
        assert!(entries.iter().any(|e| e.ends_with(".hrdps.sgrid")));
    }
}