molframe-core 0.5.1

//! Content and suffix based format recognition.

use super::super::input::InputBuffer;
use crate::diagnostic::{Code, Diagnostic};

const PDBML_PROBE_WINDOWS: usize = 2_048;
const MESSAGEPACK_PROBE_WINDOWS: usize = 512;

/// A format molframe can read or write.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
#[non_exhaustive]
pub enum Format {
    /// Decide from the content, then from the file name.
    #[default]
    Auto,
    /// The archive's structured text format.
    Mmcif,
    /// XML serialization of the `PDBx` data model.
    Pdbml,
    /// The archive's `MessagePack`-based binary format.
    BinaryCif,
    /// The deprecated but still encountered macromolecular transmission format.
    Mmtf,
    /// The legacy fixed-column format.
    Pdb,
    /// PDB-shaped coordinates carrying partial charge and radius.
    Pqr,
    /// `AutoDock`'s PDB-shaped coordinates carrying charge and atom type.
    Pdbqt,
    /// MDL structure-data file: one or more `CTfile` records.
    Sdf,
    /// Tripos MOL2 molecule records.
    Mol2,
    /// Core-CIF small-molecule crystal data with fractional coordinates.
    SmallCif,
}

impl Format {
    /// The formats a caller may name, excluding automatic detection.
    pub const NAMED: [Self; 10] = [
        Self::Mmcif,
        Self::Pdbml,
        Self::BinaryCif,
        Self::Mmtf,
        Self::Pdbqt,
        Self::Pqr,
        Self::Pdb,
        Self::Sdf,
        Self::Mol2,
        Self::SmallCif,
    ];

    /// The format's short name.
    #[must_use]
    pub const fn name(self) -> &'static str {
        match self {
            Self::Auto => "auto",
            Self::Mmcif => "mmcif",
            Self::Pdbml => "pdbml",
            Self::BinaryCif => "bcif",
            Self::Mmtf => "mmtf",
            Self::Pdb => "pdb",
            Self::Pqr => "pqr",
            Self::Pdbqt => "pdbqt",
            Self::Sdf => "sdf",
            Self::Mol2 => "mol2",
            Self::SmallCif => "smallcif",
        }
    }

    /// The file suffixes conventionally used for this format.
    #[must_use]
    pub const fn extensions(self) -> &'static [&'static str] {
        match self {
            Self::Auto | Self::SmallCif => &[],
            Self::Mmcif => &["cif", "mmcif"],
            Self::Pdbml => &["xml", "pdbml"],
            Self::BinaryCif => &["bcif"],
            Self::Mmtf => &["mmtf"],
            Self::Pdb => &["pdb", "ent"],
            Self::Pqr => &["pqr"],
            Self::Pdbqt => &["pdbqt"],
            Self::Sdf => &["sdf", "mol"],
            Self::Mol2 => &["mol2"],
        }
    }

    /// Parses a format name without allocating a normalized copy.
    #[must_use]
    pub fn parse(name: &str) -> Option<Self> {
        if name.eq_ignore_ascii_case("auto") {
            Some(Self::Auto)
        } else if name.eq_ignore_ascii_case("mmcif") || name.eq_ignore_ascii_case("cif") {
            Some(Self::Mmcif)
        } else if name.eq_ignore_ascii_case("pdbml") || name.eq_ignore_ascii_case("xml") {
            Some(Self::Pdbml)
        } else if name.eq_ignore_ascii_case("bcif") || name.eq_ignore_ascii_case("binarycif") {
            Some(Self::BinaryCif)
        } else if name.eq_ignore_ascii_case("mmtf") {
            Some(Self::Mmtf)
        } else if name.eq_ignore_ascii_case("pdb") || name.eq_ignore_ascii_case("ent") {
            Some(Self::Pdb)
        } else if name.eq_ignore_ascii_case("pqr") {
            Some(Self::Pqr)
        } else if name.eq_ignore_ascii_case("pdbqt") {
            Some(Self::Pdbqt)
        } else if name.eq_ignore_ascii_case("sdf") || name.eq_ignore_ascii_case("mol") {
            Some(Self::Sdf)
        } else if name.eq_ignore_ascii_case("mol2") {
            Some(Self::Mol2)
        } else if name.eq_ignore_ascii_case("smallcif") {
            Some(Self::SmallCif)
        } else {
            None
        }
    }

    /// Returns true when this format's content is recognisable in `bytes`.
    #[must_use]
    pub fn recognises(self, bytes: &[u8]) -> bool {
        match self {
            Self::Auto | Self::Pqr => false,
            Self::Mmcif => first_lines(bytes, 8).any(|line| line.starts_with(b"data_")),
            Self::Pdbml => contains_in_probe(bytes, b"PDBx:datablock", PDBML_PROBE_WINDOWS),
            Self::BinaryCif => recognises_messagepack_map(bytes, b"dataBlocks"),
            Self::Mmtf => recognises_messagepack_map(bytes, b"mmtfVersion"),
            Self::Pdbqt => recognises_pdbqt(bytes),
            Self::Pdb => recognises_pdb(bytes),
            Self::Sdf => recognises_sdf(bytes),
            Self::Mol2 => recognises_mol2(bytes),
            Self::SmallCif => recognises_small_cif(bytes),
        }
    }

    /// Decides the format of `input`, given what the caller asked for and the
    /// name the bytes came from.
    ///
    /// # Errors
    ///
    /// Returns a diagnostic when neither content nor the optional name identifies
    /// a supported format.
    pub fn detect(
        requested: Self,
        input: &InputBuffer,
        name: Option<&str>,
    ) -> Result<Self, Diagnostic> {
        if requested != Self::Auto {
            return Ok(requested);
        }

        let bytes = input.as_bytes();
        for format in [
            Self::SmallCif,
            Self::Mmcif,
            Self::Pdbml,
            Self::BinaryCif,
            Self::Mmtf,
            Self::Pdbqt,
            Self::Sdf,
            Self::Mol2,
        ] {
            if format.recognises(bytes) {
                return Ok(format);
            }
        }

        let named = name.and_then(Self::from_name);
        if recognises_pdb(bytes) {
            if let Some(format @ (Self::Pqr | Self::Pdbqt)) = named {
                return Ok(format);
            }
            return Ok(Self::Pdb);
        }
        if let Some(format) = named {
            return Ok(format);
        }

        let Some(name) = name else {
            return Err(Diagnostic::new(Code::E1001)
                .with_context("tried", "content, then file name")
                .with_context("name", "(none)"));
        };
        Err(Diagnostic::new(Code::E1001)
            .with_context("tried", "content, then file name")
            .with_context("name", name))
    }

    /// The format a file name suggests, ignoring any compression suffix.
    #[must_use]
    pub fn from_name(name: &str) -> Option<Self> {
        let mut suffixes = name.rsplit('.');
        let mut candidate = suffixes.next()?;
        if is_compression_suffix(candidate) {
            candidate = suffixes.next()?;
        }
        Self::from_extension(candidate)
    }

    /// The format a bare extension names (`pdb`, `cif`, `bcif`, ...), compared
    /// without regard to case.
    #[must_use]
    pub fn from_extension(extension: &str) -> Option<Self> {
        Self::NAMED.into_iter().find(|format| {
            format
                .extensions()
                .iter()
                .any(|known| extension.eq_ignore_ascii_case(known))
        })
    }
}

