threecrate-io 0.8.0

//! STL (STereoLithography) file format support
//!
//! Supports both ASCII and binary STL variants for reading and writing
//! triangle meshes. STL stores each triangle as three independent vertices
//! plus a face normal; on read, identical vertices are deduplicated so the
//! resulting [`TriangleMesh`] has a shared vertex buffer.

use crate::{MeshReader, MeshWriter};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
use std::path::Path;
use threecrate_core::{Error, Point3f, Result, TriangleMesh, Vector3f};

const BINARY_HEADER_SIZE: usize = 80;
const BINARY_TRIANGLE_SIZE: usize = 50;

/// STL reader (registry handle).
pub struct StlReader;

/// STL writer (registry handle).
pub struct StlWriter;

/// Options for writing STL files.
#[derive(Debug, Clone)]
pub struct StlWriteOptions {
    /// If true, write a binary STL file; otherwise write ASCII.
    pub binary: bool,
    /// Solid name used in ASCII output (ignored for binary).
    pub solid_name: String,
    /// Header text for binary output (truncated/padded to 80 bytes).
    pub binary_header: String,
}

impl Default for StlWriteOptions {
    fn default() -> Self {
        Self {
            binary: true,
            solid_name: "threecrate".to_string(),
            binary_header: "Binary STL generated by ThreeCrate".to_string(),
        }
    }
}

/// Detect whether a file is binary or ASCII STL.
///
/// ASCII files begin with the literal token `solid`, but some binary writers
/// also begin their 80-byte header with `solid`. The reliable check is to
/// compare the file size to what the binary triangle count would imply.
fn is_binary_stl(path: &Path) -> Result<bool> {
    let mut file = File::open(path)?;
    let file_size = file.metadata()?.len();

    if file_size < (BINARY_HEADER_SIZE + 4) as u64 {
        return Ok(false);
    }

    file.seek(SeekFrom::Start(BINARY_HEADER_SIZE as u64))?;
    let triangle_count = file.read_u32::<LittleEndian>()? as u64;
    let expected_size = BINARY_HEADER_SIZE as u64 + 4 + triangle_count * BINARY_TRIANGLE_SIZE as u64;

    if expected_size == file_size {
        return Ok(true);
    }

    // Fall back to a textual sniff: if it starts with "solid" and contains
    // "facet" in the first chunk, treat as ASCII.
    file.seek(SeekFrom::Start(0))?;
    let mut buf = [0u8; 256];
    let n = file.read(&mut buf)?;
    let head = String::from_utf8_lossy(&buf[..n]);
    let head_trim = head.trim_start();
    if head_trim.starts_with("solid") && head.contains("facet") {
        return Ok(false);
    }

    Ok(true)
}

fn quantize(v: f32) -> i32 {
    // 1e-6 quantization grid for vertex deduplication.
    (v * 1_000_000.0).round() as i32
}

struct MeshBuilder {
    vertices: Vec<Point3f>,
    face_normals: Vec<Vector3f>,
    faces: Vec<[usize; 3]>,
    index: HashMap<(i32, i32, i32), usize>,
}

impl MeshBuilder {
    fn new() -> Self {
        Self {
            vertices: Vec::new(),
            face_normals: Vec::new(),
            faces: Vec::new(),
            index: HashMap::new(),
        }
    }

    fn add_vertex(&mut self, v: Point3f) -> usize {
        let key = (quantize(v.x), quantize(v.y), quantize(v.z));
        if let Some(&i) = self.index.get(&key) {
            return i;
        }
        let i = self.vertices.len();
        self.vertices.push(v);
        self.index.insert(key, i);
        i
    }

    fn add_triangle(&mut self, normal: Vector3f, a: Point3f, b: Point3f, c: Point3f) {
        let ia = self.add_vertex(a);
        let ib = self.add_vertex(b);
        let ic = self.add_vertex(c);
        if ia == ib || ib == ic || ia == ic {
            return; // degenerate
        }
        self.faces.push([ia, ib, ic]);
        self.face_normals.push(normal);
    }

    fn into_mesh(self) -> TriangleMesh {
        let MeshBuilder { vertices, face_normals, faces, .. } = self;
        let mut mesh = TriangleMesh::from_vertices_and_faces(vertices, faces);

        // Derive per-vertex normals by averaging adjacent face normals.
        if !mesh.faces.is_empty() && !face_normals.is_empty() {
            let mut acc = vec![Vector3f::new(0.0, 0.0, 0.0); mesh.vertices.len()];
            for (face, n) in mesh.faces.iter().zip(face_normals.iter()) {
                for &vi in face {
                    acc[vi] += *n;
                }
            }
            for v in &mut acc {
                let len = v.norm();
                if len > 0.0 {
                    *v /= len;
                }
            }
            mesh.set_normals(acc);
        }

        mesh
    }
}

/// Read an STL file (auto-detects ASCII vs. binary).
pub fn read_stl<P: AsRef<Path>>(path: P) -> Result<TriangleMesh> {
    let path = path.as_ref();
    if is_binary_stl(path)? {
        read_binary_stl(path)
    } else {
        read_ascii_stl(path)
    }
}

