pixelflow-core 0.1.0

//! Y4M stream writer for Phase 1 compatible formats.

use std::io::Write;

use crate::{
    ChromaSubsampling, ErrorCategory, ErrorCode, FormatDescriptor, FormatFamily, Frame,
    PixelFlowError, Rational, Result, SampleType,
};

/// Streaming Y4M writer for one fixed output clip.
pub struct Y4mWriter<W: Write> {
    writer: W,
    format: FormatDescriptor,
    width: usize,
    height: usize,
}

impl<W: Write> Y4mWriter<W> {
    /// Creates writer and emits stream header immediately.
    pub fn new(
        mut writer: W,
        format: FormatDescriptor,
        width: usize,
        height: usize,
        frame_rate: Rational,
    ) -> Result<Self> {
        let chroma = y4m_chroma_tag(&format)?;
        writeln!(
            writer,
            "YUV4MPEG2 W{width} H{height} F{}:{} Ip A0:0 {chroma}",
            frame_rate.numerator, frame_rate.denominator
        )
        .map_err(|error| write_error(&error))?;

        Ok(Self {
            writer,
            format,
            width,
            height,
        })
    }

    /// Writes one frame header and visible plane rows in storage order.
    pub fn write_frame(&mut self, frame: &Frame) -> Result<()> {
        if frame.format() != &self.format
            || frame.width() != self.width
            || frame.height() != self.height
        {
            return Err(PixelFlowError::new(
                ErrorCategory::Format,
                ErrorCode::new("format.y4m_frame_mismatch"),
                "frame properties do not match Y4M stream header",
            ));
        }

        self.writer
            .write_all(b"FRAME\n")
            .map_err(|error| write_error(&error))?;
        match self.format.sample_type() {
            SampleType::U8 => self.write_u8_planes(frame),
            SampleType::U16 => self.write_u16_planes(frame),
            SampleType::F32 => Err(unsupported_format(
                "format does not support Phase 1 Y4M float output",
            )),
        }
    }

    /// Consumes writer and returns wrapped output sink.
    #[must_use]
    pub fn into_inner(self) -> W {
        self.writer
    }

    fn write_u8_planes(&mut self, frame: &Frame) -> Result<()> {
        for plane_index in 0..self.format.planes().len() {
            let plane = frame.plane::<u8>(plane_index)?;
            for row in plane.rows() {
                self.writer
                    .write_all(row)
                    .map_err(|error| write_error(&error))?;
            }
        }

        Ok(())
    }

    fn write_u16_planes(&mut self, frame: &Frame) -> Result<()> {
        for plane_index in 0..self.format.planes().len() {
            let plane = frame.plane::<u16>(plane_index)?;
            for row in plane.rows() {
                for sample in row {
                    self.writer
                        .write_all(&sample.to_le_bytes())
                        .map_err(|error| write_error(&error))?;
                }
            }
        }

        Ok(())
    }
}

/// Returns Y4M chroma tag for one Phase 1 compatible format.
pub fn y4m_chroma_tag(format: &FormatDescriptor) -> Result<&'static str> {
    match (
        format.family(),
        format.subsampling(),
        format.bits_per_sample(),
        format.sample_type(),
    ) {
        (FormatFamily::Gray, None, 8, SampleType::U8) => Ok("Cmono"),
        (FormatFamily::Gray, None, 10, SampleType::U16) => Ok("Cmono10"),
        (FormatFamily::Gray, None, 12, SampleType::U16) => Ok("Cmono12"),
        (FormatFamily::Gray, None, 16, SampleType::U16) => Ok("Cmono16"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 8, SampleType::U8) => Ok("C420"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 8, SampleType::U8) => Ok("C422"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 8, SampleType::U8) => Ok("C444"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 10, SampleType::U16) => Ok("C420p10"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 10, SampleType::U16) => Ok("C422p10"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 10, SampleType::U16) => Ok("C444p10"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 12, SampleType::U16) => Ok("C420p12"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 12, SampleType::U16) => Ok("C422p12"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 12, SampleType::U16) => Ok("C444p12"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs420), 16, SampleType::U16) => Ok("C420p16"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs422), 16, SampleType::U16) => Ok("C422p16"),
        (FormatFamily::Yuv, Some(ChromaSubsampling::Cs444), 16, SampleType::U16) => Ok("C444p16"),
        _ => Err(unsupported_format(format!(
            "format '{}' cannot be written as Phase 1 Y4M",
            format.name()
        ))),
    }
}