fn is_compression_suffix(suffix: &str) -> bool {
    suffix.eq_ignore_ascii_case("gz")
        || suffix.eq_ignore_ascii_case("zst")
        || suffix.eq_ignore_ascii_case("bz2")
        || suffix.eq_ignore_ascii_case("xz")
}

fn is_messagepack_map_prefix(byte: u8) -> bool {
    (0x80..=0x8f).contains(&byte) || matches!(byte, 0xde | 0xdf)
}

fn recognises_messagepack_map(bytes: &[u8], key: &[u8]) -> bool {
    bytes
        .first()
        .copied()
        .is_some_and(is_messagepack_map_prefix)
        && contains_in_probe(bytes, key, MESSAGEPACK_PROBE_WINDOWS)
}

fn contains_in_probe(bytes: &[u8], needle: &[u8], windows: usize) -> bool {
    if needle.is_empty() {
        return true;
    }
    bytes
        .windows(needle.len())
        .take(windows)
        .any(|window| window == needle)
}

fn recognises_pdb(bytes: &[u8]) -> bool {
    first_lines(bytes, 64).any(|line| {
        line.starts_with(b"ATOM  ")
            || line.starts_with(b"HETATM")
            || line.starts_with(b"HEADER")
            || line.starts_with(b"CRYST1")
            || line.starts_with(b"MODEL ")
    })
}

fn recognises_pdbqt(bytes: &[u8]) -> bool {
    first_lines(bytes, 64).any(recognises_pdbqt_line)
}

fn recognises_pdbqt_line(line: &[u8]) -> bool {
    if matches!(line, b"ROOT" | b"ENDROOT") || line.starts_with(b"TORSDOF") {
        return true;
    }
    if !(line.starts_with(b"ATOM  ") || line.starts_with(b"HETATM")) {
        return false;
    }
    let charge = line
        .get(70..76)
        .and_then(|field| str::from_utf8(field).ok());
    let atom_type = line.get(77..).and_then(|field| str::from_utf8(field).ok());
    charge.is_some_and(|field| field.trim().parse::<f64>().is_ok())
        && atom_type.is_some_and(|field| !field.trim().is_empty())
}

fn recognises_sdf(bytes: &[u8]) -> bool {
    let mut lines = first_lines(bytes, 64);
    if lines.any(|line| line == b"$$$$") {
        return true;
    }
    first_lines(bytes, 4)
        .nth(3)
        .is_some_and(|line| line.ends_with(b"V2000") || line.ends_with(b"V3000"))
}

fn recognises_mol2(bytes: &[u8]) -> bool {
    first_lines(bytes, 64).any(|line| line.starts_with(b"@<TRIPOS>MOLECULE"))
}

/// A core-CIF block: fractional atom sites and no Cartesian ones.
fn recognises_small_cif(bytes: &[u8]) -> bool {
    let mut data_block = false;
    let mut fractional = false;
    for line in first_lines(bytes, 256) {
        data_block |= line.starts_with(b"data_");
        fractional |=
            line.starts_with(b"_atom_site_fract_x") || line.starts_with(b"_atom_site.fract_x");
        if line.starts_with(b"_atom_site.Cartn_x") {
            return false;
        }
    }
    data_block && fractional
}

fn first_lines(bytes: &[u8], count: usize) -> impl Iterator<Item = &[u8]> {
    bytes
        .split(|byte| *byte == b'\n')
        .take(count)
        .map(strip_carriage_return)
}

fn strip_carriage_return(line: &[u8]) -> &[u8] {
    if line.last().copied() == Some(b'\r') {
        &line[..line.len() - 1]
    } else {
        line
    }
}