geotiff-core 0.8.1

//! GeoKey directory parsing and construction (TIFF tag 34735).
//!
//! The GeoKey directory is stored as a TIFF SHORT array with the structure:
//! - Header: KeyDirectoryVersion, KeyRevision, MinorRevision, NumberOfKeys
//! - Entries: KeyID, TIFFTagLocation, Count, ValueOffset (repeated)
//!
//! GeoKeys reference values either inline (location=0), from the
//! GeoDoubleParams tag (34736), or from the GeoAsciiParams tag (34737).

use std::collections::HashSet;
use std::error::Error;
use std::fmt;

// Well-known GeoKey IDs.
pub const GT_MODEL_TYPE: u16 = 1024;
pub const GT_RASTER_TYPE: u16 = 1025;
pub const GT_CITATION: u16 = 1026;
pub const GEODETIC_CRS_TYPE: u16 = 2048;
pub const GEOGRAPHIC_TYPE: u16 = GEODETIC_CRS_TYPE;
pub const GEODETIC_CITATION: u16 = 2049;
pub const GEOG_CITATION: u16 = 2049;
pub const GEODETIC_DATUM: u16 = 2050;
pub const GEOG_GEODETIC_DATUM: u16 = 2050;
pub const GEOG_ANGULAR_UNITS: u16 = 2054;
pub const PROJECTED_CRS_TYPE: u16 = 3072;
pub const PROJECTED_CS_TYPE: u16 = 3072;
pub const PROJ_CITATION: u16 = 3073;
pub const PROJECTION: u16 = 3074;
pub const PROJ_COORD_TRANS: u16 = 3075;
pub const PROJ_LINEAR_UNITS: u16 = 3076;
pub const VERTICAL_CITATION: u16 = 4097;
pub const VERTICAL_CS_TYPE: u16 = 4096;
pub const VERTICAL_DATUM: u16 = 4098;
pub const VERTICAL_UNITS: u16 = 4099;
const GEO_DOUBLE_PARAMS_TAG: u16 = 34736;
const GEO_ASCII_PARAMS_TAG: u16 = 34737;
pub const GEO_KEY_DIRECTORY_VERSION: u16 = 1;
pub const GEO_KEY_REVISION: u16 = 1;
pub const GEO_KEY_MINOR_REVISION_1_0: u16 = 0;
pub const GEO_KEY_MINOR_REVISION_1_1: u16 = 1;

/// Error returned when a GeoKey directory cannot be represented in the
/// GeoTIFF SHORT-based key directory format.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GeoKeySerializeError {
    /// KeyDirectoryVersion must be 1.
    InvalidDirectoryVersion { version: u16 },
    /// KeyRevision must be 1.
    InvalidMajorRevision { major_revision: u16 },
    /// The GeoKey directory minor revision must be 0 or 1.
    InvalidMinorRevision { minor_revision: u16 },
    /// A key ID may occur only once in a directory.
    DuplicateKey { key_id: u16 },
    /// GeoAsciiParams values must be delimiter-free ASCII.
    InvalidAsciiValue { key_id: u16 },
    /// The GeoKey directory header stores the key count as a SHORT.
    TooManyKeys { count: usize },
    /// A key references more parameter values than fit in a SHORT count.
    ValueCountTooLarge { key_id: u16, tag: u16, count: usize },
    /// A key's parameter start offset does not fit in a SHORT value offset.
    ParameterOffsetTooLarge {
        key_id: u16,
        tag: u16,
        offset: usize,
    },
}

impl fmt::Display for GeoKeySerializeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidDirectoryVersion { version } => write!(
                f,
                "GeoKey directory version {version} is invalid; expected 1"
            ),
            Self::InvalidMajorRevision { major_revision } => write!(
                f,
                "GeoKey directory major revision {major_revision} is invalid; expected 1"
            ),
            Self::InvalidMinorRevision { minor_revision } => write!(
                f,
                "GeoKey directory minor revision {minor_revision} is invalid; expected 0 or 1"
            ),
            Self::DuplicateKey { key_id } => {
                write!(f, "GeoKey directory contains duplicate key ID {key_id}")
            }
            Self::InvalidAsciiValue { key_id } => write!(
                f,
                "GeoKey {key_id} ASCII value must contain only ASCII characters and no '|' delimiter"
            ),
            Self::TooManyKeys { count } => {
                write!(
                    f,
                    "GeoKey directory contains {count} keys, exceeding u16::MAX"
                )
            }
            Self::ValueCountTooLarge { key_id, tag, count } => write!(
                f,
                "GeoKey {key_id} references {count} values in tag {tag}, exceeding u16::MAX"
            ),
            Self::ParameterOffsetTooLarge {
                key_id,
                tag,
                offset,
            } => write!(
                f,
                "GeoKey {key_id} parameter offset {offset} in tag {tag} exceeds u16::MAX"
            ),
        }
    }
}

