readcon-core 0.20.0

use crate::types::{
    ConFrame, SECTION_CHARGES, SECTION_DISPLACEMENTS, SECTION_ENERGIES, SECTION_FORCES,
    SECTION_MAGMOMS, SECTION_SPINS, SECTION_SPREADS, SECTION_VELOCITIES, encode_fixed_bitmask,
    meta,
};
use serde_json::{Value, json};
use std::collections::BTreeMap;
use std::fs::{File, OpenOptions};
use std::io::{self, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::Path;

/// Failure from [`ConFrameWriter`]. I/O stays distinct from format errors
/// so `InvalidUnits` and `InvalidCount` remain distinguishable (CON-error.1).
#[derive(Debug)]
pub enum WriteError {
    Io(io::Error),
    Format(crate::error::ParseError),
}

impl std::fmt::Display for WriteError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            WriteError::Io(e) => write!(f, "{e}"),
            WriteError::Format(e) => write!(f, "{e}"),
        }
    }
}

impl std::error::Error for WriteError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            WriteError::Io(e) => Some(e),
            WriteError::Format(e) => Some(e),
        }
    }
}

impl From<io::Error> for WriteError {
    fn from(e: io::Error) -> Self {
        WriteError::Io(e)
    }
}

impl From<WriteError> for io::Error {
    fn from(e: WriteError) -> Self {
        match e {
            WriteError::Io(e) => e,
            WriteError::Format(err) => io::Error::new(io::ErrorKind::InvalidData, err.to_string()),
        }
    }
}

/// Numeric text representation for CON frame fields.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[derive(Default)]
pub enum FloatFormat {
    /// A fixed number of digits after the decimal point.
    DecimalPlaces(usize),
    /// The shortest scientific representation preserving each finite f64.
    /// Signed zero is preserved; decimal precision does not limit this mode.
    /// CON-num.6 (D13): the default writer round-trips binary64.
    #[default]
    RoundTrip,
}


/// A writer that can serialize and write `ConFrame` objects to any output stream.
///
/// This struct encapsulates a writer (like a file) and provides a high-level API
/// for writing simulation frames in the `.con` format.
///
/// # Example
/// ```no_run
/// # use std::fs::File;
/// # use readcon_core::types::ConFrame;
/// # use readcon_core::writer::ConFrameWriter;
/// # let frames: Vec<ConFrame> = Vec::new();
/// let mut writer = ConFrameWriter::from_path("output.con").unwrap();
/// writer.extend(frames.iter()).unwrap();
/// ```
pub struct ConFrameWriter<W: Write> {
    writer: BufWriter<W>,
    float_format: FloatFormat,
    /// When true: sort metadata keys in JSON and emit sections in canonical
    /// order (velocities, forces, energies). Opt-in. Float text follows
    /// `float_format` (round-trip unless the caller asked for decimals).
    canonical: bool,
    /// Cache for the JSON metadata line: when consecutive frames share
    /// the same (spec_version, sections-set, metadata) triple the
    /// serialized JSON object is identical, so reusing the cached
    /// string skips the per-frame `serde_json::Map::insert` rebuild
    /// and re-serialisation. Hot for trajectory writes where every
    /// frame has the same `units` / `potential` / `validate` keys.
    metadata_cache: Option<MetadataCacheEntry>,
    /// One frame of CON text. Filled then flushed with a single `write_all`
    /// so the sink is not entered once per atom line.
    scratch: Vec<u8>,
    /// Reused shortest-round-trip buffer (ryu). One per writer, not per float.
    ryu: ryu::Buffer,
    /// Last x/y/z spellings. A value is formatted on the second sight and
    /// copied after that, which is the common lattice column.
    fmt_slot: [FmtSlot; 3],
}

/// One column of the round-trip float cache.
#[derive(Clone, Copy)]
struct FmtSlot {
    bits: u64,
    text: [u8; 24],
    len: u8,
    /// 0 empty, 1 seen once, 2 `text` holds the spelling.
    state: u8,
}

impl FmtSlot {
    const fn new() -> Self {
        Self {
            bits: 0,
            text: [0; 24],
            len: 0,
            state: 0,
        }
    }
}

#[derive(Debug)]
struct MetadataCacheEntry {
    /// Snapshot of the inputs that fully determine the serialized JSON
    /// metadata line. Cheap to clone; cheaper than re-serialising the
    /// whole map on every frame.
    spec_version: u32,
    has_velocities: bool,
    has_forces: bool,
    has_energies: bool,
    has_charges: bool,
    has_spins: bool,
    has_magmoms: bool,
    has_displacements: bool,
    has_spreads: bool,
    metadata: std::collections::BTreeMap<String, serde_json::Value>,
    /// Cached serialised metadata line (without trailing newline).
    serialized: String,
}

impl MetadataCacheEntry {
    fn matches(
        &self,
        spec_version: u32,
        sections: [bool; 8],
        metadata: &std::collections::BTreeMap<String, serde_json::Value>,
    ) -> bool {
        let [
            has_velocities,
            has_forces,
            has_energies,
            has_charges,
            has_spins,
            has_magmoms,
            has_displacements,
            has_spreads,
        ] = sections;
        self.spec_version == spec_version
            && self.has_velocities == has_velocities
            && self.has_forces == has_forces
            && self.has_energies == has_energies
            && self.has_charges == has_charges
            && self.has_spins == has_spins
            && self.has_magmoms == has_magmoms
            && self.has_displacements == has_displacements
            && self.has_spreads == has_spreads
            && self.metadata.len() == metadata.len()
            && self
                .metadata
                .iter()
                .zip(metadata)
                .all(|((ka, a), (kb, b))| ka == kb && metadata_value_eq(a, b))
    }
}

