//! Heightmap loader — reads the canonical 16-bit grayscale elevation PNG (D-202).
//!
//! Per D-202 (amended #963), each body's canonical elevation is a per-body
//! **16-bit grayscale `heightmap.png`** stored next to the color `reliefmap.png`
//! in the body's wiki directory (path from `bodies.terrain_reference`). The
//! grayscale luminance is the normalized elevation; this is the single source
//! of truth — no `systems.db` BLOB. The file is canonically 1024×512 (D-202
//! amended): native elevation detail for the lower cascade layers
//! (region/block/tile), while **Layer 1** (continental drainage/basins/
//! mountain-ranges) calls [`BodyHeightmap::downsample`] to the
//! `GRID_W × GRID_H = 512×256` working grid first, decoupling continental
//! compute cost from the stored resolution.
//!
//! Elevation values are normalized to [0.0, 1.0]. `sea_level` is body metadata
//! (carried alongside, not in the PNG) below which terrain is ocean.

use std::fs::File;
use std::io::Read;
use std::path::Path;

use thiserror::Error;

/// Layer-1 working resolution: continental drainage/feature extraction run at
/// this size (the stored heightmap is higher-res and downsampled to it).
pub const GRID_W: u32 = 512;
pub const GRID_H: u32 = 256;

/// A loaded heightmap for one planetary body.
#[derive(Debug, Clone)]
pub struct BodyHeightmap {
    pub body_id: String,
    pub width: u32,
    pub height: u32,
    /// Row-major elevation values, normalized to [0.0, 1.0].
    pub data: Vec<f32>,
    /// Elevation fraction below which terrain is ocean/sea.
    pub sea_level: f32,
}

impl BodyHeightmap {
    /// Returns the elevation at (row, col), or `None` if out of bounds.
    #[inline]
    pub fn get(&self, row: u32, col: u32) -> Option<f32> {
        if row < self.height && col < self.width {
            Some(self.data[(row * self.width + col) as usize])
        } else {
            None
        }
    }

    /// Returns `true` if the cell at (row, col) is land (above sea level).
    #[inline]
    pub fn is_land(&self, row: u32, col: u32) -> bool {
        self.get(row, col).is_some_and(|e| e >= self.sea_level)
    }

    /// Box-average downsample to `target_w × target_h` (used by Layer 1 to drop
    /// the high-res stored heightmap to the continental working resolution).
    /// Deterministic (f64 accumulation, fixed footprint). Returns a clone when
    /// the target is not strictly smaller (no upsampling here).
    pub fn downsample(&self, target_w: u32, target_h: u32) -> BodyHeightmap {
        if target_w == 0 || target_h == 0 || (target_w >= self.width && target_h >= self.height) {
            return self.clone();
        }
        let (sw, sh) = (self.width as u64, self.height as u64);
        let mut out = vec![0f32; (target_w * target_h) as usize];
        for ty in 0..target_h {
            let sy0 = (ty as u64 * sh / target_h as u64) as u32;
            let sy1 = (((ty as u64 + 1) * sh / target_h as u64).max(sy0 as u64 + 1)) as u32;
            for tx in 0..target_w {
                let sx0 = (tx as u64 * sw / target_w as u64) as u32;
                let sx1 = (((tx as u64 + 1) * sw / target_w as u64).max(sx0 as u64 + 1)) as u32;
                let mut sum = 0f64;
                let mut cnt = 0u32;
                for sy in sy0..sy1 {
                    for sx in sx0..sx1 {
                        sum += self.data[(sy * self.width + sx) as usize] as f64;
                        cnt += 1;
                    }
                }
                out[(ty * target_w + tx) as usize] = (sum / cnt as f64) as f32;
            }
        }
        BodyHeightmap {
            body_id: self.body_id.clone(),
            width: target_w,
            height: target_h,
            data: out,
            sea_level: self.sea_level,
        }
    }
}