fn unsupported_format(message: impl Into<String>) -> PixelFlowError {
    PixelFlowError::new(
        ErrorCategory::Format,
        ErrorCode::new("format.unsupported_y4m"),
        message,
    )
}

fn write_error(error: &std::io::Error) -> PixelFlowError {
    PixelFlowError::new(
        ErrorCategory::Io,
        ErrorCode::new("io.write_y4m"),
        format!("failed to write Y4M output: {error}"),
    )
}

#[cfg(test)]
mod tests {
    #![expect(clippy::indexing_slicing, reason = "allow in tests")]

    use crate::{AllocatorConfig, FrameBuilder, MetadataSchema, Rational, resolve_format_alias};

    use super::{Y4mWriter, y4m_chroma_tag};

    #[test]
    fn y4m_chroma_tag_maps_phase1_formats() {
        assert_eq!(
            y4m_chroma_tag(&resolve_format_alias("gray8").expect("gray8 format should resolve"))
                .expect("gray8 should map to Y4M"),
            "Cmono"
        );
        assert_eq!(
            y4m_chroma_tag(&resolve_format_alias("gray10").expect("gray10 format should resolve"))
                .expect("gray10 should map to Y4M"),
            "Cmono10"
        );
        assert_eq!(
            y4m_chroma_tag(
                &resolve_format_alias("yuv420p8").expect("yuv420p8 format should resolve")
            )
            .expect("yuv420p8 should map to Y4M"),
            "C420"
        );
        assert_eq!(
            y4m_chroma_tag(
                &resolve_format_alias("yuv422p10").expect("yuv422p10 format should resolve")
            )
            .expect("yuv422p10 should map to Y4M"),
            "C422p10"
        );
        assert_eq!(
            y4m_chroma_tag(
                &resolve_format_alias("yuv444p16").expect("yuv444p16 format should resolve")
            )
            .expect("yuv444p16 should map to Y4M"),
            "C444p16"
        );
        assert!(
            y4m_chroma_tag(&resolve_format_alias("rgbp8").expect("rgbp8 format should resolve"))
                .is_err()
        );
    }

    #[test]
    fn y4m_writer_writes_header_and_u8_planes() {
        let format = resolve_format_alias("gray8").expect("gray8 format should resolve");
        let mut builder = FrameBuilder::new(
            format.clone(),
            2,
            2,
            &MetadataSchema::core(),
            AllocatorConfig::default(),
        )
        .expect("frame builder should construct");
        {
            let mut plane = builder.plane_mut::<u8>(0).expect("plane should exist");
            plane
                .row_mut(0)
                .expect("row 0 should exist")
                .copy_from_slice(&[1, 2]);
            plane
                .row_mut(1)
                .expect("row 1 should exist")
                .copy_from_slice(&[3, 4]);
        }
        let frame = builder.finish();
        let mut output = Vec::new();
        let mut writer = Y4mWriter::new(
            &mut output,
            format,
            2,
            2,
            Rational {
                numerator: 24,
                denominator: 1,
            },
        )
        .expect("writer should construct");

        writer
            .write_frame(&frame)
            .expect("frame should write successfully");

        assert_eq!(
            output,
            b"YUV4MPEG2 W2 H2 F24:1 Ip A0:0 Cmono\nFRAME\n\x01\x02\x03\x04".to_vec()
        );
    }

    #[test]
    fn y4m_writer_writes_u16_samples_little_endian() {
        let format = resolve_format_alias("gray10").expect("gray10 format should resolve");
        let mut builder = FrameBuilder::new(
            format.clone(),
            1,
            1,
            &MetadataSchema::core(),
            AllocatorConfig::default(),
        )
        .expect("frame builder should construct");
        builder
            .plane_mut::<u16>(0)
            .expect("plane should exist")
            .row_mut(0)
            .expect("row 0 should exist")[0] = 0x0201;
        let frame = builder.finish();
        let mut output = Vec::new();
        let mut writer = Y4mWriter::new(
            &mut output,
            format,
            1,
            1,
            Rational {
                numerator: 1,
                denominator: 1,
            },
        )
        .expect("writer should construct");

        writer
            .write_frame(&frame)
            .expect("frame should write successfully");

        assert_eq!(
            output,
            b"YUV4MPEG2 W1 H1 F1:1 Ip A0:0 Cmono10\nFRAME\n\x01\x02".to_vec()
        );
    }
}