impl Error for GeoKeySerializeError {}

/// A parsed GeoKey entry.
#[derive(Debug, Clone)]
pub struct GeoKey {
    pub id: u16,
    pub value: GeoKeyValue,
}

/// The value of a GeoKey.
#[derive(Debug, Clone)]
pub enum GeoKeyValue {
    /// Short value stored inline.
    Short(u16),
    /// Double value(s) from GeoDoubleParams.
    Double(Vec<f64>),
    /// ASCII string from GeoAsciiParams.
    Ascii(String),
}

/// Parsed GeoKey directory.
#[derive(Debug, Clone)]
pub struct GeoKeyDirectory {
    pub version: u16,
    pub major_revision: u16,
    pub minor_revision: u16,
    pub keys: Vec<GeoKey>,
}

impl GeoKeyDirectory {
    /// Create an empty directory with a GeoTIFF 1.1 writer header.
    pub fn new() -> Self {
        Self {
            version: GEO_KEY_DIRECTORY_VERSION,
            major_revision: GEO_KEY_REVISION,
            minor_revision: GEO_KEY_MINOR_REVISION_1_1,
            keys: Vec::new(),
        }
    }

    /// Parse the GeoKey directory from the three GeoTIFF tags.
    ///
    /// - `directory`: contents of tag 34735 (SHORT array)
    /// - `double_params`: contents of tag 34736 (DOUBLE array), may be empty
    /// - `ascii_params`: contents of tag 34737 (ASCII), may be empty
    pub fn parse(directory: &[u16], double_params: &[f64], ascii_params: &str) -> Option<Self> {
        if directory.len() < 4 {
            return None;
        }

        let version = directory[0];
        let major_revision = directory[1];
        let minor_revision = directory[2];
        let num_keys = directory[3] as usize;

        if version != GEO_KEY_DIRECTORY_VERSION
            || major_revision != GEO_KEY_REVISION
            || !matches!(
                minor_revision,
                GEO_KEY_MINOR_REVISION_1_0 | GEO_KEY_MINOR_REVISION_1_1
            )
            || directory.len() != 4 + num_keys * 4
            || !ascii_params.is_ascii()
        {
            return None;
        }

        let mut keys = Vec::with_capacity(num_keys);
        let mut seen_ids = HashSet::with_capacity(num_keys);
        for i in 0..num_keys {
            let base = 4 + i * 4;
            let key_id = directory[base];
            let location = directory[base + 1];
            let count = directory[base + 2] as usize;
            let value_offset = directory[base + 3];
            if !seen_ids.insert(key_id) {
                return None;
            }

            let value = match location {
                0 => {
                    // Value is the offset itself (short).
                    if count != 1 {
                        return None;
                    }
                    GeoKeyValue::Short(value_offset)
                }
                34736 => {
                    // Value is in GeoDoubleParams.
                    if count == 0 {
                        return None;
                    }
                    let start = value_offset as usize;
                    let end = start.checked_add(count)?;
                    if end <= double_params.len() {
                        GeoKeyValue::Double(double_params[start..end].to_vec())
                    } else {
                        return None;
                    }
                }
                34737 => {
                    // Value is in GeoAsciiParams.
                    if count == 0 {
                        return None;
                    }
                    let start = value_offset as usize;
                    let end = start.checked_add(count)?;
                    if let Some(raw) = ascii_params.get(start..end) {
                        if !raw.ends_with('|') {
                            return None;
                        }
                        let s = raw.trim_end_matches('|').trim_end_matches('\0').to_string();
                        GeoKeyValue::Ascii(s)
                    } else {
                        return None;
                    }
                }
                _ => return None,
            };

            keys.push(GeoKey { id: key_id, value });
        }

        Some(Self {
            version,
            major_revision,
            minor_revision,
            keys,
        })
    }

    /// Look up a GeoKey by ID.
    pub fn get(&self, id: u16) -> Option<&GeoKey> {
        self.keys.iter().find(|k| k.id == id)
    }

    /// Get a short value for a key.
    pub fn get_short(&self, id: u16) -> Option<u16> {
        self.get(id).and_then(|k| match &k.value {
            GeoKeyValue::Short(v) => Some(*v),
            _ => None,
        })
    }