#[derive(Debug, Error)]
pub enum HeightmapLoadError {
    #[error("heightmap file unreadable: {0}")]
    Io(#[from] std::io::Error),
    #[error("PNG decode error: {0}")]
    Png(#[from] png::DecodingError),
    #[error("unsupported heightmap PNG: {bit_depth} {color_type} (expected 16-bit grayscale)")]
    UnsupportedFormat {
        bit_depth: String,
        color_type: String,
    },
}

/// Load the canonical heightmap PNG for `body_id` from `path`.
///
/// `sea_level` is body metadata (e.g. from the `bodies` table / body def), not
/// stored in the PNG.
pub fn load_heightmap_png(
    path: &Path,
    body_id: &str,
    default_sea_level: f32,
) -> Result<BodyHeightmap, HeightmapLoadError> {
    let file = File::open(path)?;
    load_heightmap_reader(file, body_id, default_sea_level)
}

/// Decode a heightmap PNG from any reader. 16-bit grayscale is the canonical
/// format; 8-bit grayscale is accepted (coarse — viewable/test only).
///
/// `sea_level` is read from the PNG's `sea_level` tEXt chunk (written by the
/// bake); `default_sea_level` is the fallback when the chunk is absent.
pub fn load_heightmap_reader<R: Read>(
    reader: R,
    body_id: &str,
    default_sea_level: f32,
) -> Result<BodyHeightmap, HeightmapLoadError> {
    let mut png_reader = png::Decoder::new(reader).read_info()?;
    let (width, height, bit_depth, color_type, sea_level) = {
        let info = png_reader.info();
        let sea_level = info
            .uncompressed_latin1_text
            .iter()
            .find(|c| c.keyword == "sea_level")
            .and_then(|c| c.text.trim().parse::<f32>().ok())
            .unwrap_or(default_sea_level);
        (info.width, info.height, info.bit_depth, info.color_type, sea_level)
    };
    let mut buf = vec![0u8; png_reader.output_buffer_size()];
    let frame = png_reader.next_frame(&mut buf)?;
    let bytes = &buf[..frame.buffer_size()];

    let data: Vec<f32> = match (bit_depth, color_type) {
        (png::BitDepth::Sixteen, png::ColorType::Grayscale) => bytes
            .chunks_exact(2)
            .map(|p| u16::from_be_bytes([p[0], p[1]]) as f32 / 65535.0)
            .collect(),
        (png::BitDepth::Eight, png::ColorType::Grayscale) => {
            bytes.iter().map(|&b| b as f32 / 255.0).collect()
        }
        (bd, ct) => {
            return Err(HeightmapLoadError::UnsupportedFormat {
                bit_depth: format!("{bd:?}"),
                color_type: format!("{ct:?}"),
            })
        }
    };

    Ok(BodyHeightmap {
        body_id: body_id.to_string(),
        width,
        height,
        data,
        sea_level,
    })
}

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

    /// Encode a 16-bit grayscale PNG (row-major u16 elevation) to bytes.
    fn encode_gray16(w: u32, h: u32, vals: &[u16]) -> Vec<u8> {
        let mut out = Vec::new();
        {
            let mut enc = png::Encoder::new(&mut out, w, h);
            enc.set_color(png::ColorType::Grayscale);
            enc.set_depth(png::BitDepth::Sixteen);
            let mut writer = enc.write_header().unwrap();
            // png expects big-endian 16-bit samples.
            let mut be = Vec::with_capacity(vals.len() * 2);
            for &v in vals {
                be.extend_from_slice(&v.to_be_bytes());
            }
            writer.write_image_data(&be).unwrap();
        }
        out
    }

    #[test]
    fn round_trips_16bit_grayscale() {
        let (w, h) = (4u32, 2u32);
        // Elevations 0, 1/7, ... 7/7 across 8 cells.
        let vals: Vec<u16> = (0..(w * h))
            .map(|i| (i as f32 / (w * h - 1) as f32 * 65535.0) as u16)
            .collect();
        let png_bytes = encode_gray16(w, h, &vals);
        let hm = load_heightmap_reader(png_bytes.as_slice(), "T", 0.3).unwrap();
        assert_eq!((hm.width, hm.height), (w, h));
        assert_eq!(hm.data.len(), 8);
        assert!((hm.data[0] - 0.0).abs() < 1e-6);
        assert!((hm.data[7] - 1.0).abs() < 1e-4);
        assert!(hm.is_land(1, 3)); // last cell = 1.0 > sea 0.3
        assert!(!hm.is_land(0, 0)); // first cell = 0.0 < sea
    }

    #[test]
    fn downsample_halves_and_averages() {
        // 4×2 → 2×1: each output cell averages a 2×2 block.
        let vals: Vec<u16> = vec![0, 0, 65535, 65535, 0, 0, 65535, 65535];
        let hm = load_heightmap_reader(encode_gray16(4, 2, &vals).as_slice(), "T", 0.3).unwrap();
        let ds = hm.downsample(2, 1);
        assert_eq!((ds.width, ds.height), (2, 1));
        assert!((ds.data[0] - 0.0).abs() < 1e-4); // left 2×2 block all 0
        assert!((ds.data[1] - 1.0).abs() < 1e-4); // right 2×2 block all 1
    }

    #[test]
    fn downsample_is_deterministic() {
        let vals: Vec<u16> = (0..(16 * 8)).map(|i| (i * 257) as u16).collect();
        let hm = load_heightmap_reader(encode_gray16(16, 8, &vals).as_slice(), "T", 0.3).unwrap();
        let a = hm.downsample(4, 2);
        let b = hm.downsample(4, 2);
        assert_eq!(a.data, b.data);
    }

    #[test]
    fn rejects_rgb() {
        // An RGB PNG (like the reliefmap) must be rejected, not misread.
        let mut out = Vec::new();
        {
            let mut enc = png::Encoder::new(&mut out, 2, 2);
            enc.set_color(png::ColorType::Rgb);
            enc.set_depth(png::BitDepth::Eight);
            let mut w = enc.write_header().unwrap();
            w.write_image_data(&[0u8; 12]).unwrap();
        }
        let err = load_heightmap_reader(out.as_slice(), "T", 0.3).unwrap_err();
        assert!(matches!(err, HeightmapLoadError::UnsupportedFormat { .. }));
    }

    #[test]
    fn reads_sea_level_from_text_chunk() {
        // The bake writes sea_level as a tEXt chunk; the loader must prefer it
        // over the supplied default.
        let mut out = Vec::new();
        {
            let mut enc = png::Encoder::new(&mut out, 2, 1);
            enc.set_color(png::ColorType::Grayscale);
            enc.set_depth(png::BitDepth::Sixteen);
            enc.add_text_chunk("sea_level".to_string(), "0.42".to_string())
                .unwrap();
            let mut w = enc.write_header().unwrap();
            w.write_image_data(&[0u8; 4]).unwrap(); // 2×1 × 2 bytes
        }
        let hm = load_heightmap_reader(out.as_slice(), "T", 0.1).unwrap();
        assert!(
            (hm.sea_level - 0.42).abs() < 1e-6,
            "sea_level must come from the tEXt chunk, got {}",
            hm.sea_level
        );
    }
}