// Cache equality preserves the numeric representation, including signed zero,
// throughout nested JSON values.
fn metadata_value_eq(a: &serde_json::Value, b: &serde_json::Value) -> bool {
    use serde_json::Value;
    match (a, b) {
        (Value::Number(a), Value::Number(b)) if a.is_f64() && b.is_f64() => {
            a.as_f64().map(f64::to_bits) == b.as_f64().map(f64::to_bits)
        }
        (Value::Array(a), Value::Array(b)) => {
            a.len() == b.len() && a.iter().zip(b).all(|(a, b)| metadata_value_eq(a, b))
        }
        (Value::Object(a), Value::Object(b)) => {
            a.len() == b.len()
                && a.iter()
                    .all(|(key, a)| b.get(key).is_some_and(|b| metadata_value_eq(a, b)))
        }
        _ => a == b,
    }
}

/// CON-section.16. Minor 2 and later must list, in `section_arity`, every
/// emitted section that minor 1 does not reserve.
fn require_minor2_section_arity(
    minor: u64,
    sections: &[serde_json::Value],
    metadata: &serde_json::Map<String, serde_json::Value>,
) -> Result<(), WriteError> {
    if minor < 2 {
        return Ok(());
    }
    for name in sections.iter().filter_map(|value| value.as_str()) {
        if crate::types::reserved_section_value_arity(name).is_some() {
            continue;
        }
        let listed = metadata
            .get(meta::SECTION_ARITY)
            .and_then(|value| value.as_object())
            .and_then(|obj| obj.get(name))
            .and_then(|value| value.as_u64())
            .is_some_and(|n| n >= 1);
        if !listed {
            return Err(WriteError::Format(
                crate::error::ParseError::ValidationError(format!(
                    "minor >= 2 must list {name} in section_arity"
                )),
            ));
        }
    }
    Ok(())
}

// General implementation for any type that implements `Write`.
impl<W: Write> ConFrameWriter<W> {
    /// Creates a new `ConFrameWriter` that wraps a given writer.
    ///
    /// # Arguments
    ///
    /// * `writer` - Any type that implements `std::io::Write`, e.g., a `File`.
    pub fn new(writer: W) -> Self {
        Self::with_float_format(writer, FloatFormat::default())
    }

    /// Creates a new `ConFrameWriter` with a custom floating-point precision.
    ///
    /// # Arguments
    ///
    /// * `writer` - Any type that implements `std::io::Write`.
    /// * `precision` - Number of decimal places for floating-point output.
    pub fn with_precision(writer: W, precision: usize) -> Self {
        Self::with_float_format(writer, FloatFormat::DecimalPlaces(precision))
    }

    /// Creates a writer with an explicit numeric text representation.
    ///
    /// `FloatFormat::RoundTrip` retains small forces and coordinates exactly.
    /// Metadata floats use the JSON serializer independently of this setting.
    pub fn with_float_format(writer: W, float_format: FloatFormat) -> Self {
        Self {
            writer: BufWriter::new(writer),
            float_format,
            canonical: false,
            metadata_cache: None,
            scratch: Vec::with_capacity(16 * 1024),
            ryu: ryu::Buffer::new(),
            fmt_slot: [FmtSlot::new(); 3],
        }
    }

    /// Opt-in **canonical** serialization: BTree-ordered metadata keys in JSON,
    /// fixed section order. Use for corpus
    /// materialization and content-stable hashes; not required for on-disk fidelity
    /// of spans preserved from `next_with_raw_span`.
    pub fn canonical(mut self, on: bool) -> Self {
        self.set_canonical(on);
        self
    }

    /// Set or clear canonical mode on an existing writer (C ABI / FFI).
    pub fn set_canonical(&mut self, on: bool) {
        self.canonical = on;
        if on {
            self.metadata_cache = None;
        }
    }

    /// Whether canonical serialization is enabled.
    pub fn is_canonical(&self) -> bool {
        self.canonical
    }