    /// Get an ASCII value for a key.
    pub fn get_ascii(&self, id: u16) -> Option<&str> {
        self.get(id).and_then(|k| match &k.value {
            GeoKeyValue::Ascii(s) => Some(s.as_str()),
            _ => None,
        })
    }

    /// Get double value(s) for a key.
    pub fn get_double(&self, id: u16) -> Option<&[f64]> {
        self.get(id).and_then(|k| match &k.value {
            GeoKeyValue::Double(v) => Some(v.as_slice()),
            _ => None,
        })
    }

    /// Insert or replace a GeoKey.
    pub fn set(&mut self, id: u16, value: GeoKeyValue) {
        if let Some(existing) = self.keys.iter_mut().find(|k| k.id == id) {
            existing.value = value;
        } else {
            self.keys.push(GeoKey { id, value });
        }
    }

    /// Remove a GeoKey by ID.
    pub fn remove(&mut self, id: u16) {
        self.keys.retain(|k| k.id != id);
    }

    /// Serialize the directory into the three TIFF tag payloads.
    ///
    /// Returns `(directory_shorts, double_params, ascii_params)`.
    /// Keys are sorted by ID per spec. Short values go inline (location=0),
    /// Double values reference the double_params array (location=34736),
    /// Ascii values reference the ascii_params string (location=34737).
    pub fn serialize(&self) -> Result<(Vec<u16>, Vec<f64>, String), GeoKeySerializeError> {
        if self.version != GEO_KEY_DIRECTORY_VERSION {
            return Err(GeoKeySerializeError::InvalidDirectoryVersion {
                version: self.version,
            });
        }
        if self.major_revision != GEO_KEY_REVISION {
            return Err(GeoKeySerializeError::InvalidMajorRevision {
                major_revision: self.major_revision,
            });
        }

        let mut sorted_keys = self.keys.clone();
        sorted_keys.sort_by_key(|k| k.id);
        if let Some(duplicate) = sorted_keys.windows(2).find(|keys| keys[0].id == keys[1].id) {
            return Err(GeoKeySerializeError::DuplicateKey {
                key_id: duplicate[0].id,
            });
        }
        let key_count =
            u16::try_from(sorted_keys.len()).map_err(|_| GeoKeySerializeError::TooManyKeys {
                count: sorted_keys.len(),
            })?;

        let mut directory = Vec::new();
        let mut double_params = Vec::new();
        let mut ascii_params = String::new();

        let minor_revision = self.serialized_minor_revision()?;

        // Header: version, major_revision, minor_revision, num_keys
        directory.push(self.version);
        directory.push(self.major_revision);
        directory.push(minor_revision);
        directory.push(key_count);

        for key in &sorted_keys {
            directory.push(key.id);
            match &key.value {
                GeoKeyValue::Short(v) => {
                    directory.push(0); // location: inline
                    directory.push(1); // count
                    directory.push(*v); // value
                }
                GeoKeyValue::Double(v) => {
                    let count = checked_u16_len(key.id, GEO_DOUBLE_PARAMS_TAG, v.len())?;
                    let offset =
                        checked_u16_offset(key.id, GEO_DOUBLE_PARAMS_TAG, double_params.len())?;
                    directory.push(GEO_DOUBLE_PARAMS_TAG); // location: GeoDoubleParams
                    directory.push(count);
                    directory.push(offset);
                    double_params.extend_from_slice(v);
                }
                GeoKeyValue::Ascii(s) => {
                    if !s.is_ascii() || s.contains('|') {
                        return Err(GeoKeySerializeError::InvalidAsciiValue { key_id: key.id });
                    }
                    let ascii_with_pipe = format!("{}|", s);
                    let count =
                        checked_u16_len(key.id, GEO_ASCII_PARAMS_TAG, ascii_with_pipe.len())?;
                    let offset =
                        checked_u16_offset(key.id, GEO_ASCII_PARAMS_TAG, ascii_params.len())?;
                    directory.push(GEO_ASCII_PARAMS_TAG); // location: GeoAsciiParams
                    directory.push(count);
                    directory.push(offset);
                    ascii_params.push_str(&ascii_with_pipe);
                }
            }
        }

        Ok((directory, double_params, ascii_params))
    }