fn read_binary_stl(path: &Path) -> Result<TriangleMesh> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);

    let mut header = [0u8; BINARY_HEADER_SIZE];
    reader.read_exact(&mut header)?;
    let triangle_count = reader.read_u32::<LittleEndian>()?;

    let mut builder = MeshBuilder::new();
    for i in 0..triangle_count {
        let nx = reader.read_f32::<LittleEndian>()
            .map_err(|_| Error::InvalidData(format!("Failed to read normal of triangle {}", i)))?;
        let ny = reader.read_f32::<LittleEndian>()?;
        let nz = reader.read_f32::<LittleEndian>()?;

        let mut verts = [Point3f::new(0.0, 0.0, 0.0); 3];
        for v in &mut verts {
            let x = reader.read_f32::<LittleEndian>()?;
            let y = reader.read_f32::<LittleEndian>()?;
            let z = reader.read_f32::<LittleEndian>()?;
            *v = Point3f::new(x, y, z);
        }
        let _attr = reader.read_u16::<LittleEndian>()?;

        builder.add_triangle(Vector3f::new(nx, ny, nz), verts[0], verts[1], verts[2]);
    }

    Ok(builder.into_mesh())
}

fn parse_vec3(parts: &[&str], line_num: usize, what: &str) -> Result<(f32, f32, f32)> {
    if parts.len() < 3 {
        return Err(Error::InvalidData(format!(
            "Invalid {} at line {}: expected 3 numbers",
            what, line_num
        )));
    }
    let x = parts[0].parse::<f32>().map_err(|_| {
        Error::InvalidData(format!("Invalid {} x at line {}", what, line_num))
    })?;
    let y = parts[1].parse::<f32>().map_err(|_| {
        Error::InvalidData(format!("Invalid {} y at line {}", what, line_num))
    })?;
    let z = parts[2].parse::<f32>().map_err(|_| {
        Error::InvalidData(format!("Invalid {} z at line {}", what, line_num))
    })?;
    Ok((x, y, z))
}

fn read_ascii_stl(path: &Path) -> Result<TriangleMesh> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);

    let mut builder = MeshBuilder::new();
    let mut current_normal = Vector3f::new(0.0, 0.0, 0.0);
    let mut current_verts: Vec<Point3f> = Vec::with_capacity(3);

    for (line_num, line) in reader.lines().enumerate() {
        let line = line?;
        let trimmed = line.trim();
        if trimmed.is_empty() {
            continue;
        }
        let parts: Vec<&str> = trimmed.split_whitespace().collect();
        match parts[0] {
            "facet" => {
                // "facet normal nx ny nz"
                if parts.len() >= 5 && parts[1] == "normal" {
                    let (nx, ny, nz) = parse_vec3(&parts[2..5], line_num + 1, "normal")?;
                    current_normal = Vector3f::new(nx, ny, nz);
                } else {
                    current_normal = Vector3f::new(0.0, 0.0, 0.0);
                }
                current_verts.clear();
            }
            "vertex" => {
                if parts.len() < 4 {
                    return Err(Error::InvalidData(format!(
                        "Invalid vertex at line {}: expected 3 coordinates",
                        line_num + 1
                    )));
                }
                let (x, y, z) = parse_vec3(&parts[1..4], line_num + 1, "vertex")?;
                current_verts.push(Point3f::new(x, y, z));
            }
            "endfacet" => {
                if current_verts.len() == 3 {
                    builder.add_triangle(
                        current_normal,
                        current_verts[0],
                        current_verts[1],
                        current_verts[2],
                    );
                }
                current_verts.clear();
            }
            _ => {}
        }
    }

    Ok(builder.into_mesh())
}

/// Write an STL file using the given options.
pub fn write_stl<P: AsRef<Path>>(
    mesh: &TriangleMesh,
    path: P,
    options: &StlWriteOptions,
) -> Result<()> {
    if options.binary {
        write_binary_stl(mesh, path.as_ref(), options)
    } else {
        write_ascii_stl(mesh, path.as_ref(), options)
    }
}

fn compute_face_normal(a: Point3f, b: Point3f, c: Point3f) -> Vector3f {
    let u = b - a;
    let v = c - a;
    let n = u.cross(&v);
    let len = n.norm();
    if len > 0.0 {
        n / len
    } else {
        Vector3f::new(0.0, 0.0, 0.0)
    }
}

fn write_binary_stl(mesh: &TriangleMesh, path: &Path, options: &StlWriteOptions) -> Result<()> {
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    let mut header = [0u8; BINARY_HEADER_SIZE];
    let bytes = options.binary_header.as_bytes();
    let n = bytes.len().min(BINARY_HEADER_SIZE);
    header[..n].copy_from_slice(&bytes[..n]);
    writer.write_all(&header)?;
    writer.write_u32::<LittleEndian>(mesh.faces.len() as u32)?;

    for face in &mesh.faces {
        let a = mesh.vertices[face[0]];
        let b = mesh.vertices[face[1]];
        let c = mesh.vertices[face[2]];
        let normal = compute_face_normal(a, b, c);
        writer.write_f32::<LittleEndian>(normal.x)?;
        writer.write_f32::<LittleEndian>(normal.y)?;
        writer.write_f32::<LittleEndian>(normal.z)?;
        for v in [a, b, c] {
            writer.write_f32::<LittleEndian>(v.x)?;
            writer.write_f32::<LittleEndian>(v.y)?;
            writer.write_f32::<LittleEndian>(v.z)?;
        }
        writer.write_u16::<LittleEndian>(0)?;
    }

    writer.flush()?;
    Ok(())
}