    fn refresh_metadata_cache(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
        let spec_version = frame.header.spec_version;
        let has_vel = frame.has_velocities();
        let has_frc = frame.has_forces();
        let has_eng = frame.has_energies();
        let has_chg = frame.has_charges();
        let has_spn = frame.has_spins();
        let has_mm = frame.has_magmoms();
        let has_dsp = frame.has_displacements();
        let has_spr = frame.has_spreads();

        let cache_hit = !self.canonical
            && self.metadata_cache.as_ref().is_some_and(|c| {
                c.matches(
                    spec_version,
                    [
                        has_vel, has_frc, has_eng, has_chg, has_spn, has_mm, has_dsp, has_spr,
                    ],
                    &frame.header.metadata,
                )
            });
        if cache_hit {
            return Ok(());
        }

        let mut meta_obj = serde_json::Map::new();
        meta_obj.insert(meta::CON_SPEC_VERSION.into(), json!(spec_version));
        let mut sections = Vec::new();
        if has_vel {
            sections.push(json!(SECTION_VELOCITIES));
        }
        if has_frc {
            sections.push(json!(SECTION_FORCES));
        }
        if has_eng {
            sections.push(json!(SECTION_ENERGIES));
        }
        if has_chg {
            sections.push(json!(SECTION_CHARGES));
        }
        if has_spn {
            sections.push(json!(SECTION_SPINS));
        }
        if has_mm {
            sections.push(json!(SECTION_MAGMOMS));
        }
        if has_dsp {
            sections.push(json!(SECTION_DISPLACEMENTS));
        }
        if has_spr {
            sections.push(json!(SECTION_SPREADS));
        }
        let validate = frame
            .header
            .metadata
            .get(meta::VALIDATE)
            .and_then(|value| value.as_bool())
            .unwrap_or(false);
        if !sections.is_empty() || validate {
            meta_obj.insert(meta::SECTIONS.into(), json!(sections));
        }
        for (k, v) in &frame.header.metadata {
            if k == meta::CON_SPEC_VERSION || k == meta::SECTIONS {
                continue;
            }
            meta_obj.insert(k.clone(), v.clone());
        }
        if let Some(u) = meta_obj.get(meta::UNITS).cloned() {
            // D11: an invalid unit fails the write. Do not substitute a default.
            let c = crate::units::canonicalize_units_object(&u).map_err(|e| {
                WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
            })?;
            meta_obj.insert(meta::UNITS.into(), c);
        }
        if spec_version >= 3 {
            match meta_obj.get(meta::UNITS) {
                None => {
                    return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
                        "con_spec_version >= 3 requires units.length and units.energy".into(),
                    )));
                }
                Some(u) => {
                    crate::units::validate_v3_units_metadata(u).map_err(|e| {
                        WriteError::Format(crate::error::ParseError::InvalidUnits(e.to_string()))
                    })?;
                }
            }
        }
        let minor = frame
            .header
            .metadata
            .get(meta::CON_SPEC_MINOR)
            .and_then(|v| v.as_u64())
            .unwrap_or(0);
        if spec_version >= 3 && minor >= 1 && has_vel {
            let time_ok = meta_obj
                .get(meta::UNITS)
                .and_then(|u| u.get("time"))
                .and_then(|v| v.as_str())
                .is_some_and(|s| !s.trim().is_empty());
            if !time_ok {
                return Err(WriteError::Format(crate::error::ParseError::InvalidUnits(
                    "minor >= 1 requires units.time when a velocities section is present".into(),
                )));
            }
        }
        // A minor of 2 or later lists every section minor 1 does not reserve
        // (CON-section.16). Reserved names need no entry.
        require_minor2_section_arity(minor, &sections, &meta_obj)?;
        let serialized = serde_json::Value::Object(meta_obj).to_string();
        self.metadata_cache = Some(MetadataCacheEntry {
            spec_version,
            has_velocities: has_vel,
            has_forces: has_frc,
            has_energies: has_eng,
            has_charges: has_chg,
            has_spins: has_spn,
            has_magmoms: has_mm,
            has_displacements: has_dsp,
            has_spreads: has_spr,
            metadata: frame.header.metadata.clone(),
            serialized,
        });
        Ok(())
    }

    /// Writes `frame` with the binding metadata stamp and the original atoms.
    ///
    /// The stamp is installed for the duration of the write and then put back.
    /// Atoms are not rebuilt.
    pub fn write_frame_binding(&mut self, frame: &mut ConFrame) -> Result<(), WriteError> {
        let stamped = binding_metadata(frame);
        let saved_meta = std::mem::replace(&mut frame.header.metadata, stamped);
        let saved_spec = frame.header.spec_version;
        if frame.header.spec_version < crate::CON_SPEC_VERSION {
            frame.header.spec_version = crate::CON_SPEC_VERSION;
        }
        let result = self.write_frame(frame);
        frame.header.metadata = saved_meta;
        frame.header.spec_version = saved_spec;
        result
    }

    /// Writes a single `ConFrame` to the output stream.
    pub fn write_frame(&mut self, frame: &ConFrame) -> Result<(), WriteError> {
        let zero_frame = frame.header.natm_types == 0
            && frame.atom_data.is_empty()
            && frame.header.natms_per_type.is_empty()
            && frame.header.masses_per_type.is_empty();
        if zero_frame && !frame.header.sections.is_empty() {
            return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
                "a zero-atom frame cannot declare sections".into(),
            )));
        }
        if !zero_frame
            && (frame.header.natm_types == 0
                || frame.atom_data.is_empty()
                || frame.header.natms_per_type.contains(&0))
        {
            return Err(WriteError::Format(crate::error::ParseError::InvalidCount(
                "each per-type count must be >= 1".into(),
            )));
        }
        ensure_unique_ids(&frame.atom_data)?;
        let prec = self.float_format;
        if frame.header.spec_version >= 2 {
            self.refresh_metadata_cache(frame)?;
        }
        let (boxl, angles) = match frame.header.lattice_vectors() {
            Some(vecs) => crate::types::lengths_and_angles(vecs),
            None => (frame.header.boxl, frame.header.angles),
        };
        let atoms = frame.atom_data.len();
        let mut blocks = 1usize;
        if frame.has_velocities() {
            blocks += 1;
        }
        if frame.has_forces() {
            blocks += 1;
        }
        if frame.has_energies() {
            blocks += 1;
        }
        if frame.has_charges() {
            blocks += 1;
        }
        if frame.has_spins() {
            blocks += 1;
        }
        if frame.has_magmoms() {
            blocks += 1;
        }
        if frame.has_displacements() {
            blocks += 1;
        }
        if frame.has_spreads() {
            blocks += 1;
        }
        self.scratch
            .reserve(256 + atoms.saturating_mul(64).saturating_mul(blocks));
        self.scratch.clear();
        push_str_line(&mut self.scratch, &frame.header.prebox_header.user);
        if frame.header.spec_version < 2 {
            // Spec 1 line 2 is free text. Do not re-emit JSON major 0 or 1 (D15).
            push_str_line(
                &mut self.scratch,
                frame.header.prebox_header.metadata_line(),
            );
        } else {
            let line = self
                .metadata_cache
                .as_ref()
                .expect("metadata_cache populated above")
                .serialized
                .as_bytes();
            self.scratch.extend_from_slice(line);
            self.scratch.push(b'\n');
        }
        {
            let buf = &mut self.scratch;
            let ryu = &mut self.ryu;
            let slots = &mut self.fmt_slot;

            // --- Write the 9-line Header ---
            push_f64(buf, ryu, boxl[0], prec);
            buf.push(b' ');
            push_f64(buf, ryu, boxl[1], prec);
            buf.push(b' ');
            push_f64(buf, ryu, boxl[2], prec);
            buf.push(b'\n');
            push_f64(buf, ryu, angles[0], prec);
            buf.push(b' ');
            push_f64(buf, ryu, angles[1], prec);
            buf.push(b' ');
            push_f64(buf, ryu, angles[2], prec);
            buf.push(b'\n');
            push_str_line(buf, &frame.header.postbox_header[0]);
            push_str_line(buf, &frame.header.postbox_header[1]);
            push_u64(buf, frame.header.natm_types as u64);
            buf.push(b'\n');

            for (i, n) in frame.header.natms_per_type.iter().enumerate() {
                if i > 0 {
                    buf.push(b' ');
                }
                push_u64(buf, *n as u64);
            }
            buf.push(b'\n');

            for (i, m) in frame.header.masses_per_type.iter().enumerate() {
                if i > 0 {
                    buf.push(b' ');
                }
                push_f64(buf, ryu, *m, prec);
            }
            buf.push(b'\n');

            // --- Write the Atom Data ---
            let mut atom_idx_offset = 0;
            for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate() {
                let symbol = &frame.atom_data[atom_idx_offset].symbol;
                push_str_line(buf, symbol);
                push_component_line(buf, "Coordinates of Component", type_idx);

                for i in 0..num_atoms_in_type {
                    let atom = &frame.atom_data[atom_idx_offset + i];
                    push_xyz_line(
                        buf,
                        ryu,
                        slots,
                        [atom.x, atom.y, atom.z],
                        prec,
                        written_fixed_flag(atom.fixed, frame.header.spec_version),
                        atom.atom_id,
                    );
                }
                atom_idx_offset += num_atoms_in_type;
            }

            // --- Write optional velocity section ---
            if frame.has_velocities() {
                buf.push(b'\n');

                let mut vel_idx_offset = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[vel_idx_offset].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Velocities of Component", type_idx);

                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[vel_idx_offset + i];
                        let [vx, vy, vz] = atom.velocity().unwrap_or([0.0; 3]);
                        push_xyz_line(
                            buf,
                            ryu,
                            slots,
                            [vx, vy, vz],
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    vel_idx_offset += num_atoms_in_type;
                }
            }

            // --- Write optional force section ---
            if frame.has_forces() {
                buf.push(b'\n');

                let mut force_idx_offset = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[force_idx_offset].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Forces of Component", type_idx);

                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[force_idx_offset + i];
                        let [fx, fy, fz] = atom.force().unwrap_or([0.0; 3]);
                        push_xyz_line(
                            buf,
                            ryu,
                            slots,
                            [fx, fy, fz],
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    force_idx_offset += num_atoms_in_type;
                }
            }

            // --- Write optional energies section ---
            if frame.has_energies() {
                buf.push(b'\n');

                let mut energy_idx_offset = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[energy_idx_offset].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Energies of Component", type_idx);

                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[energy_idx_offset + i];
                        let e = atom.energy().unwrap_or(0.0);
                        push_scalar_line(
                            buf,
                            ryu,
                            slots,
                            e,
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    energy_idx_offset += num_atoms_in_type;
                }
            }

            if frame.has_charges() {
                buf.push(b'\n');
                let mut off = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[off].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Charges of Component", type_idx);
                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[off + i];
                        let q = atom.charge().unwrap_or(0.0);
                        push_scalar_line(
                            buf,
                            ryu,
                            slots,
                            q,
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    off += num_atoms_in_type;
                }
            }

            if frame.has_spins() {
                buf.push(b'\n');
                let mut off = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[off].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Spins of Component", type_idx);
                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[off + i];
                        let s = atom.spin().unwrap_or(0.0);
                        push_scalar_line(
                            buf,
                            ryu,
                            slots,
                            s,
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    off += num_atoms_in_type;
                }
            }

            if frame.has_magmoms() {
                buf.push(b'\n');
                let mut off = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[off].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Magmoms of Component", type_idx);
                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[off + i];
                        let [mx, my, mz] = atom.magmom().unwrap_or([0.0; 3]);
                        push_xyz_line(
                            buf,
                            ryu,
                            slots,
                            [mx, my, mz],
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    off += num_atoms_in_type;
                }
            }

            if frame.has_displacements() {
                buf.push(b'\n');
                let mut off = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[off].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Displacements of Component", type_idx);
                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[off + i];
                        let [dx, dy, dz] = atom.displacement().unwrap_or([0.0; 3]);
                        push_xyz_line(
                            buf,
                            ryu,
                            slots,
                            [dx, dy, dz],
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    off += num_atoms_in_type;
                }
            }

            if frame.has_spreads() {
                buf.push(b'\n');
                let mut off = 0;
                for (type_idx, &num_atoms_in_type) in frame.header.natms_per_type.iter().enumerate()
                {
                    let symbol = &frame.atom_data[off].symbol;
                    push_str_line(buf, symbol);
                    push_component_line(buf, "Spreads of Component", type_idx);
                    for i in 0..num_atoms_in_type {
                        let atom = &frame.atom_data[off + i];
                        let [sx, sy, sz] = atom.spread().unwrap_or([0.0; 3]);
                        if !crate::types::spread_row_ok([sx, sy, sz]) {
                            let err = if sx.is_finite() && sy.is_finite() && sz.is_finite() {
                                crate::types::invalid_spread(off + i)
                            } else {
                                crate::error::ParseError::NonFiniteValue(format!(
                                    "spreads atom {}",
                                    off + i
                                ))
                            };
                            return Err(WriteError::Format(err));
                        }
                        push_xyz_line(
                            buf,
                            ryu,
                            slots,
                            [sx, sy, sz],
                            prec,
                            written_fixed_flag(atom.fixed, frame.header.spec_version),
                            atom.atom_id,
                        );
                    }
                    off += num_atoms_in_type;
                }
            }
        }

        self.writer.write_all(&self.scratch)?;
        Ok(())
    }

    /// Flushes buffered frames to the underlying writer.
    ///
    /// `write_frame` fills a buffer and does not always reach the sink.
    /// After `flush` returns, a reader sees a complete, parseable prefix.
    /// Dropping the writer flushes as well.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }

    /// Writes all frames from an iterator to the output stream.
    ///
    /// This is the most convenient way to write a multi-frame file.
    pub fn extend<'a>(
        &mut self,
        frames: impl Iterator<Item = &'a ConFrame>,
    ) -> Result<(), WriteError> {
        for frame in frames {
            self.write_frame(frame)?;
        }
        Ok(())
    }
}