    fn serialized_minor_revision(&self) -> Result<u16, GeoKeySerializeError> {
        match self.minor_revision {
            GEO_KEY_MINOR_REVISION_1_0 if self.requires_geotiff_1_1() => {
                Ok(GEO_KEY_MINOR_REVISION_1_1)
            }
            GEO_KEY_MINOR_REVISION_1_0 | GEO_KEY_MINOR_REVISION_1_1 => Ok(self.minor_revision),
            minor_revision => Err(GeoKeySerializeError::InvalidMinorRevision { minor_revision }),
        }
    }

    fn requires_geotiff_1_1(&self) -> bool {
        self.keys.iter().any(|key| {
            matches!(
                key.id,
                VERTICAL_CS_TYPE | VERTICAL_CITATION | VERTICAL_DATUM | VERTICAL_UNITS
            )
        })
    }
}

fn checked_u16_len(key_id: u16, tag: u16, count: usize) -> Result<u16, GeoKeySerializeError> {
    u16::try_from(count).map_err(|_| GeoKeySerializeError::ValueCountTooLarge {
        key_id,
        tag,
        count,
    })
}

fn checked_u16_offset(key_id: u16, tag: u16, offset: usize) -> Result<u16, GeoKeySerializeError> {
    u16::try_from(offset).map_err(|_| GeoKeySerializeError::ParameterOffsetTooLarge {
        key_id,
        tag,
        offset,
    })
}

impl Default for GeoKeyDirectory {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn parse_roundtrip() {
        let mut dir = GeoKeyDirectory::new();
        dir.set(GT_MODEL_TYPE, GeoKeyValue::Short(2));
        dir.set(GEOGRAPHIC_TYPE, GeoKeyValue::Short(4326));
        dir.set(GEOG_CITATION, GeoKeyValue::Ascii("WGS 84".into()));

        let (shorts, doubles, ascii) = dir.serialize().unwrap();
        assert_eq!(shorts[..4], [1, 1, 1, 3]);
        let parsed = GeoKeyDirectory::parse(&shorts, &doubles, &ascii).unwrap();

        assert_eq!(parsed.get_short(GT_MODEL_TYPE), Some(2));
        assert_eq!(parsed.get_short(GEOGRAPHIC_TYPE), Some(4326));
        assert_eq!(parsed.get_ascii(GEOG_CITATION), Some("WGS 84"));
    }

    #[test]
    fn serialize_preserves_legacy_minor_revision_zero_when_compatible() {
        let mut dir = GeoKeyDirectory::new();
        dir.minor_revision = GEO_KEY_MINOR_REVISION_1_0;
        dir.set(GT_MODEL_TYPE, GeoKeyValue::Short(2));

        let (shorts, _, _) = dir.serialize().unwrap();
        assert_eq!(shorts[..4], [1, 1, 0, 1]);
    }

    #[test]
    fn serialize_promotes_vertical_geokeys_to_geotiff_1_1_minor_revision() {
        let mut dir = GeoKeyDirectory::new();
        dir.minor_revision = GEO_KEY_MINOR_REVISION_1_0;
        dir.set(VERTICAL_CS_TYPE, GeoKeyValue::Short(5703));

        let (shorts, _, _) = dir.serialize().unwrap();
        assert_eq!(shorts[..4], [1, 1, 1, 1]);
    }

    #[test]
    fn serialize_rejects_invalid_minor_revision() {
        let mut dir = GeoKeyDirectory::new();
        dir.minor_revision = 2;

        let err = dir.serialize().unwrap_err();
        assert_eq!(
            err,
            GeoKeySerializeError::InvalidMinorRevision { minor_revision: 2 }
        );
    }

    #[test]
    fn parse_rejects_invalid_headers_duplicate_keys_and_inline_counts() {
        assert!(GeoKeyDirectory::parse(&[2, 1, 1, 0], &[], "").is_none());
        assert!(GeoKeyDirectory::parse(&[1, 2, 1, 0], &[], "").is_none());
        assert!(GeoKeyDirectory::parse(&[1, 1, 2, 0], &[], "").is_none());
        assert!(GeoKeyDirectory::parse(&[1, 1, 1, 1, GT_MODEL_TYPE, 0, 2, 1], &[], "").is_none());
        assert!(GeoKeyDirectory::parse(
            &[1, 1, 1, 2, GT_MODEL_TYPE, 0, 1, 1, GT_MODEL_TYPE, 0, 1, 2],
            &[],
            ""
        )
        .is_none());
    }