fn write_ascii_stl(mesh: &TriangleMesh, path: &Path, options: &StlWriteOptions) -> Result<()> {
    let file = File::create(path)?;
    let mut writer = BufWriter::new(file);

    writeln!(writer, "solid {}", options.solid_name)?;
    for face in &mesh.faces {
        let a = mesh.vertices[face[0]];
        let b = mesh.vertices[face[1]];
        let c = mesh.vertices[face[2]];
        let n = compute_face_normal(a, b, c);
        writeln!(writer, "  facet normal {} {} {}", n.x, n.y, n.z)?;
        writeln!(writer, "    outer loop")?;
        writeln!(writer, "      vertex {} {} {}", a.x, a.y, a.z)?;
        writeln!(writer, "      vertex {} {} {}", b.x, b.y, b.z)?;
        writeln!(writer, "      vertex {} {} {}", c.x, c.y, c.z)?;
        writeln!(writer, "    endloop")?;
        writeln!(writer, "  endfacet")?;
    }
    writeln!(writer, "endsolid {}", options.solid_name)?;
    Ok(())
}

impl crate::registry::MeshReader for StlReader {
    fn read_mesh(&self, path: &Path) -> Result<TriangleMesh> {
        read_stl(path)
    }

    fn can_read(&self, path: &Path) -> bool {
        is_binary_stl(path).is_ok()
    }

    fn format_name(&self) -> &'static str {
        "stl"
    }
}

impl crate::registry::MeshWriter for StlWriter {
    fn write_mesh(&self, mesh: &TriangleMesh, path: &Path) -> Result<()> {
        write_stl(mesh, path, &StlWriteOptions::default())
    }

    fn format_name(&self) -> &'static str {
        "stl"
    }
}

impl MeshReader for StlReader {
    fn read_mesh<P: AsRef<Path>>(path: P) -> Result<TriangleMesh> {
        read_stl(path)
    }
}

impl MeshWriter for StlWriter {
    fn write_mesh<P: AsRef<Path>>(mesh: &TriangleMesh, path: P) -> Result<()> {
        write_stl(mesh, path, &StlWriteOptions::default())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    use threecrate_core::Point3f;

    fn sample_mesh() -> TriangleMesh {
        // Simple tetrahedron.
        let vertices = vec![
            Point3f::new(0.0, 0.0, 0.0),
            Point3f::new(1.0, 0.0, 0.0),
            Point3f::new(0.0, 1.0, 0.0),
            Point3f::new(0.0, 0.0, 1.0),
        ];
        let faces = vec![[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]];
        TriangleMesh::from_vertices_and_faces(vertices, faces)
    }

    #[test]
    fn binary_round_trip() {
        let mesh = sample_mesh();
        let dir = tempdir().unwrap();
        let path = dir.path().join("tet.stl");

        write_stl(&mesh, &path, &StlWriteOptions { binary: true, ..Default::default() }).unwrap();
        assert!(is_binary_stl(&path).unwrap());

        let read = read_stl(&path).unwrap();
        assert_eq!(read.vertex_count(), 4);
        assert_eq!(read.face_count(), 4);
    }

    #[test]
    fn ascii_round_trip() {
        let mesh = sample_mesh();
        let dir = tempdir().unwrap();
        let path = dir.path().join("tet_ascii.stl");

        write_stl(
            &mesh,
            &path,
            &StlWriteOptions { binary: false, ..Default::default() },
        )
        .unwrap();
        assert!(!is_binary_stl(&path).unwrap());

        let read = read_stl(&path).unwrap();
        assert_eq!(read.vertex_count(), 4);
        assert_eq!(read.face_count(), 4);
    }

    #[test]
    fn ascii_fixture_parses() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("tri.stl");
        let content = "\
solid foo
  facet normal 0 0 1
    outer loop
      vertex 0 0 0
      vertex 1 0 0
      vertex 0 1 0
    endloop
  endfacet
endsolid foo
";
        std::fs::write(&path, content).unwrap();
        let mesh = read_stl(&path).unwrap();
        assert_eq!(mesh.face_count(), 1);
        assert_eq!(mesh.vertex_count(), 3);
    }

    #[test]
    fn auto_dispatch_via_read_mesh() {
        let mesh = sample_mesh();
        let dir = tempdir().unwrap();
        let path = dir.path().join("auto.stl");
        crate::write_mesh(&mesh, &path).unwrap();
        let read = crate::read_mesh(&path).unwrap();
        assert_eq!(read.face_count(), 4);
    }
}