fn ensure_unique_ids(atoms: &[crate::types::AtomDatum]) -> Result<(), WriteError> {
    let n = atoms.len();
    if n == 0 {
        return Ok(());
    }
    // Trajectory frames usually number atoms 0..n-1. That needs no set.
    if atoms
        .iter()
        .enumerate()
        .all(|(i, atom)| atom.atom_id == i as u64)
    {
        return Ok(());
    }
    let mut seen_ids = std::collections::HashSet::with_capacity(n);
    for atom in atoms {
        if !seen_ids.insert(atom.atom_id) {
            return Err(WriteError::Format(
                crate::error::ParseError::ValidationError(format!(
                    "atom id {} is repeated in the frame",
                    atom.atom_id
                )),
            ));
        }
    }
    Ok(())
}

fn written_fixed_flag(fixed: [bool; 3], spec_version: u32) -> u8 {
    if spec_version < 2 {
        if fixed.iter().any(|axis| *axis) { 1 } else { 0 }
    } else {
        encode_fixed_bitmask(fixed)
    }
}

fn push_u64(buf: &mut Vec<u8>, n: u64) {
    // Two digits per division. u64 stays in a hardware remainder; the old
    // path widened every atom id to u128 and called the software divider.
    const DIGITS: &[u8; 200] = b"00010203040506070809101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899";
    let mut tmp = [0u8; 20];
    let mut i = 20;
    let mut n = n;
    while n >= 100 {
        let q = n / 100;
        let r = (n % 100) as usize;
        i -= 2;
        tmp[i] = DIGITS[r * 2];
        tmp[i + 1] = DIGITS[r * 2 + 1];
        n = q;
    }
    if n < 10 {
        i -= 1;
        tmp[i] = b'0' + n as u8;
    } else {
        let r = n as usize;
        i -= 2;
        tmp[i] = DIGITS[r * 2];
        tmp[i + 1] = DIGITS[r * 2 + 1];
    }
    buf.extend_from_slice(&tmp[i..]);
}

