wblidar 0.1.1

//! PLY writer — ASCII and binary encodings.

use std::io::Write;
use crate::io::PointWriter;
use crate::ply::PlyEncoding;
use crate::point::PointRecord;
use crate::Result;

/// A PLY writer.  The header is written on construction; points are streamed
/// via `write_point`.
///
/// PLY requires the element count in the header *before* the data, so we
/// need to know the total point count up front.  If it is not known, write
/// to a `Vec<u8>` and then copy; alternatively use a seekable writer and
/// back-patch the header.
pub struct PlyWriter<W: Write> {
    inner: W,
    encoding: PlyEncoding,
    has_color: bool,
    has_normals: bool,
    points_written: u64,
}

impl<W: Write> PlyWriter<W> {
    /// Create the writer and write the PLY header immediately.
    ///
    /// # Parameters
    /// * `point_count` – total number of points that will be written.
    /// * `encoding`    – file encoding.
    /// * `has_color`   – include r/g/b properties.
    /// * `has_normals` – include nx/ny/nz properties.
    pub fn new(
        mut inner: W,
        point_count: u64,
        encoding: PlyEncoding,
        has_color: bool,
        has_normals: bool,
    ) -> Result<Self> {
        write_ply_header(&mut inner, point_count, encoding, has_color, has_normals)?;
        Ok(PlyWriter { inner, encoding, has_color, has_normals, points_written: 0 })
    }
}

impl<W: Write> PointWriter for PlyWriter<W> {
    fn write_point(&mut self, p: &PointRecord) -> Result<()> {
        match self.encoding {
            PlyEncoding::Ascii => write_ascii_point(&mut self.inner, p, self.has_color, self.has_normals)?,
            PlyEncoding::BinaryLittleEndian => {
                write_binary_point(&mut self.inner, p, self.has_color, self.has_normals, false)?;
            }
            PlyEncoding::BinaryBigEndian => {
                write_binary_point(&mut self.inner, p, self.has_color, self.has_normals, true)?;
            }
        }
        self.points_written += 1;
        Ok(())
    }

    fn finish(&mut self) -> Result<()> {
        Ok(self.inner.flush()?)
    }
}

// ── Header writer ─────────────────────────────────────────────────────────────

fn write_ply_header<W: Write>(
    w: &mut W, count: u64, enc: PlyEncoding, color: bool, normals: bool,
) -> Result<()> {
    let enc_str = match enc {
        PlyEncoding::Ascii => "ascii",
        PlyEncoding::BinaryLittleEndian => "binary_little_endian",
        PlyEncoding::BinaryBigEndian    => "binary_big_endian",
    };

    write!(w, "ply\nformat {enc_str} 1.0\ncomment generated by wblidar\n")?;
    write!(w, "element vertex {count}\n")?;
    write!(w, "property double x\nproperty double y\nproperty double z\n")?;
    write!(w, "property uint16 intensity\n")?;
    if color {
        write!(w, "property uchar red\nproperty uchar green\nproperty uchar blue\n")?;
    }
    if normals {
        write!(w, "property float nx\nproperty float ny\nproperty float nz\n")?;
    }
    write!(w, "property uint8 classification\n")?;
    write!(w, "end_header\n")?;
    Ok(())
}

// ── ASCII point writer ────────────────────────────────────────────────────────

fn write_ascii_point<W: Write>(
    w: &mut W, p: &PointRecord, color: bool, normals: bool,
) -> Result<()> {
    write!(w, "{} {} {} {}", p.x, p.y, p.z, p.intensity)?;
    if color {
        let c = p.color.map(|c| (c.red >> 8) as u8).unwrap_or(0);
        let g = p.color.map(|c| (c.green >> 8) as u8).unwrap_or(0);
        let b = p.color.map(|c| (c.blue >> 8) as u8).unwrap_or(0);
        write!(w, " {c} {g} {b}")?;
    }
    if normals {
        write!(w, " {} {} {}", p.normal_x.unwrap_or(0.0), p.normal_y.unwrap_or(0.0), p.normal_z.unwrap_or(0.0))?;
    }
    writeln!(w, " {}", p.classification)?;
    Ok(())
}

// ── Binary point writer ───────────────────────────────────────────────────────

fn write_binary_point<W: Write>(
    w: &mut W, p: &PointRecord, color: bool, normals: bool, big: bool,
) -> Result<()> {
    write_f64(w, p.x, big)?;
    write_f64(w, p.y, big)?;
    write_f64(w, p.z, big)?;
    write_u16(w, p.intensity, big)?;
    if color {
        let r = p.color.map(|c| (c.red   >> 8) as u8).unwrap_or(0);
        let g = p.color.map(|c| (c.green >> 8) as u8).unwrap_or(0);
        let b = p.color.map(|c| (c.blue  >> 8) as u8).unwrap_or(0);
        w.write_all(&[r, g, b])?;
    }
    if normals {
        write_f32(w, p.normal_x.unwrap_or(0.0), big)?;
        write_f32(w, p.normal_y.unwrap_or(0.0), big)?;
        write_f32(w, p.normal_z.unwrap_or(0.0), big)?;
    }
    w.write_all(&[p.classification])?;
    Ok(())
}

fn write_f64<W: Write>(w: &mut W, v: f64, big: bool) -> std::io::Result<()> {
    if big { w.write_all(&v.to_be_bytes()) } else { w.write_all(&v.to_le_bytes()) }
}
fn write_f32<W: Write>(w: &mut W, v: f32, big: bool) -> std::io::Result<()> {
    if big { w.write_all(&v.to_be_bytes()) } else { w.write_all(&v.to_le_bytes()) }
}
fn write_u16<W: Write>(w: &mut W, v: u16, big: bool) -> std::io::Result<()> {
    if big { w.write_all(&v.to_be_bytes()) } else { w.write_all(&v.to_le_bytes()) }
}