    #[test]
    fn serialize_rejects_invalid_headers_duplicates_and_non_ascii_values() {
        let mut dir = GeoKeyDirectory::new();
        dir.version = 2;
        assert!(matches!(
            dir.serialize(),
            Err(GeoKeySerializeError::InvalidDirectoryVersion { version: 2 })
        ));

        let mut dir = GeoKeyDirectory::new();
        dir.keys = vec![
            GeoKey {
                id: GT_MODEL_TYPE,
                value: GeoKeyValue::Short(1),
            },
            GeoKey {
                id: GT_MODEL_TYPE,
                value: GeoKeyValue::Short(2),
            },
        ];
        assert!(matches!(
            dir.serialize(),
            Err(GeoKeySerializeError::DuplicateKey {
                key_id: GT_MODEL_TYPE
            })
        ));

        let mut dir = GeoKeyDirectory::new();
        dir.set(GEOG_CITATION, GeoKeyValue::Ascii("WGS 84 | invalid".into()));
        assert!(matches!(
            dir.serialize(),
            Err(GeoKeySerializeError::InvalidAsciiValue {
                key_id: GEOG_CITATION
            })
        ));
    }

    #[test]
    fn set_replaces_existing() {
        let mut dir = GeoKeyDirectory::new();
        dir.set(GT_MODEL_TYPE, GeoKeyValue::Short(1));
        dir.set(GT_MODEL_TYPE, GeoKeyValue::Short(2));
        assert_eq!(dir.get_short(GT_MODEL_TYPE), Some(2));
        assert_eq!(dir.keys.len(), 1);
    }

    #[test]
    fn remove_key() {
        let mut dir = GeoKeyDirectory::new();
        dir.set(GT_MODEL_TYPE, GeoKeyValue::Short(1));
        dir.remove(GT_MODEL_TYPE);
        assert!(dir.get(GT_MODEL_TYPE).is_none());
    }

    #[test]
    fn parse_rejects_invalid_parameter_references_without_panicking() {
        let directory = [
            1u16,
            1,
            0,
            1, // header
            GEOG_CITATION,
            34737,
            1,
            1, // byte offsets that are invalid for lossy UTF-8
        ];
        let ascii = String::from_utf8_lossy(&[0xff, b'|']).into_owned();

        assert!(GeoKeyDirectory::parse(&directory, &[], &ascii).is_none());
        assert!(
            GeoKeyDirectory::parse(&[1, 1, 1, 1, GEOG_CITATION, 34737, 2, 0], &[], "x|").is_some()
        );
        assert!(
            GeoKeyDirectory::parse(&[1, 1, 1, 1, GEOG_CITATION, 34737, 1, 0], &[], "x").is_none()
        );
        assert!(
            GeoKeyDirectory::parse(&[1, 1, 1, 1, GEOG_CITATION, 65000, 1, 0], &[], "").is_none()
        );
    }

    #[test]
    fn serialize_rejects_too_many_keys() {
        let mut dir = GeoKeyDirectory::new();
        dir.keys = (0..=u16::MAX as usize)
            .map(|index| GeoKey {
                id: index as u16,
                value: GeoKeyValue::Short(1),
            })
            .collect();

        let err = dir.serialize().unwrap_err();
        assert_eq!(
            err,
            GeoKeySerializeError::TooManyKeys {
                count: u16::MAX as usize + 1
            }
        );
    }

    #[test]
    fn serialize_rejects_oversized_double_value_count() {
        let mut dir = GeoKeyDirectory::new();
        dir.set(
            GT_CITATION,
            GeoKeyValue::Double(vec![1.0; u16::MAX as usize + 1]),
        );

        let err = dir.serialize().unwrap_err();
        assert_eq!(
            err,
            GeoKeySerializeError::ValueCountTooLarge {
                key_id: GT_CITATION,
                tag: GEO_DOUBLE_PARAMS_TAG,
                count: u16::MAX as usize + 1
            }
        );
    }

    #[test]
    fn serialize_rejects_oversized_ascii_parameter_offset() {
        let mut dir = GeoKeyDirectory::new();
        dir.set(
            GEOG_CITATION,
            GeoKeyValue::Ascii("a".repeat(u16::MAX as usize - 1)),
        );
        dir.set(PROJ_CITATION, GeoKeyValue::Ascii("b".to_string()));
        dir.set(VERTICAL_CITATION, GeoKeyValue::Ascii("c".to_string()));

        let err = dir.serialize().unwrap_err();
        assert_eq!(
            err,
            GeoKeySerializeError::ParameterOffsetTooLarge {
                key_id: VERTICAL_CITATION,
                tag: GEO_ASCII_PARAMS_TAG,
                offset: u16::MAX as usize + 2
            }
        );
    }
}