fn push_u128(buf: &mut Vec<u8>, mut n: u128) {
    if n <= u128::from(u64::MAX) {
        push_u64(buf, n as u64);
        return;
    }
    let mut tmp = [0u8; 40];
    let mut i = 40;
    while n > 0 {
        i -= 1;
        tmp[i] = b'0' + (n % 10) as u8;
        n /= 10;
    }
    buf.extend_from_slice(&tmp[i..]);
}

fn push_str_line(buf: &mut Vec<u8>, text: &str) {
    buf.extend_from_slice(text.as_bytes());
    buf.push(b'\n');
}

fn push_component_line(buf: &mut Vec<u8>, label: &str, index: usize) {
    buf.extend_from_slice(label.as_bytes());
    buf.push(b' ');
    push_u64(buf, (index + 1) as u64);
    buf.push(b'\n');
}

fn push_f64(buf: &mut Vec<u8>, ryu: &mut ryu::Buffer, value: f64, format: FloatFormat) {
    match format {
        FloatFormat::DecimalPlaces(precision) => push_f64_prec(buf, value, precision),
        FloatFormat::RoundTrip => push_round_trip(buf, ryu, None, value),
    }
}

/// Shortest round-trip spelling. Finite values are written with `ryu::raw`
/// straight into `buf`. `slot` remembers a column so a repeated lattice
/// value is copied instead of formatted again.
fn push_round_trip(
    buf: &mut Vec<u8>,
    ryu: &mut ryu::Buffer,
    slot: Option<&mut FmtSlot>,
    value: f64,
) {
    if !value.is_finite() {
        let text = ryu.format(value);
        buf.extend_from_slice(text.as_bytes());
        return;
    }
    if let Some(slot) = slot {
        let bits = value.to_bits();
        if slot.bits == bits {
            if slot.state == 2 {
                buf.extend_from_slice(&slot.text[..slot.len as usize]);
                return;
            }
            if slot.state == 1 {
                let n = unsafe { ryu::raw::format64(value, slot.text.as_mut_ptr()) };
                debug_assert!(n <= 24);
                slot.len = n as u8;
                slot.state = 2;
                buf.extend_from_slice(&slot.text[..n]);
                return;
            }
        }
        write_finite(buf, value);
        slot.bits = bits;
        slot.state = 1;
        return;
    }
    write_finite(buf, value);
}

fn write_finite(buf: &mut Vec<u8>, value: f64) {
    buf.reserve(24);
    // `format64` writes at most 24 bytes and does not handle non-finite values.
    let n = unsafe { ryu::raw::format64(value, buf.as_mut_ptr().add(buf.len())) };
    debug_assert!(n <= 24);
    unsafe { buf.set_len(buf.len() + n) };
}

/// Fixed-point `f64` with `prec` digits after the decimal, matching `{:.prec$}`.
/// Non-finite values and `prec > 17` fall back to `std::fmt`.
fn push_f64_prec(buf: &mut Vec<u8>, v: f64, prec: usize) {
    // The six-decimal path uses fixed-point arithmetic; other precisions
    // use the standard formatter.
    if !v.is_finite() || prec != 6 {
        let _ = write!(buf, "{v:.prec$}");
        return;
    }
    if v.is_sign_negative() {
        buf.push(b'-');
    }
    let ax = v.abs();
    if prec == 0 {
        push_u128(buf, ax.round() as u128);
        return;
    }
    let scale = 10u128.pow(prec as u32);
    let n = (ax * scale as f64).round() as u128;
    let int_part = n / scale;
    let frac = n % scale;
    push_u128(buf, int_part);
    buf.push(b'.');
    let mut tmp = [b'0'; 20];
    let mut x = frac;
    let mut i = prec;
    while i > 0 {
        i -= 1;
        tmp[i] = b'0' + (x % 10) as u8;
        x /= 10;
    }
    buf.extend_from_slice(&tmp[..prec]);
}

fn push_xyz_line(
    buf: &mut Vec<u8>,
    ryu: &mut ryu::Buffer,
    slots: &mut [FmtSlot; 3],
    xyz: [f64; 3],
    prec: FloatFormat,
    fixed: u8,
    atom_id: u64,
) {
    let [x, y, z] = xyz;
    match prec {
        FloatFormat::RoundTrip => {
            push_round_trip(buf, ryu, Some(&mut slots[0]), x);
            buf.push(b' ');
            push_round_trip(buf, ryu, Some(&mut slots[1]), y);
            buf.push(b' ');
            push_round_trip(buf, ryu, Some(&mut slots[2]), z);
        }
        FloatFormat::DecimalPlaces(_) => {
            push_f64(buf, ryu, x, prec);
            buf.push(b' ');
            push_f64(buf, ryu, y, prec);
            buf.push(b' ');
            push_f64(buf, ryu, z, prec);
        }
    }
    buf.push(b' ');
    // Column 4 is a mask in 0..=7, one digit.
    buf.push(b'0' + fixed);
    buf.push(b' ');
    push_u64(buf, atom_id);
    buf.push(b'\n');
}

fn push_scalar_line(
    buf: &mut Vec<u8>,
    ryu: &mut ryu::Buffer,
    slots: &mut [FmtSlot; 3],
    v: f64,
    prec: FloatFormat,
    fixed: u8,
    atom_id: u64,
) {
    match prec {
        FloatFormat::RoundTrip => push_round_trip(buf, ryu, Some(&mut slots[0]), v),
        FloatFormat::DecimalPlaces(_) => push_f64(buf, ryu, v, prec),
    }
    buf.push(b' ');
    push_u64(buf, u64::from(fixed));
    buf.push(b' ');
    push_u64(buf, atom_id);
    buf.push(b'\n');
}

/// Record `units.time` for a v1 or v2 frame that carries velocities.
///
/// Those velocities are already in [`crate::units::NATURAL_TIME_UNIT`]
/// (CON-units.8). [`ConFrameBuilder::build`] stamps minor 1, and a minor
/// of 1 or later rejects velocities with no `units.time` (CON-units.11).
/// A 3.0 frame is treated the same way, overwriting any `units.time`
/// label. A 3.1 frame is left alone so a caller who omitted time still fails.
pub(crate) fn stamp_legacy_velocity_time(
    meta: &mut BTreeMap<String, Value>,
    spec_version: u32,
    has_velocity: bool,
) {
    if !has_velocity {
        return;
    }
    let minor = meta
        .get(meta::CON_SPEC_MINOR)
        .and_then(Value::as_u64)
        .unwrap_or(0);
    if spec_version >= 3 && minor >= 1 {
        return;
    }
    // Versions 1, 2 and 3.0 store velocities in the natural unit whatever
    // `units.time` says (CON-units.8, CON-units.9). Writing minor 1 makes
    // the label authoritative (CON-units.10), so replace any label rather
    // than keep one the old file never honoured.
    let mut units = match meta.get(meta::UNITS).cloned() {
        Some(Value::Object(map)) => Value::Object(map),
        Some(_) => return,
        None => crate::units::default_v3_units_json(),
    };
    if let Some(object) = units.as_object_mut() {
        object.insert(
            "time".to_string(),
            Value::String(crate::units::NATURAL_TIME_UNIT.to_string()),
        );
        meta.insert(meta::UNITS.to_string(), units);
        // The label now holds, so the frame is written as minor 1.
        if meta.contains_key(meta::CON_SPEC_MINOR) {
            meta.insert(meta::CON_SPEC_MINOR.to_string(), Value::from(1u64));
        }
    }
}

/// Metadata a language binding emits for `frame`, without copying atoms.
///
/// [`ConFrameBuilder::build`] stamps `con_spec_minor` 1 when the key is
/// absent. A parsed frame does not carry that key unless the file did.
/// The stamp lives on the header. The atoms stay where the parser put them.
pub fn binding_metadata(frame: &ConFrame) -> BTreeMap<String, Value> {
    let mut meta = frame.header.metadata.clone();
    meta.remove(meta::CON_SPEC_VERSION);
    meta.remove(meta::SECTIONS);
    let has_velocity = frame.atom_data.iter().any(|atom| atom.has_velocity());
    stamp_legacy_velocity_time(&mut meta, frame.header.spec_version, has_velocity);
    // The builder writes the current spec and fills units when the file
    // omitted them. A binding rewrite has to emit that same header.
    if !meta.contains_key(meta::UNITS) {
        meta.insert(
            meta::UNITS.to_string(),
            crate::units::default_v3_units_json(),
        );
    }
    if !meta.contains_key(meta::CON_SPEC_MINOR) {
        meta.insert(meta::CON_SPEC_MINOR.to_string(), Value::from(1u64));
    }
    meta
}

/// Same bytes as a binding rewrite, without rebuilding each atom.
pub fn binding_rewrite_frame(frame: &ConFrame) -> Result<ConFrame, crate::error::ParseError> {
    let mut out = frame.clone();
    out.header.metadata = binding_metadata(frame);
    // A parsed spec 1 frame keeps line 2 as free text. The builder emits
    // the current spec as JSON, so the rewrite does too.
    if out.header.spec_version < crate::CON_SPEC_VERSION {
        out.header.spec_version = crate::CON_SPEC_VERSION;
    }
    Ok(out)
}

// Implementation block specifically for when the writer is a `File`.
impl ConFrameWriter<File> {
    /// Creates a new `ConFrameWriter` that writes to a file at the given path.
    ///
    /// This is a convenience function that creates the file and wraps it.
    pub fn from_path<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let file = File::create(path)?;
        Ok(Self::new(file))
    }

    /// Creates a new `ConFrameWriter` that writes to a file with a custom precision.
    pub fn from_path_with_precision<P: AsRef<Path>>(path: P, precision: usize) -> io::Result<Self> {
        let file = File::create(path)?;
        Ok(Self::with_precision(file, precision))
    }

    /// Opens an uncompressed file for appending.
    ///
    /// The file is created when it does not exist. A non-empty file that does
    /// not end in a newline receives one, so the next frame header starts on
    /// its own line. A file this writer produced already ends in a newline,
    /// and appending the remaining frames reproduces one whole-file write.
    /// Gzip and zstd paths are refused: a compressed member cannot be extended.
    ///
    /// `write_frame` fills a one-frame scratch buffer and a [`std::io::BufWriter`].
    /// [`Self::flush`] pushes that buffer to the file before the writer is dropped.
    ///
    /// ```
    /// use readcon_core::iterators::read_all_frames;
    /// use readcon_core::types::ConFrameBuilder;
    /// use readcon_core::writer::ConFrameWriter;
    /// let path = std::env::temp_dir().join(format!(
    ///     "readcon-append-doc-{}",
    ///     std::process::id()
    /// ));
    /// let mut builder = ConFrameBuilder::new([10.0, 10.0, 10.0], [90.0, 90.0, 90.0]);
    /// builder.add_atom("H", 0.0, 0.0, 0.0, [false; 3], 1, 1.0);
    /// let frame = builder.build().unwrap();
    /// {
    ///     let mut writer = ConFrameWriter::from_path(&path).unwrap();
    ///     writer.write_frame(&frame).unwrap();
    ///     writer.flush().unwrap();
    /// }
    /// {
    ///     let mut writer = ConFrameWriter::from_path_append(&path).unwrap();
    ///     writer.write_frame(&frame).unwrap();
    ///     writer.flush().unwrap();
    /// }
    /// let frames = read_all_frames(&path).unwrap();
    /// assert_eq!(frames.len(), 2);
    /// let _ = std::fs::remove_file(&path);
    /// ```
    pub fn from_path_append<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        Ok(Self::new(open_uncompressed_append(path.as_ref())?))
    }

    /// Opens an uncompressed file for appending with a fixed decimal precision.
    pub fn from_path_append_with_precision<P: AsRef<Path>>(
        path: P,
        precision: usize,
    ) -> io::Result<Self> {
        Ok(Self::with_precision(
            open_uncompressed_append(path.as_ref())?,
            precision,
        ))
    }
}

/// Opens `path` at the end. See [`ConFrameWriter::from_path_append`].
pub(crate) fn open_uncompressed_append(path: &Path) -> io::Result<File> {
    if let Some(ext) = path.extension().and_then(|ext| ext.to_str())
        && (ext == "gz" || ext == "zst")
    {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "append is only supported for an uncompressed file",
        ));
    }
    let mut file = OpenOptions::new()
        .create(true)
        .read(true)
        .append(true)
        .open(path)?;
    let len = file.seek(SeekFrom::End(0))?;
    if len > 0 {
        file.seek(SeekFrom::End(-1))?;
        let mut last = [0u8; 1];
        file.read_exact(&mut last)?;
        if last[0] != b'\n' {
            file.write_all(b"\n")?;
        }
    }
    Ok(file)
}

/// Serializes `frames` to CON bytes. No file is created.
///
/// The bytes match [`ConFrameWriter::from_path`] writing the same frames.
pub fn write_frames_to_vec(frames: &[ConFrame]) -> Result<Vec<u8>, WriteError> {
    let mut buffer = Vec::new();
    {
        let mut writer = ConFrameWriter::new(&mut buffer);
        writer.extend(frames.iter())?;
    }
    Ok(buffer)
}

/// Serializes `frames` to a CON string. No file is created.
///
/// The bytes match [`ConFrameWriter::from_path`] writing the same frames.
pub fn write_frames_to_string(frames: &[ConFrame]) -> Result<String, WriteError> {
    let bytes = write_frames_to_vec(frames)?;
    String::from_utf8(bytes).map_err(|err| {
        WriteError::Io(io::Error::new(
            io::ErrorKind::InvalidData,
            err.utf8_error().to_string(),
        ))
    })
}

// Gzip-compressed writer constructors.
impl ConFrameWriter<flate2::write::GzEncoder<File>> {
    /// Creates a gzip-compressed writer for the given path.
    pub fn from_path_gzip<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let encoder = crate::compression::gzip_writer(path.as_ref())?;
        Ok(Self::new(encoder))
    }

    /// Creates a gzip-compressed writer with custom precision.
    pub fn from_path_gzip_with_precision<P: AsRef<Path>>(
        path: P,
        precision: usize,
    ) -> io::Result<Self> {
        let encoder = crate::compression::gzip_writer(path.as_ref())?;
        Ok(Self::with_precision(encoder, precision))
    }
}

// Zstd-compressed writer constructors. Available only with the `zstd`
// Cargo feature.
#[cfg(feature = "zstd")]
impl ConFrameWriter<zstd::stream::write::AutoFinishEncoder<'static, File>> {
    /// Creates a zstd-compressed writer for the given path.
    pub fn from_path_zstd<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let encoder = crate::compression::zstd_writer(path.as_ref())?;
        Ok(Self::new(encoder))
    }

    /// Creates a zstd-compressed writer with custom precision.
    pub fn from_path_zstd_with_precision<P: AsRef<Path>>(
        path: P,
        precision: usize,
    ) -> io::Result<Self> {
        let encoder = crate::compression::zstd_writer(path.as_ref())?;
        Ok(Self::with_precision(encoder, precision))
    }
}

#[cfg(test)]
mod float_format_tests {
    use super::{push_f64_prec, push_u64};

    fn formatted(v: f64, prec: usize) -> String {
        let mut buf = Vec::new();
        push_f64_prec(&mut buf, v, prec);
        String::from_utf8(buf).expect("utf8")
    }

    #[test]
    fn matches_std_fixed_precision() {
        let vals = [
            0.0,
            -0.0,
            1.0,
            -1.0,
            0.9045,
            6.975_299_999_999_995,
            63.546,
            1.008,
            15.3456,
            90.0,
            218.0,
            1e-6,
            -1.23456789,
            10.0,
        ];
        for prec in [0usize, 6] {
            for v in vals {
                let got = formatted(v, prec);
                let exp = format!("{v:.prec$}");
                assert_eq!(got, exp, "v={v:?} prec={prec}");
            }
        }
    }

    #[test]
    fn round_trip_spelling_matches_ryu_on_repeated_and_unique_rows() {
        use crate::iterators::ConFrameIterator;
        use crate::writer::ConFrameWriter;
        let mut text = String::from(
            "bench\n{\"con_spec_version\":2}\n10.0 10.0 10.0\n90.0 90.0 90.0\n0 0\n0 0 1\n1\n6\n12.0\nCu\nCoordinates of Component 1\n",
        );
        for i in 0..6 {
            let x = (i % 2) as f64 * 1.5 + 0.25;
            let y = 3.25;
            let z = if i == 3 {
                -0.0
            } else {
                (i as f64) * 0.0010000001
            };
            text.push_str(&format!("   {x:.17}    {y:.17}    {z:.17} 0 {i}\n"));
        }
        // Repeated velocities share the column cache with the coordinates.
        text.push_str("\nCu\nVelocities of Component 1\n");
        for i in 0..6 {
            let vx = 0.125;
            let vy = if i % 2 == 0 { 0.5 } else { -0.5 };
            let vz = (i as f64) * 1.0e-4;
            text.push_str(&format!("   {vx:.17}    {vy:.17}    {vz:.17} 0 {i}\n"));
        }
        let frames: Vec<_> = ConFrameIterator::new(&text)
            .map(|frame| frame.expect("parse"))
            .collect();
        let mut values = Vec::new();
        for atom in &frames[0].atom_data {
            values.extend([atom.x, atom.y, atom.z]);
        }
        for atom in &frames[0].atom_data {
            if let Some(v) = atom.velocity() {
                values.extend(v);
            }
        }
        let mut buf = Vec::new();
        {
            let mut writer = ConFrameWriter::new(&mut buf);
            writer.extend(frames.iter()).unwrap();
        }
        let written = String::from_utf8(buf).unwrap();
        let mut ryu = ryu::Buffer::new();
        let mut spelled = Vec::new();
        for line in written.lines() {
            let toks: Vec<_> = line.split_whitespace().collect();
            if toks.len() >= 5
                && toks[0].contains('.')
                && toks[3].bytes().all(|b| b.is_ascii_digit())
            {
                spelled.extend(toks[..3].iter().map(|tok| (*tok).to_string()));
            }
        }
        assert_eq!(spelled.len(), values.len());
        for (got, value) in spelled.iter().zip(values) {
            let expect = if value.is_finite() {
                ryu.format_finite(value).to_string()
            } else {
                ryu.format(value).to_string()
            };
            assert_eq!(got, &expect, "value {value:?}");
        }
    }

    #[test]
    fn u64_matches_decimal_display() {
        let mut values = vec![0, 1, 9, 10, 11, 99, 100, 101, 999, 1000, 50000, u64::MAX];
        values.extend(0..200);
        for n in values {
            let mut buf = Vec::new();
            push_u64(&mut buf, n);
            assert_eq!(std::str::from_utf8(&buf).unwrap(), n.to_string());
        }
    }
}

#[cfg(test)]
mod section_arity_tests {
    use super::require_minor2_section_arity;
    use serde_json::json;

    #[test]
    fn minor2_lists_a_new_section_in_section_arity() {
        let widgets = vec![json!("widgets")];
        let err = require_minor2_section_arity(2, &widgets, &serde_json::Map::new()).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("section_arity"), "{msg}");

        let mut listed = serde_json::Map::new();
        listed.insert("section_arity".into(), json!({"widgets": 3}));
        require_minor2_section_arity(2, &widgets, &listed).unwrap();

        let reserved = vec![json!("velocities"), json!("energies")];
        require_minor2_section_arity(2, &reserved, &serde_json::Map::new()).unwrap();
        require_minor2_section_arity(1, &widgets, &serde_json::Map::new()).unwrap();
    }
}