jpegli/encode/
byte_encoders.rs

1//! Encoder implementations for v2 API.
2
3use core::marker::PhantomData;
4
5#[cfg(feature = "std")]
6use std::io::Write;
7
8use enough::Stop;
9
10use super::encoder_config::EncoderConfig;
11use super::encoder_types::{PixelLayout, YCbCrPlanes};
12use super::extras::inject_encoder_segments;
13use super::streaming::StreamingEncoder;
14use crate::error::{Error, Result};
15
16/// Encoder for raw byte input with explicit pixel layout.
17///
18/// This encoder wraps `StreamingEncoder` to provide true streaming encoding
19/// without buffering the entire image in memory.
20pub struct BytesEncoder {
21    /// v2 config (kept for ICC profile injection)
22    config: EncoderConfig,
23    /// Pixel layout
24    layout: PixelLayout,
25    /// Image dimensions
26    width: u32,
27    height: u32,
28    /// Inner streaming encoder (handles actual encoding)
29    inner: StreamingEncoder,
30}
31
32impl BytesEncoder {
33    pub(crate) fn new(
34        config: EncoderConfig,
35        width: u32,
36        height: u32,
37        layout: PixelLayout,
38    ) -> Result<Self> {
39        // Validate dimensions
40        if width == 0 || height == 0 {
41            return Err(Error::invalid_dimensions(
42                width,
43                height,
44                "dimensions cannot be zero",
45            ));
46        }
47
48        // Check for overflow
49        let pixel_count = (width as u64) * (height as u64);
50        if pixel_count > u32::MAX as u64 {
51            return Err(Error::invalid_dimensions(
52                width,
53                height,
54                "dimensions too large",
55            ));
56        }
57
58        // Build and start the streaming encoder with config from v2
59        let inner = Self::build_streaming_encoder(&config, width, height, layout)?;
60
61        Ok(Self {
62            config,
63            layout,
64            width,
65            height,
66            inner,
67        })
68    }
69
70    /// Build a StreamingEncoder from v2 config.
71    fn build_streaming_encoder(
72        config: &EncoderConfig,
73        width: u32,
74        height: u32,
75        layout: PixelLayout,
76    ) -> Result<StreamingEncoder> {
77        use crate::encode::streaming::StreamingEncoder as SE;
78        use crate::types::PixelFormat;
79
80        let pixel_format: PixelFormat = layout.into();
81        let subsampling = match config.color_mode {
82            super::encoder_types::ColorMode::YCbCr { subsampling } => subsampling.into(),
83            super::encoder_types::ColorMode::Xyb { .. } => crate::types::Subsampling::S444,
84            super::encoder_types::ColorMode::Grayscale => crate::types::Subsampling::S444,
85        };
86
87        let mut builder = SE::new(width, height)
88            .quality(config.quality)
89            .pixel_format(pixel_format)
90            .subsampling(subsampling)
91            .optimize_huffman(config.optimize_huffman)
92            .chroma_downsampling(config.downsampling_method)
93            .restart_interval(config.restart_interval);
94
95        // Apply custom encoding tables if configured
96        if let Some(ref tables) = config.tables {
97            builder = builder.encoding_tables(tables.clone());
98        }
99
100        if config.progressive {
101            builder = builder.progressive(true);
102        }
103
104        if matches!(
105            config.color_mode,
106            super::encoder_types::ColorMode::Xyb { .. }
107        ) {
108            builder = builder.use_xyb(true);
109        }
110
111        // Always pass deringing setting (StreamingEncoder defaults to true)
112        builder = builder.deringing(config.deringing);
113
114        builder = builder.allow_16bit_quant_tables(config.allow_16bit_quant_tables);
115        builder = builder.separate_chroma_tables(config.separate_chroma_tables);
116
117        #[cfg(feature = "parallel")]
118        if config.parallel.is_some() {
119            // ParallelEncoding::Auto means enable parallel encoding
120            // Future variants may have different behaviors
121            builder = builder.parallel(true);
122        }
123
124        // Apply trellis quantization config if set
125        #[cfg(feature = "experimental-hybrid-trellis")]
126        if let Some(ref trellis) = config.trellis {
127            builder = builder.trellis(*trellis);
128        }
129
130        builder.start()
131    }
132
133    /// Push rows with explicit stride.
134    ///
135    /// - `data`: Raw pixel bytes
136    /// - `rows`: Number of scanlines to push
137    /// - `stride_bytes`: Bytes per row in buffer (>= width * bytes_per_pixel)
138    /// - `stop`: Cancellation token (use `enough::Unstoppable` if not needed)
139    pub fn push(
140        &mut self,
141        data: &[u8],
142        rows: usize,
143        stride_bytes: usize,
144        stop: impl Stop,
145    ) -> Result<()> {
146        // Check cancellation
147        if stop.should_stop() {
148            return Err(Error::cancelled());
149        }
150
151        let bpp = self.layout.bytes_per_pixel();
152        let min_stride = self.width as usize * bpp;
153
154        // Validate stride
155        if stride_bytes < min_stride {
156            return Err(Error::stride_too_small(self.width, stride_bytes));
157        }
158
159        // Validate row count
160        let current_rows = self.inner.rows_pushed() as u32;
161        let new_total = current_rows + rows as u32;
162        if new_total > self.height {
163            return Err(Error::too_many_rows(self.height, new_total));
164        }
165
166        // Validate buffer size
167        let expected_size = rows * stride_bytes;
168        if data.len() < expected_size {
169            return Err(Error::invalid_buffer_size(expected_size, data.len()));
170        }
171
172        // Push rows to streaming encoder
173        if stride_bytes == min_stride {
174            // Packed data - can push directly
175            self.inner
176                .push_rows_with_stop(&data[..rows * min_stride], rows, &stop)?;
177        } else {
178            // Strided data - push row by row
179            for row in 0..rows {
180                if stop.should_stop() {
181                    return Err(Error::cancelled());
182                }
183
184                let src_start = row * stride_bytes;
185                let src_end = src_start + min_stride;
186                self.inner
187                    .push_row_with_stop(&data[src_start..src_end], &stop)?;
188            }
189        }
190
191        Ok(())
192    }
193
194    /// Push contiguous (packed) data.
195    ///
196    /// Stride is assumed to be `width * bytes_per_pixel`.
197    /// Rows inferred from `data.len() / (width * bytes_per_pixel)`.
198    pub fn push_packed(&mut self, data: &[u8], stop: impl Stop) -> Result<()> {
199        let bpp = self.layout.bytes_per_pixel();
200        let row_bytes = self.width as usize * bpp;
201
202        if row_bytes == 0 {
203            return Err(Error::invalid_dimensions(
204                self.width,
205                self.height,
206                "row size is zero",
207            ));
208        }
209
210        let rows = data.len() / row_bytes;
211        if rows == 0 && !data.is_empty() {
212            return Err(Error::invalid_buffer_size(row_bytes, data.len()));
213        }
214
215        self.push(data, rows, row_bytes, stop)
216    }
217
218    // === Status ===
219
220    /// Get image width.
221    #[must_use]
222    pub fn width(&self) -> u32 {
223        self.width
224    }
225
226    /// Get image height.
227    #[must_use]
228    pub fn height(&self) -> u32 {
229        self.height
230    }
231
232    /// Get number of rows pushed so far.
233    #[must_use]
234    pub fn rows_pushed(&self) -> u32 {
235        self.inner.rows_pushed() as u32
236    }
237
238    /// Get number of rows remaining.
239    #[must_use]
240    pub fn rows_remaining(&self) -> u32 {
241        self.height - self.inner.rows_pushed() as u32
242    }
243
244    /// Get the pixel layout.
245    #[must_use]
246    pub fn layout(&self) -> PixelLayout {
247        self.layout
248    }
249
250    // === Finish ===
251
252    /// Finish encoding, return JPEG bytes.
253    pub fn finish(mut self) -> Result<Vec<u8>> {
254        let rows_pushed = self.inner.rows_pushed() as u32;
255        if rows_pushed != self.height {
256            return Err(Error::incomplete_image(self.height, rows_pushed));
257        }
258
259        // Finish streaming encoder
260        let mut jpeg = self.inner.finish()?;
261
262        // When using EncoderSegments with MPF images, we need to merge individual
263        // metadata fields (xmp_data, exif_data, icc_profile) into segments BEFORE
264        // calling inject_encoder_segments. This ensures the MPF offset calculation
265        // includes all metadata that will be in the final output.
266        if let Some(mut segments) = self.config.segments.take() {
267            // Merge individual metadata into segments if MPF is present
268            // (MPF offset calculation needs all data to be accounted for)
269            if segments.has_mpf_images() {
270                if let Some(ref xmp_data) = self.config.xmp_data {
271                    if !xmp_data.is_empty() {
272                        // Convert raw XMP bytes to string for set_xmp
273                        if let Ok(xmp_str) = core::str::from_utf8(xmp_data) {
274                            segments = segments.set_xmp(xmp_str);
275                        }
276                    }
277                    self.config.xmp_data = None; // Mark as handled
278                }
279                if let Some(ref exif) = self.config.exif_data {
280                    if let Some(exif_bytes) = exif.to_bytes() {
281                        segments.set_exif_mut(exif_bytes);
282                    }
283                    self.config.exif_data = None; // Mark as handled
284                }
285                if let Some(ref icc_data) = self.config.icc_profile {
286                    if !icc_data.is_empty() {
287                        segments = segments.set_icc(icc_data.clone());
288                    }
289                    self.config.icc_profile = None; // Mark as handled
290                }
291            }
292            jpeg = inject_encoder_segments(jpeg, &segments);
293        }
294
295        // Fall back to individual metadata fields for backwards compatibility
296        // These are applied after EncoderSegments if both are provided
297        // (allows override of specific fields while keeping bulk segments)
298        if let Some(ref exif) = self.config.exif_data {
299            if let Some(exif_bytes) = exif.to_bytes() {
300                jpeg = inject_exif(jpeg, &exif_bytes);
301            }
302        }
303
304        if let Some(ref xmp_data) = self.config.xmp_data {
305            jpeg = inject_xmp(jpeg, xmp_data);
306        }
307
308        if let Some(ref icc_data) = self.config.icc_profile {
309            jpeg = inject_icc_profile(jpeg, icc_data);
310        }
311
312        Ok(jpeg)
313    }
314
315    /// Finish encoding to Write destination.
316    #[cfg(feature = "std")]
317    pub fn finish_to<W: Write>(self, mut output: W) -> Result<W> {
318        let jpeg = self.finish()?;
319        output.write_all(&jpeg)?;
320        Ok(output)
321    }
322
323    /// Finish encoding, appending JPEG bytes to an existing Vec.
324    ///
325    /// Useful for no_std environments or buffer reuse.
326    pub fn finish_to_vec(self, output: &mut Vec<u8>) -> Result<()> {
327        let jpeg = self.finish()?;
328        output.extend_from_slice(&jpeg);
329        Ok(())
330    }
331}
332
333/// ICC profile signature for APP2 marker.
334const ICC_PROFILE_SIGNATURE: &[u8; 12] = b"ICC_PROFILE\0";
335
336/// Maximum ICC profile bytes per APP2 marker segment.
337/// APP2 max length is 65535, minus 2 (length) - 12 (signature) - 2 (chunk info) = 65519.
338const MAX_ICC_BYTES_PER_MARKER: usize = 65519;
339
340/// Inject an ICC profile into a JPEG, writing proper APP2 marker chunks.
341///
342/// Inserts APP2 markers right after SOI (and any existing APP0/APP1 markers).
343/// Large profiles are automatically chunked per ICC spec.
344fn inject_icc_profile(jpeg: Vec<u8>, icc_data: &[u8]) -> Vec<u8> {
345    if icc_data.is_empty() {
346        return jpeg;
347    }
348
349    // Find insertion point: after SOI and any APP0/APP1 markers
350    let insert_pos = find_icc_insert_position(&jpeg);
351
352    // Build ICC APP2 marker segments
353    let icc_markers = build_icc_markers(icc_data);
354
355    // Construct new JPEG with ICC markers inserted
356    let mut result = Vec::with_capacity(jpeg.len() + icc_markers.len());
357    result.extend_from_slice(&jpeg[..insert_pos]);
358    result.extend_from_slice(&icc_markers);
359    result.extend_from_slice(&jpeg[insert_pos..]);
360
361    result
362}
363
364/// Find the position to insert ICC markers (after SOI and APP0/APP1).
365fn find_icc_insert_position(jpeg: &[u8]) -> usize {
366    // Start after SOI marker (2 bytes)
367    let mut pos = 2;
368
369    // Skip any existing APP0 (JFIF) and APP1 (EXIF) markers
370    while pos + 4 <= jpeg.len() {
371        if jpeg[pos] != 0xFF {
372            break;
373        }
374
375        let marker = jpeg[pos + 1];
376        // APP0 = 0xE0, APP1 = 0xE1
377        if marker == 0xE0 || marker == 0xE1 {
378            // Get segment length (big-endian, includes length bytes)
379            let length = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
380            pos += 2 + length;
381        } else {
382            break;
383        }
384    }
385
386    pos
387}
388
389/// Build ICC profile APP2 marker segments with proper chunking.
390fn build_icc_markers(icc_data: &[u8]) -> Vec<u8> {
391    let num_chunks = (icc_data.len() + MAX_ICC_BYTES_PER_MARKER - 1) / MAX_ICC_BYTES_PER_MARKER;
392    let mut markers = Vec::new();
393
394    let mut offset = 0;
395    for chunk_num in 0..num_chunks {
396        let chunk_size = (icc_data.len() - offset).min(MAX_ICC_BYTES_PER_MARKER);
397
398        // APP2 marker
399        markers.push(0xFF);
400        markers.push(0xE2); // APP2
401
402        // Length: 2 (length field) + 12 (signature) + 2 (chunk info) + data
403        let segment_length = 2 + 12 + 2 + chunk_size;
404        markers.push((segment_length >> 8) as u8);
405        markers.push(segment_length as u8);
406
407        // ICC_PROFILE signature
408        markers.extend_from_slice(ICC_PROFILE_SIGNATURE);
409
410        // Chunk number (1-based) and total chunks
411        markers.push((chunk_num + 1) as u8);
412        markers.push(num_chunks as u8);
413
414        // ICC data chunk
415        markers.extend_from_slice(&icc_data[offset..offset + chunk_size]);
416
417        offset += chunk_size;
418    }
419
420    markers
421}
422
423/// EXIF signature for APP1 marker.
424const EXIF_SIGNATURE: &[u8; 6] = b"Exif\0\0";
425
426/// Maximum EXIF data bytes per APP1 marker segment.
427/// APP1 max length is 65535, minus 2 (length) - 6 (signature) = 65527.
428const MAX_EXIF_BYTES: usize = 65527;
429
430/// XMP namespace signature for APP1 marker.
431const XMP_NAMESPACE: &[u8; 29] = b"http://ns.adobe.com/xap/1.0/\0";
432
433/// Maximum XMP data bytes per APP1 marker segment.
434/// APP1 max length is 65535, minus 2 (length) - 29 (namespace) = 65504.
435const MAX_XMP_BYTES: usize = 65504;
436
437/// Inject EXIF data into a JPEG as APP1 marker, right after SOI.
438fn inject_exif(jpeg: Vec<u8>, exif_data: &[u8]) -> Vec<u8> {
439    if exif_data.is_empty() {
440        return jpeg;
441    }
442
443    // Truncate if too large
444    let exif_len = exif_data.len().min(MAX_EXIF_BYTES);
445
446    // Build EXIF APP1 marker
447    let mut marker = Vec::with_capacity(4 + 6 + exif_len);
448    marker.push(0xFF);
449    marker.push(0xE1); // APP1
450
451    // Length: 2 (length field) + 6 (signature) + data
452    let segment_length = 2 + 6 + exif_len;
453    marker.push((segment_length >> 8) as u8);
454    marker.push(segment_length as u8);
455
456    // EXIF signature
457    marker.extend_from_slice(EXIF_SIGNATURE);
458
459    // EXIF data
460    marker.extend_from_slice(&exif_data[..exif_len]);
461
462    // Insert after SOI (2 bytes)
463    let mut result = Vec::with_capacity(jpeg.len() + marker.len());
464    result.extend_from_slice(&jpeg[..2]); // SOI
465    result.extend_from_slice(&marker);
466    result.extend_from_slice(&jpeg[2..]);
467
468    result
469}
470
471/// Inject XMP data into a JPEG as APP1 marker.
472///
473/// Inserts after SOI and any existing APP1 (EXIF) markers.
474fn inject_xmp(jpeg: Vec<u8>, xmp_data: &[u8]) -> Vec<u8> {
475    if xmp_data.is_empty() {
476        return jpeg;
477    }
478
479    // Truncate if too large
480    let xmp_len = xmp_data.len().min(MAX_XMP_BYTES);
481
482    // Build XMP APP1 marker
483    let mut marker = Vec::with_capacity(4 + 29 + xmp_len);
484    marker.push(0xFF);
485    marker.push(0xE1); // APP1
486
487    // Length: 2 (length field) + 29 (namespace) + data
488    let segment_length = 2 + 29 + xmp_len;
489    marker.push((segment_length >> 8) as u8);
490    marker.push(segment_length as u8);
491
492    // XMP namespace
493    marker.extend_from_slice(XMP_NAMESPACE);
494
495    // XMP data
496    marker.extend_from_slice(&xmp_data[..xmp_len]);
497
498    // Find insertion point: after SOI and any existing EXIF APP1 markers
499    let insert_pos = find_xmp_insert_position(&jpeg);
500
501    // Construct new JPEG with XMP marker inserted
502    let mut result = Vec::with_capacity(jpeg.len() + marker.len());
503    result.extend_from_slice(&jpeg[..insert_pos]);
504    result.extend_from_slice(&marker);
505    result.extend_from_slice(&jpeg[insert_pos..]);
506
507    result
508}
509
510/// Find the position to insert XMP marker (after SOI and EXIF APP1).
511fn find_xmp_insert_position(jpeg: &[u8]) -> usize {
512    // Start after SOI marker (2 bytes)
513    let mut pos = 2;
514
515    // Skip any existing EXIF APP1 markers
516    while pos + 4 <= jpeg.len() {
517        if jpeg[pos] != 0xFF {
518            break;
519        }
520
521        let marker = jpeg[pos + 1];
522        // APP1 = 0xE1
523        if marker == 0xE1 {
524            // Check if it's EXIF (not XMP)
525            if pos + 10 <= jpeg.len() && &jpeg[pos + 4..pos + 10] == b"Exif\0\0" {
526                // Get segment length (big-endian, includes length bytes)
527                let length = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
528                pos += 2 + length;
529                continue;
530            }
531        }
532        break;
533    }
534
535    pos
536}
537
538/// Marker trait for supported rgb crate pixel types.
539pub trait Pixel: Copy + 'static + bytemuck::Pod {
540    /// Equivalent PixelLayout for this type.
541    const LAYOUT: PixelLayout;
542}
543
544// Implement Pixel for rgb crate types
545impl Pixel for rgb::RGB<u8> {
546    const LAYOUT: PixelLayout = PixelLayout::Rgb8Srgb;
547}
548impl Pixel for rgb::RGBA<u8> {
549    const LAYOUT: PixelLayout = PixelLayout::Rgbx8Srgb;
550}
551impl Pixel for rgb::Bgr<u8> {
552    const LAYOUT: PixelLayout = PixelLayout::Bgr8Srgb;
553}
554impl Pixel for rgb::Bgra<u8> {
555    const LAYOUT: PixelLayout = PixelLayout::Bgrx8Srgb;
556}
557impl Pixel for rgb::Gray<u8> {
558    const LAYOUT: PixelLayout = PixelLayout::Gray8Srgb;
559}
560
561impl Pixel for rgb::RGB<u16> {
562    const LAYOUT: PixelLayout = PixelLayout::Rgb16Linear;
563}
564impl Pixel for rgb::RGBA<u16> {
565    const LAYOUT: PixelLayout = PixelLayout::Rgbx16Linear;
566}
567impl Pixel for rgb::Gray<u16> {
568    const LAYOUT: PixelLayout = PixelLayout::Gray16Linear;
569}
570
571impl Pixel for rgb::RGB<f32> {
572    const LAYOUT: PixelLayout = PixelLayout::RgbF32Linear;
573}
574impl Pixel for rgb::RGBA<f32> {
575    const LAYOUT: PixelLayout = PixelLayout::RgbxF32Linear;
576}
577impl Pixel for rgb::Gray<f32> {
578    const LAYOUT: PixelLayout = PixelLayout::GrayF32Linear;
579}
580
581/// Encoder for rgb crate pixel types.
582///
583/// Type parameter P determines pixel layout at compile time.
584/// For RGBA/BGRA types, 4th channel is ignored.
585pub struct RgbEncoder<P: Pixel> {
586    inner: BytesEncoder,
587    _marker: PhantomData<P>,
588}
589
590impl<P: Pixel> RgbEncoder<P> {
591    pub(crate) fn new(config: EncoderConfig, width: u32, height: u32) -> Result<Self> {
592        let inner = BytesEncoder::new(config, width, height, P::LAYOUT)?;
593        Ok(Self {
594            inner,
595            _marker: PhantomData,
596        })
597    }
598
599    /// Push rows with explicit stride (in pixels).
600    ///
601    /// - `data`: Pixel slice
602    /// - `rows`: Number of scanlines to push
603    /// - `stride`: Pixels per row in buffer (>= width)
604    /// - `stop`: Cancellation token
605    pub fn push(&mut self, data: &[P], rows: usize, stride: usize, stop: impl Stop) -> Result<()> {
606        let stride_bytes = stride * core::mem::size_of::<P>();
607        let bytes = bytemuck::cast_slice(data);
608        self.inner.push(bytes, rows, stride_bytes, stop)
609    }
610
611    /// Push contiguous (packed) data.
612    ///
613    /// Stride assumed to be `width`. Rows inferred from `data.len() / width`.
614    pub fn push_packed(&mut self, data: &[P], stop: impl Stop) -> Result<()> {
615        let bytes = bytemuck::cast_slice(data);
616        self.inner.push_packed(bytes, stop)
617    }
618
619    // === Status ===
620
621    /// Get image width.
622    #[must_use]
623    pub fn width(&self) -> u32 {
624        self.inner.width()
625    }
626
627    /// Get image height.
628    #[must_use]
629    pub fn height(&self) -> u32 {
630        self.inner.height()
631    }
632
633    /// Get number of rows pushed so far.
634    #[must_use]
635    pub fn rows_pushed(&self) -> u32 {
636        self.inner.rows_pushed()
637    }
638
639    /// Get number of rows remaining.
640    #[must_use]
641    pub fn rows_remaining(&self) -> u32 {
642        self.inner.rows_remaining()
643    }
644
645    // === Finish ===
646
647    /// Finish encoding, return JPEG bytes.
648    pub fn finish(self) -> Result<Vec<u8>> {
649        self.inner.finish()
650    }
651
652    /// Finish encoding to Write destination.
653    #[cfg(feature = "std")]
654    pub fn finish_to<W: Write>(self, output: W) -> Result<W> {
655        self.inner.finish_to(output)
656    }
657
658    /// Finish encoding, appending JPEG bytes to an existing Vec.
659    ///
660    /// Useful for no_std environments or buffer reuse.
661    pub fn finish_to_vec(self, output: &mut Vec<u8>) -> Result<()> {
662        self.inner.finish_to_vec(output)
663    }
664}
665
666/// Encoder for planar f32 YCbCr input.
667///
668/// Use when you have pre-converted YCbCr from video decoders, etc.
669/// Skips RGB->YCbCr conversion entirely.
670///
671/// Only valid with `ColorMode::YCbCr`. XYB mode requires RGB input.
672///
673/// # YCbCr Value Range
674///
675/// Input values should be in the centered range:
676/// - Y: 0.0 to 255.0 (luma)
677/// - Cb, Cr: -128.0 to 127.0 (centered chroma)
678///
679/// This matches the output of standard RGB→YCbCr conversion with BT.601 coefficients.
680///
681/// # Streaming
682///
683/// Data can be pushed in any row count - the encoder buffers partial strips
684/// internally and flushes when a complete strip is accumulated.
685pub struct YCbCrPlanarEncoder {
686    /// v2 config (kept for metadata injection)
687    config: EncoderConfig,
688    /// Image width
689    width: u32,
690    /// Image height
691    height: u32,
692    /// Chroma subsampling configuration
693    subsampling: super::encoder_types::ChromaSubsampling,
694    /// Strip height for MCU alignment (8 for 4:4:4, 16 for 4:2:0)
695    strip_height: usize,
696    /// Total rows received so far
697    total_rows_pushed: usize,
698    /// Y plane buffer (accumulates until strip_height rows)
699    y_buffer: Vec<f32>,
700    /// Cb plane buffer
701    cb_buffer: Vec<f32>,
702    /// Cr plane buffer
703    cr_buffer: Vec<f32>,
704    /// Number of rows currently buffered
705    buffered_rows: usize,
706    /// Inner streaming encoder (handles actual encoding)
707    inner: StreamingEncoder,
708}
709
710impl YCbCrPlanarEncoder {
711    pub(crate) fn new(config: EncoderConfig, width: u32, height: u32) -> Result<Self> {
712        // Validate dimensions
713        if width == 0 || height == 0 {
714            return Err(Error::invalid_dimensions(
715                width,
716                height,
717                "dimensions cannot be zero",
718            ));
719        }
720
721        // Check for overflow
722        let pixel_count = (width as u64) * (height as u64);
723        if pixel_count > u32::MAX as u64 {
724            return Err(Error::invalid_dimensions(
725                width,
726                height,
727                "dimensions too large",
728            ));
729        }
730
731        // Extract subsampling from color mode
732        let subsampling = match config.color_mode {
733            super::encoder_types::ColorMode::YCbCr { subsampling } => subsampling,
734            _ => {
735                return Err(Error::invalid_config(
736                    "YCbCrPlanarEncoder requires YCbCr color mode".into(),
737                ))
738            }
739        };
740
741        // Build the streaming encoder
742        let inner = Self::build_streaming_encoder(&config, width, height)?;
743
744        // Get strip height from inner encoder (8 for 4:4:4, 16 for 4:2:0)
745        let strip_height = inner.strip_height();
746
747        // Allocate buffers for one strip
748        let width_usize = width as usize;
749        let buffer_size = width_usize * strip_height;
750
751        Ok(Self {
752            config,
753            width,
754            height,
755            subsampling,
756            strip_height,
757            total_rows_pushed: 0,
758            y_buffer: vec![0.0f32; buffer_size],
759            cb_buffer: vec![0.0f32; buffer_size],
760            cr_buffer: vec![0.0f32; buffer_size],
761            buffered_rows: 0,
762            inner,
763        })
764    }
765
766    /// Build a StreamingEncoder from v2 config for YCbCr planar input.
767    fn build_streaming_encoder(
768        config: &EncoderConfig,
769        width: u32,
770        height: u32,
771    ) -> Result<StreamingEncoder> {
772        use crate::types::PixelFormat;
773
774        let subsampling = match config.color_mode {
775            super::encoder_types::ColorMode::YCbCr { subsampling } => subsampling.into(),
776            _ => crate::types::Subsampling::S444,
777        };
778
779        // Use RGB pixel format - the streaming encoder will accept YCbCr data
780        // via push_ycbcr_strip_f32, but needs a pixel format for buffer sizing
781        let mut builder = StreamingEncoder::new(width, height)
782            .quality(config.quality)
783            .pixel_format(PixelFormat::Rgb) // Buffer sizing only
784            .subsampling(subsampling)
785            .optimize_huffman(config.optimize_huffman)
786            .chroma_downsampling(config.downsampling_method)
787            .restart_interval(config.restart_interval);
788
789        // Apply custom encoding tables if configured
790        if let Some(ref tables) = config.tables {
791            builder = builder.encoding_tables(tables.clone());
792        }
793
794        if config.progressive {
795            builder = builder.progressive(true);
796        }
797
798        // Always pass deringing setting (StreamingEncoder defaults to true)
799        builder = builder.deringing(config.deringing);
800
801        builder = builder.allow_16bit_quant_tables(config.allow_16bit_quant_tables);
802        builder = builder.separate_chroma_tables(config.separate_chroma_tables);
803
804        #[cfg(feature = "parallel")]
805        if config.parallel.is_some() {
806            builder = builder.parallel(true);
807        }
808
809        builder.start()
810    }
811
812    /// Push full-resolution planes. Encoder subsamples chroma as needed.
813    ///
814    /// All three planes must be at full luma resolution (width × rows).
815    /// The encoder will perform chroma subsampling according to the configured
816    /// `ChromaSubsampling` mode.
817    ///
818    /// Data can be pushed in any amount - the encoder buffers partial strips
819    /// internally and flushes when a complete strip is accumulated.
820    ///
821    /// # Arguments
822    /// - `planes`: Y, Cb, Cr plane data with per-plane strides
823    /// - `rows`: Number of luma rows to push
824    /// - `stop`: Cancellation token (use `Unstoppable` if not needed)
825    ///
826    /// # Value Range
827    /// - Y: 0.0 to 255.0
828    /// - Cb, Cr: -128.0 to 127.0
829    pub fn push(&mut self, planes: &YCbCrPlanes<'_>, rows: usize, stop: impl Stop) -> Result<()> {
830        if stop.should_stop() {
831            return Err(Error::cancelled());
832        }
833
834        let width = self.width as usize;
835
836        // Validate row count
837        let new_total = self.total_rows_pushed + rows;
838        if new_total > self.height as usize {
839            return Err(Error::too_many_rows(self.height, new_total as u32));
840        }
841
842        let mut src_row = 0;
843        while src_row < rows {
844            if stop.should_stop() {
845                return Err(Error::cancelled());
846            }
847
848            // How many rows can we add to the buffer?
849            let rows_to_add = (rows - src_row).min(self.strip_height - self.buffered_rows);
850
851            // Copy rows to buffer
852            for i in 0..rows_to_add {
853                let buf_offset = (self.buffered_rows + i) * width;
854                let src_row_idx = src_row + i;
855
856                // Copy Y
857                let y_src_start = src_row_idx * planes.y_stride;
858                let y_src_end = y_src_start + width;
859                if y_src_end > planes.y.len() {
860                    return Err(Error::invalid_buffer_size(y_src_end, planes.y.len()));
861                }
862                self.y_buffer[buf_offset..buf_offset + width]
863                    .copy_from_slice(&planes.y[y_src_start..y_src_end]);
864
865                // Copy Cb
866                let cb_src_start = src_row_idx * planes.cb_stride;
867                let cb_src_end = cb_src_start + width;
868                if cb_src_end > planes.cb.len() {
869                    return Err(Error::invalid_buffer_size(cb_src_end, planes.cb.len()));
870                }
871                self.cb_buffer[buf_offset..buf_offset + width]
872                    .copy_from_slice(&planes.cb[cb_src_start..cb_src_end]);
873
874                // Copy Cr
875                let cr_src_start = src_row_idx * planes.cr_stride;
876                let cr_src_end = cr_src_start + width;
877                if cr_src_end > planes.cr.len() {
878                    return Err(Error::invalid_buffer_size(cr_src_end, planes.cr.len()));
879                }
880                self.cr_buffer[buf_offset..buf_offset + width]
881                    .copy_from_slice(&planes.cr[cr_src_start..cr_src_end]);
882            }
883
884            self.buffered_rows += rows_to_add;
885            src_row += rows_to_add;
886            self.total_rows_pushed += rows_to_add;
887
888            // Flush if we have a complete strip or this is the final strip
889            let remaining_image_rows = self.height as usize - self.inner.rows_pushed();
890            if self.buffered_rows >= self.strip_height || self.buffered_rows >= remaining_image_rows
891            {
892                self.flush_buffer()?;
893            }
894        }
895
896        Ok(())
897    }
898
899    /// Flush buffered rows to the inner encoder.
900    fn flush_buffer(&mut self) -> Result<()> {
901        if self.buffered_rows == 0 {
902            return Ok(());
903        }
904
905        let width = self.width as usize;
906        let data_len = self.buffered_rows * width;
907
908        self.inner.push_ycbcr_strip_f32(
909            &self.y_buffer[..data_len],
910            &self.cb_buffer[..data_len],
911            &self.cr_buffer[..data_len],
912            self.buffered_rows,
913        )?;
914
915        self.buffered_rows = 0;
916        Ok(())
917    }
918
919    /// Push with pre-subsampled chroma.
920    ///
921    /// Use this when your chroma planes are already at the target subsampled resolution.
922    /// Y plane is still at full resolution.
923    ///
924    /// **Note:** Unlike `push()`, this method does not buffer partial strips.
925    /// For best results, push complete strips (multiples of `strip_height` rows,
926    /// which is 8 for 4:4:4 or 16 for 4:2:0).
927    ///
928    /// # Arguments
929    /// - `planes`: Y at full resolution, Cb/Cr at subsampled resolution
930    /// - `y_rows`: Number of luma rows to push
931    /// - `stop`: Cancellation token
932    ///
933    /// # Chroma Dimensions
934    /// The expected chroma dimensions depend on the subsampling mode:
935    /// - 4:4:4 (None): cb/cr at full width × full height
936    /// - 4:2:2 (HalfHorizontal): cb/cr at width/2 × full height
937    /// - 4:2:0 (Quarter): cb/cr at width/2 × height/2
938    pub fn push_subsampled(
939        &mut self,
940        planes: &YCbCrPlanes<'_>,
941        y_rows: usize,
942        stop: impl Stop,
943    ) -> Result<()> {
944        if stop.should_stop() {
945            return Err(Error::cancelled());
946        }
947
948        // Check that we don't have buffered full-resolution data
949        // (can't mix push() and push_subsampled())
950        if self.buffered_rows > 0 {
951            return Err(Error::internal(
952                "cannot mix push() and push_subsampled() - flush first",
953            ));
954        }
955
956        let width = self.width as usize;
957
958        // Calculate chroma dimensions based on subsampling
959        let (chroma_width, chroma_v_factor) = match self.subsampling {
960            super::encoder_types::ChromaSubsampling::None => (width, 1),
961            super::encoder_types::ChromaSubsampling::HalfHorizontal => ((width + 1) / 2, 1),
962            super::encoder_types::ChromaSubsampling::Quarter => ((width + 1) / 2, 2),
963            super::encoder_types::ChromaSubsampling::HalfVertical => (width, 2),
964        };
965        let chroma_rows = (y_rows + chroma_v_factor - 1) / chroma_v_factor;
966
967        // Validate row count
968        let new_total = self.total_rows_pushed + y_rows;
969        if new_total > self.height as usize {
970            return Err(Error::too_many_rows(self.height, new_total as u32));
971        }
972
973        // Check if input is already contiguous
974        let y_contiguous = planes.y_stride == width;
975        let cb_contiguous = planes.cb_stride == chroma_width;
976        let cr_contiguous = planes.cr_stride == chroma_width;
977
978        if y_contiguous && cb_contiguous && cr_contiguous {
979            // Fast path: data is already contiguous
980            let y_len = width * y_rows;
981            let c_len = chroma_width * chroma_rows;
982
983            if planes.y.len() < y_len {
984                return Err(Error::invalid_buffer_size(y_len, planes.y.len()));
985            }
986            if planes.cb.len() < c_len {
987                return Err(Error::invalid_buffer_size(c_len, planes.cb.len()));
988            }
989            if planes.cr.len() < c_len {
990                return Err(Error::invalid_buffer_size(c_len, planes.cr.len()));
991            }
992
993            self.inner.push_ycbcr_strip_f32_subsampled(
994                &planes.y[..y_len],
995                &planes.cb[..c_len],
996                &planes.cr[..c_len],
997                y_rows,
998            )?;
999        } else {
1000            // Slow path: copy strided data to contiguous buffers
1001            let mut y_buf = vec![0.0f32; width * y_rows];
1002            let mut cb_buf = vec![0.0f32; chroma_width * chroma_rows];
1003            let mut cr_buf = vec![0.0f32; chroma_width * chroma_rows];
1004
1005            // Copy Y plane
1006            for row in 0..y_rows {
1007                if stop.should_stop() {
1008                    return Err(Error::cancelled());
1009                }
1010                let dst_start = row * width;
1011                let dst_end = dst_start + width;
1012                let src_start = row * planes.y_stride;
1013                let src_end = src_start + width;
1014                if src_end > planes.y.len() {
1015                    return Err(Error::invalid_buffer_size(src_end, planes.y.len()));
1016                }
1017                y_buf[dst_start..dst_end].copy_from_slice(&planes.y[src_start..src_end]);
1018            }
1019
1020            // Copy chroma planes
1021            for row in 0..chroma_rows {
1022                let dst_start = row * chroma_width;
1023                let dst_end = dst_start + chroma_width;
1024
1025                let cb_src_start = row * planes.cb_stride;
1026                let cb_src_end = cb_src_start + chroma_width;
1027                if cb_src_end > planes.cb.len() {
1028                    return Err(Error::invalid_buffer_size(cb_src_end, planes.cb.len()));
1029                }
1030                cb_buf[dst_start..dst_end].copy_from_slice(&planes.cb[cb_src_start..cb_src_end]);
1031
1032                let cr_src_start = row * planes.cr_stride;
1033                let cr_src_end = cr_src_start + chroma_width;
1034                if cr_src_end > planes.cr.len() {
1035                    return Err(Error::invalid_buffer_size(cr_src_end, planes.cr.len()));
1036                }
1037                cr_buf[dst_start..dst_end].copy_from_slice(&planes.cr[cr_src_start..cr_src_end]);
1038            }
1039
1040            self.inner
1041                .push_ycbcr_strip_f32_subsampled(&y_buf, &cb_buf, &cr_buf, y_rows)?;
1042        }
1043
1044        self.total_rows_pushed += y_rows;
1045        Ok(())
1046    }
1047
1048    // === Status ===
1049
1050    /// Get image width.
1051    #[must_use]
1052    pub fn width(&self) -> u32 {
1053        self.width
1054    }
1055
1056    /// Get image height.
1057    #[must_use]
1058    pub fn height(&self) -> u32 {
1059        self.height
1060    }
1061
1062    /// Get number of rows pushed so far (including buffered rows).
1063    #[must_use]
1064    pub fn rows_pushed(&self) -> u32 {
1065        self.total_rows_pushed as u32
1066    }
1067
1068    /// Get number of rows remaining.
1069    #[must_use]
1070    pub fn rows_remaining(&self) -> u32 {
1071        self.height - self.rows_pushed()
1072    }
1073
1074    // === Finish ===
1075
1076    /// Finish encoding, return JPEG bytes.
1077    pub fn finish(mut self) -> Result<Vec<u8>> {
1078        // Check if all rows were pushed
1079        if self.total_rows_pushed != self.height as usize {
1080            return Err(Error::incomplete_image(
1081                self.height,
1082                self.total_rows_pushed as u32,
1083            ));
1084        }
1085
1086        // Flush any remaining buffered rows
1087        self.flush_buffer()?;
1088
1089        // Finish streaming encoder
1090        let mut jpeg = self.inner.finish()?;
1091
1092        // When using EncoderSegments with MPF images, we need to merge individual
1093        // metadata fields (xmp_data, exif_data, icc_profile) into segments BEFORE
1094        // calling inject_encoder_segments. This ensures the MPF offset calculation
1095        // includes all metadata that will be in the final output.
1096        if let Some(mut segments) = self.config.segments.take() {
1097            // Merge individual metadata into segments if MPF is present
1098            // (MPF offset calculation needs all data to be accounted for)
1099            if segments.has_mpf_images() {
1100                if let Some(ref xmp_data) = self.config.xmp_data {
1101                    if !xmp_data.is_empty() {
1102                        // Convert raw XMP bytes to string for set_xmp
1103                        if let Ok(xmp_str) = core::str::from_utf8(xmp_data) {
1104                            segments = segments.set_xmp(xmp_str);
1105                        }
1106                    }
1107                    self.config.xmp_data = None; // Mark as handled
1108                }
1109                if let Some(ref exif) = self.config.exif_data {
1110                    if let Some(exif_bytes) = exif.to_bytes() {
1111                        segments.set_exif_mut(exif_bytes);
1112                    }
1113                    self.config.exif_data = None; // Mark as handled
1114                }
1115                if let Some(ref icc_data) = self.config.icc_profile {
1116                    if !icc_data.is_empty() {
1117                        segments = segments.set_icc(icc_data.clone());
1118                    }
1119                    self.config.icc_profile = None; // Mark as handled
1120                }
1121            }
1122            jpeg = inject_encoder_segments(jpeg, &segments);
1123        }
1124
1125        // Fall back to individual metadata fields for backwards compatibility
1126        if let Some(ref exif) = self.config.exif_data {
1127            if let Some(exif_bytes) = exif.to_bytes() {
1128                jpeg = inject_exif(jpeg, &exif_bytes);
1129            }
1130        }
1131
1132        if let Some(ref xmp_data) = self.config.xmp_data {
1133            jpeg = inject_xmp(jpeg, xmp_data);
1134        }
1135
1136        if let Some(ref icc_data) = self.config.icc_profile {
1137            jpeg = inject_icc_profile(jpeg, icc_data);
1138        }
1139
1140        Ok(jpeg)
1141    }
1142
1143    /// Finish encoding to Write destination.
1144    #[cfg(feature = "std")]
1145    pub fn finish_to<W: Write>(self, mut output: W) -> Result<W> {
1146        let jpeg = self.finish()?;
1147        output.write_all(&jpeg)?;
1148        Ok(output)
1149    }
1150
1151    /// Finish encoding, appending JPEG bytes to an existing Vec.
1152    ///
1153    /// Useful for no_std environments or buffer reuse.
1154    pub fn finish_to_vec(self, output: &mut Vec<u8>) -> Result<()> {
1155        let jpeg = self.finish()?;
1156        output.extend_from_slice(&jpeg);
1157        Ok(())
1158    }
1159}
1160
1161#[cfg(test)]
1162mod tests {
1163    use super::*;
1164    use crate::encode::ChromaSubsampling;
1165    use crate::error::ErrorKind;
1166    use enough::Unstoppable;
1167    use rgb::RGB;
1168
1169    #[test]
1170    fn test_bytes_encoder_basic() {
1171        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1172        let mut enc = config
1173            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1174            .unwrap();
1175
1176        // Create 8x8 red image
1177        let pixels = [255u8, 0, 0].repeat(64);
1178        enc.push_packed(&pixels, Unstoppable).unwrap();
1179
1180        let jpeg = enc.finish().unwrap();
1181        assert!(!jpeg.is_empty());
1182        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]); // JPEG SOI marker
1183    }
1184
1185    #[test]
1186    fn test_rgb_encoder_basic() {
1187        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1188        let mut enc = config.encode_from_rgb::<RGB<u8>>(8, 8).unwrap();
1189
1190        // Create 8x8 green image
1191        let pixels: Vec<RGB<u8>> = vec![RGB::new(0, 255, 0); 64];
1192        enc.push_packed(&pixels, Unstoppable).unwrap();
1193
1194        let jpeg = enc.finish().unwrap();
1195        assert!(!jpeg.is_empty());
1196        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]); // JPEG SOI marker
1197    }
1198
1199    #[test]
1200    fn test_stride_validation() {
1201        let config = EncoderConfig::ycbcr(90.0, ChromaSubsampling::None);
1202        let mut enc = config
1203            .encode_from_bytes(100, 10, PixelLayout::Rgb8Srgb)
1204            .unwrap();
1205
1206        // Stride too small (less than width * 3)
1207        let result = enc.push(&[0u8; 100], 1, 100, Unstoppable);
1208        assert!(matches!(
1209            result.as_ref().map_err(|e| e.kind()),
1210            Err(ErrorKind::StrideTooSmall { .. })
1211        ));
1212    }
1213
1214    #[test]
1215    fn test_too_many_rows() {
1216        let config = EncoderConfig::ycbcr(90.0, ChromaSubsampling::None);
1217        let mut enc = config
1218            .encode_from_bytes(8, 4, PixelLayout::Rgb8Srgb)
1219            .unwrap();
1220
1221        let row_data = vec![0u8; 8 * 3];
1222
1223        // Push all 4 rows
1224        for _ in 0..4 {
1225            enc.push_packed(&row_data, Unstoppable).unwrap();
1226        }
1227
1228        // Try to push one more
1229        let result = enc.push_packed(&row_data, Unstoppable);
1230        assert!(matches!(
1231            result.as_ref().map_err(|e| e.kind()),
1232            Err(ErrorKind::TooManyRows { .. })
1233        ));
1234    }
1235
1236    #[test]
1237    fn test_incomplete_image() {
1238        let config = EncoderConfig::ycbcr(90.0, ChromaSubsampling::None);
1239        let mut enc = config
1240            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1241            .unwrap();
1242
1243        // Only push 4 rows
1244        let rows_data = vec![0u8; 8 * 3 * 4];
1245        enc.push_packed(&rows_data, Unstoppable).unwrap();
1246
1247        // Try to finish
1248        let result = enc.finish();
1249        assert!(matches!(
1250            result.as_ref().map_err(|e| e.kind()),
1251            Err(ErrorKind::IncompleteImage { .. })
1252        ));
1253    }
1254
1255    #[test]
1256    fn test_icc_profile_injection() {
1257        // Small fake ICC profile (just for testing structure)
1258        let fake_icc = vec![0u8; 1000];
1259
1260        let config =
1261            EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).icc_profile(fake_icc.clone());
1262        let mut enc = config
1263            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1264            .unwrap();
1265
1266        let pixels = vec![128u8; 8 * 8 * 3];
1267        enc.push_packed(&pixels, Unstoppable).unwrap();
1268
1269        let jpeg = enc.finish().unwrap();
1270
1271        // Verify JPEG structure
1272        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]); // SOI
1273
1274        // Find APP2 ICC profile marker
1275        let mut found_icc = false;
1276        let mut pos = 2;
1277        while pos + 4 < jpeg.len() {
1278            if jpeg[pos] == 0xFF && jpeg[pos + 1] == 0xE2 {
1279                // APP2 marker - check for ICC signature
1280                if jpeg.len() > pos + 16 && &jpeg[pos + 4..pos + 16] == b"ICC_PROFILE\0" {
1281                    found_icc = true;
1282                    // Verify chunk numbers
1283                    assert_eq!(jpeg[pos + 16], 1); // chunk 1
1284                    assert_eq!(jpeg[pos + 17], 1); // of 1 total
1285                    break;
1286                }
1287            }
1288            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
1289                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1290                pos += 2 + len;
1291            } else {
1292                pos += 1;
1293            }
1294        }
1295        assert!(found_icc, "ICC profile APP2 marker not found");
1296    }
1297
1298    #[test]
1299    fn test_icc_profile_chunking() {
1300        // Large ICC profile that requires multiple chunks
1301        let large_icc = vec![0xABu8; 100_000]; // > 65519 bytes
1302
1303        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).icc_profile(large_icc);
1304        let mut enc = config
1305            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1306            .unwrap();
1307
1308        let pixels = vec![128u8; 8 * 8 * 3];
1309        enc.push_packed(&pixels, Unstoppable).unwrap();
1310
1311        let jpeg = enc.finish().unwrap();
1312
1313        // Count APP2 ICC chunks
1314        let mut chunk_count = 0;
1315        let mut pos = 2;
1316        while pos + 4 < jpeg.len() {
1317            if jpeg[pos] == 0xFF
1318                && jpeg[pos + 1] == 0xE2
1319                && jpeg.len() > pos + 16
1320                && &jpeg[pos + 4..pos + 16] == b"ICC_PROFILE\0"
1321            {
1322                chunk_count += 1;
1323                let chunk_num = jpeg[pos + 16];
1324                let total_chunks = jpeg[pos + 17];
1325                assert_eq!(chunk_num as usize, chunk_count);
1326                assert_eq!(total_chunks, 2); // 100000 / 65519 = 2 chunks
1327            }
1328            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
1329                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1330                pos += 2 + len;
1331            } else {
1332                pos += 1;
1333            }
1334        }
1335        assert_eq!(chunk_count, 2, "Expected 2 ICC chunks for 100KB profile");
1336    }
1337
1338    #[test]
1339    fn test_finish_to_vec() {
1340        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1341        let mut enc = config.encode_from_rgb::<RGB<u8>>(8, 8).unwrap();
1342
1343        let pixels: Vec<RGB<u8>> = vec![RGB::new(100, 150, 200); 64];
1344        enc.push_packed(&pixels, Unstoppable).unwrap();
1345
1346        // Finish to existing vec
1347        let mut output = Vec::new();
1348        enc.finish_to_vec(&mut output).unwrap();
1349
1350        assert!(!output.is_empty());
1351        assert_eq!(&output[0..2], &[0xFF, 0xD8]); // JPEG SOI marker
1352    }
1353
1354    #[test]
1355    fn test_finish_to_vec_append() {
1356        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1357        let mut enc = config.encode_from_rgb::<RGB<u8>>(8, 8).unwrap();
1358
1359        let pixels: Vec<RGB<u8>> = vec![RGB::new(100, 150, 200); 64];
1360        enc.push_packed(&pixels, Unstoppable).unwrap();
1361
1362        // Finish to vec with existing content
1363        let mut output = vec![0xDE, 0xAD, 0xBE, 0xEF];
1364        let prefix_len = output.len();
1365        enc.finish_to_vec(&mut output).unwrap();
1366
1367        // Verify prefix preserved
1368        assert_eq!(&output[0..4], &[0xDE, 0xAD, 0xBE, 0xEF]);
1369        // Verify JPEG appended
1370        assert_eq!(&output[prefix_len..prefix_len + 2], &[0xFF, 0xD8]);
1371    }
1372
1373    #[test]
1374    fn test_icc_roundtrip_extraction() {
1375        // Test that we can extract the same ICC profile we injected
1376        let original_icc: Vec<u8> = (0..=255).cycle().take(3000).collect();
1377
1378        let config =
1379            EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).icc_profile(original_icc.clone());
1380        let mut enc = config
1381            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1382            .unwrap();
1383
1384        let pixels = vec![100u8; 8 * 8 * 3];
1385        enc.push_packed(&pixels, Unstoppable).unwrap();
1386
1387        let jpeg = enc.finish().unwrap();
1388
1389        // Extract ICC profile using the existing extraction function
1390        let extracted = crate::color::icc::extract_icc_profile(&jpeg);
1391        assert!(extracted.is_some(), "Failed to extract ICC profile");
1392        assert_eq!(
1393            extracted.unwrap(),
1394            original_icc,
1395            "Extracted ICC doesn't match original"
1396        );
1397    }
1398
1399    // =========================================================================
1400    // YCbCrPlanarEncoder tests
1401    // =========================================================================
1402
1403    /// Helper: Convert RGB to YCbCr f32 using BT.601 coefficients.
1404    fn rgb_to_ycbcr_f32(r: u8, g: u8, b: u8) -> (f32, f32, f32) {
1405        let r = r as f32;
1406        let g = g as f32;
1407        let b = b as f32;
1408
1409        // BT.601 coefficients
1410        let y = 0.299 * r + 0.587 * g + 0.114 * b;
1411        let cb = -0.168736 * r - 0.331264 * g + 0.5 * b;
1412        let cr = 0.5 * r - 0.418688 * g - 0.081312 * b;
1413
1414        (y, cb, cr)
1415    }
1416
1417    #[test]
1418    fn test_ycbcr_planar_encoder_basic() {
1419        use crate::encode::YCbCrPlanes;
1420
1421        let width = 8usize;
1422        let height = 8usize;
1423
1424        // Create YCbCr data for a solid red image
1425        let mut y_plane = vec![0.0f32; width * height];
1426        let mut cb_plane = vec![0.0f32; width * height];
1427        let mut cr_plane = vec![0.0f32; width * height];
1428
1429        for i in 0..(width * height) {
1430            let (y, cb, cr) = rgb_to_ycbcr_f32(255, 0, 0); // Red
1431            y_plane[i] = y;
1432            cb_plane[i] = cb;
1433            cr_plane[i] = cr;
1434        }
1435
1436        let planes = YCbCrPlanes {
1437            y: &y_plane,
1438            y_stride: width,
1439            cb: &cb_plane,
1440            cb_stride: width,
1441            cr: &cr_plane,
1442            cr_stride: width,
1443        };
1444
1445        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1446        let mut enc = config
1447            .encode_from_ycbcr_planar(width as u32, height as u32)
1448            .unwrap();
1449
1450        enc.push(&planes, height, Unstoppable).unwrap();
1451
1452        let jpeg = enc.finish().unwrap();
1453        assert!(!jpeg.is_empty());
1454        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]); // JPEG SOI marker
1455    }
1456
1457    #[test]
1458    fn test_ycbcr_planar_encoder_gradient() {
1459        use crate::encode::YCbCrPlanes;
1460
1461        let width = 64usize;
1462        let height = 64usize;
1463
1464        // Create YCbCr data for a horizontal gradient (black to white)
1465        let mut y_plane = vec![0.0f32; width * height];
1466        let mut cb_plane = vec![0.0f32; width * height];
1467        let mut cr_plane = vec![0.0f32; width * height];
1468
1469        for row in 0..height {
1470            for col in 0..width {
1471                let gray = (col * 255 / (width - 1)) as u8;
1472                let (y, cb, cr) = rgb_to_ycbcr_f32(gray, gray, gray);
1473                let idx = row * width + col;
1474                y_plane[idx] = y;
1475                cb_plane[idx] = cb;
1476                cr_plane[idx] = cr;
1477            }
1478        }
1479
1480        let planes = YCbCrPlanes {
1481            y: &y_plane,
1482            y_stride: width,
1483            cb: &cb_plane,
1484            cb_stride: width,
1485            cr: &cr_plane,
1486            cr_stride: width,
1487        };
1488
1489        // Test with 4:4:4 subsampling
1490        let config = EncoderConfig::ycbcr(90, ChromaSubsampling::None);
1491        let mut enc = config
1492            .encode_from_ycbcr_planar(width as u32, height as u32)
1493            .unwrap();
1494
1495        enc.push(&planes, height, Unstoppable).unwrap();
1496
1497        let jpeg = enc.finish().unwrap();
1498        assert!(!jpeg.is_empty());
1499        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1500    }
1501
1502    #[test]
1503    fn test_ycbcr_planar_encoder_strided_input() {
1504        use crate::encode::YCbCrPlanes;
1505
1506        let width = 8usize;
1507        let height = 8usize;
1508        let stride = 16usize; // Larger stride than width
1509
1510        // Create YCbCr data with padding (stride > width)
1511        let mut y_plane = vec![0.0f32; stride * height];
1512        let mut cb_plane = vec![0.0f32; stride * height];
1513        let mut cr_plane = vec![0.0f32; stride * height];
1514
1515        for row in 0..height {
1516            for col in 0..width {
1517                let (y, cb, cr) = rgb_to_ycbcr_f32(0, 255, 0); // Green
1518                let idx = row * stride + col;
1519                y_plane[idx] = y;
1520                cb_plane[idx] = cb;
1521                cr_plane[idx] = cr;
1522            }
1523            // Rest of the row (padding) is zeros
1524        }
1525
1526        let planes = YCbCrPlanes {
1527            y: &y_plane,
1528            y_stride: stride,
1529            cb: &cb_plane,
1530            cb_stride: stride,
1531            cr: &cr_plane,
1532            cr_stride: stride,
1533        };
1534
1535        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1536        let mut enc = config
1537            .encode_from_ycbcr_planar(width as u32, height as u32)
1538            .unwrap();
1539
1540        enc.push(&planes, height, Unstoppable).unwrap();
1541
1542        let jpeg = enc.finish().unwrap();
1543        assert!(!jpeg.is_empty());
1544        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1545    }
1546
1547    #[test]
1548    fn test_ycbcr_planar_encoder_multiple_pushes() {
1549        use crate::encode::YCbCrPlanes;
1550
1551        // Use 4:4:4 which has 8-row strips, allowing 8-row pushes
1552        let width = 16usize;
1553        let height = 32usize;
1554        let rows_per_push = 8usize;
1555
1556        // Create full image YCbCr data
1557        let mut y_plane = vec![0.0f32; width * height];
1558        let mut cb_plane = vec![0.0f32; width * height];
1559        let mut cr_plane = vec![0.0f32; width * height];
1560
1561        for row in 0..height {
1562            for col in 0..width {
1563                // Different color for each 8-row strip
1564                let strip = row / 8;
1565                let (r, g, b) = match strip {
1566                    0 => (255, 0, 0),   // Red
1567                    1 => (0, 255, 0),   // Green
1568                    2 => (0, 0, 255),   // Blue
1569                    _ => (255, 255, 0), // Yellow
1570                };
1571                let (y, cb, cr) = rgb_to_ycbcr_f32(r, g, b);
1572                let idx = row * width + col;
1573                y_plane[idx] = y;
1574                cb_plane[idx] = cb;
1575                cr_plane[idx] = cr;
1576            }
1577        }
1578
1579        // Use 4:4:4 subsampling which has 8-row strip height
1580        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::None);
1581        let mut enc = config
1582            .encode_from_ycbcr_planar(width as u32, height as u32)
1583            .unwrap();
1584
1585        // Push in 4 chunks of 8 rows each
1586        for chunk in 0..4 {
1587            let start_row = chunk * rows_per_push;
1588            let start_idx = start_row * width;
1589            let end_idx = start_idx + rows_per_push * width;
1590
1591            let planes = YCbCrPlanes {
1592                y: &y_plane[start_idx..end_idx],
1593                y_stride: width,
1594                cb: &cb_plane[start_idx..end_idx],
1595                cb_stride: width,
1596                cr: &cr_plane[start_idx..end_idx],
1597                cr_stride: width,
1598            };
1599
1600            enc.push(&planes, rows_per_push, Unstoppable).unwrap();
1601            assert_eq!(enc.rows_pushed(), ((chunk + 1) * rows_per_push) as u32);
1602        }
1603
1604        let jpeg = enc.finish().unwrap();
1605        assert!(!jpeg.is_empty());
1606        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1607    }
1608
1609    #[test]
1610    fn test_ycbcr_planar_encoder_incomplete_image() {
1611        use crate::encode::YCbCrPlanes;
1612
1613        // Use 4:4:4 with 8-row strip height for easier testing
1614        let width = 8usize;
1615        let height = 16usize;
1616
1617        // Only create data for half the image
1618        let half_height = 8usize;
1619        let y_plane = vec![128.0f32; width * half_height];
1620        let cb_plane = vec![0.0f32; width * half_height];
1621        let cr_plane = vec![0.0f32; width * half_height];
1622
1623        let planes = YCbCrPlanes {
1624            y: &y_plane,
1625            y_stride: width,
1626            cb: &cb_plane,
1627            cb_stride: width,
1628            cr: &cr_plane,
1629            cr_stride: width,
1630        };
1631
1632        // Use 4:4:4 subsampling which has 8-row strip height
1633        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::None);
1634        let mut enc = config
1635            .encode_from_ycbcr_planar(width as u32, height as u32)
1636            .unwrap();
1637
1638        // Only push half the rows (8 of 16)
1639        enc.push(&planes, half_height, Unstoppable).unwrap();
1640
1641        // Try to finish - should fail because only 8 of 16 rows pushed
1642        let result = enc.finish();
1643        assert!(matches!(
1644            result.as_ref().map_err(|e| e.kind()),
1645            Err(ErrorKind::IncompleteImage { .. })
1646        ));
1647    }
1648
1649    #[test]
1650    fn test_ycbcr_planar_encoder_subsampled_444() {
1651        use crate::encode::YCbCrPlanes;
1652
1653        let width = 16usize;
1654        let height = 16usize;
1655
1656        // For 4:4:4, chroma is same size as luma
1657        let y_plane: Vec<f32> = (0..width * height).map(|i| (i % 256) as f32).collect();
1658        let cb_plane = vec![0.0f32; width * height];
1659        let cr_plane = vec![0.0f32; width * height];
1660
1661        let planes = YCbCrPlanes {
1662            y: &y_plane,
1663            y_stride: width,
1664            cb: &cb_plane,
1665            cb_stride: width,
1666            cr: &cr_plane,
1667            cr_stride: width,
1668        };
1669
1670        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::None);
1671        let mut enc = config
1672            .encode_from_ycbcr_planar(width as u32, height as u32)
1673            .unwrap();
1674
1675        enc.push_subsampled(&planes, height, Unstoppable).unwrap();
1676
1677        let jpeg = enc.finish().unwrap();
1678        assert!(!jpeg.is_empty());
1679        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1680    }
1681
1682    #[test]
1683    fn test_ycbcr_planar_encoder_subsampled_420() {
1684        use crate::encode::YCbCrPlanes;
1685
1686        let width = 16usize;
1687        let height = 16usize;
1688        let chroma_width = (width + 1) / 2; // 8
1689        let chroma_height = (height + 1) / 2; // 8
1690
1691        // For 4:2:0, chroma is half size in both dimensions
1692        let y_plane: Vec<f32> = (0..width * height).map(|i| (i % 256) as f32).collect();
1693        let cb_plane = vec![0.0f32; chroma_width * chroma_height];
1694        let cr_plane = vec![0.0f32; chroma_width * chroma_height];
1695
1696        let planes = YCbCrPlanes {
1697            y: &y_plane,
1698            y_stride: width,
1699            cb: &cb_plane,
1700            cb_stride: chroma_width,
1701            cr: &cr_plane,
1702            cr_stride: chroma_width,
1703        };
1704
1705        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1706        let mut enc = config
1707            .encode_from_ycbcr_planar(width as u32, height as u32)
1708            .unwrap();
1709
1710        enc.push_subsampled(&planes, height, Unstoppable).unwrap();
1711
1712        let jpeg = enc.finish().unwrap();
1713        assert!(!jpeg.is_empty());
1714        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1715    }
1716
1717    #[test]
1718    fn test_ycbcr_planar_encoder_requires_ycbcr_mode() {
1719        // Try to create planar encoder with XYB mode - should fail
1720        let config = EncoderConfig::xyb(85, crate::encode::encoder_types::XybSubsampling::BQuarter);
1721        let result = config.encode_from_ycbcr_planar(8, 8);
1722        assert!(result.is_err());
1723    }
1724
1725    #[test]
1726    fn test_ycbcr_planar_encoder_status_methods() {
1727        use crate::encode::YCbCrPlanes;
1728
1729        let width = 16u32;
1730        let height = 32u32;
1731
1732        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1733        let enc = config.encode_from_ycbcr_planar(width, height).unwrap();
1734
1735        assert_eq!(enc.width(), width);
1736        assert_eq!(enc.height(), height);
1737        assert_eq!(enc.rows_pushed(), 0);
1738        assert_eq!(enc.rows_remaining(), height);
1739
1740        // Push some rows
1741        let y_plane = vec![128.0f32; 16 * 16];
1742        let cb_plane = vec![0.0f32; 16 * 16];
1743        let cr_plane = vec![0.0f32; 16 * 16];
1744
1745        let planes = YCbCrPlanes {
1746            y: &y_plane,
1747            y_stride: 16,
1748            cb: &cb_plane,
1749            cb_stride: 16,
1750            cr: &cr_plane,
1751            cr_stride: 16,
1752        };
1753
1754        let mut enc = config.encode_from_ycbcr_planar(width, height).unwrap();
1755        enc.push(&planes, 16, Unstoppable).unwrap();
1756
1757        assert_eq!(enc.rows_pushed(), 16);
1758        assert_eq!(enc.rows_remaining(), 16);
1759    }
1760
1761    #[test]
1762    fn test_ycbcr_planar_encoder_with_icc_profile() {
1763        use crate::encode::YCbCrPlanes;
1764
1765        let width = 8usize;
1766        let height = 8usize;
1767
1768        let y_plane = vec![128.0f32; width * height];
1769        let cb_plane = vec![0.0f32; width * height];
1770        let cr_plane = vec![0.0f32; width * height];
1771
1772        let planes = YCbCrPlanes {
1773            y: &y_plane,
1774            y_stride: width,
1775            cb: &cb_plane,
1776            cb_stride: width,
1777            cr: &cr_plane,
1778            cr_stride: width,
1779        };
1780
1781        // Create config with ICC profile
1782        let fake_icc = vec![0xABu8; 1000];
1783        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).icc_profile(fake_icc);
1784        let mut enc = config
1785            .encode_from_ycbcr_planar(width as u32, height as u32)
1786            .unwrap();
1787
1788        enc.push(&planes, height, Unstoppable).unwrap();
1789
1790        let jpeg = enc.finish().unwrap();
1791
1792        // Verify ICC profile was injected
1793        let mut found_icc = false;
1794        let mut pos = 2;
1795        while pos + 4 < jpeg.len() {
1796            if jpeg[pos] == 0xFF
1797                && jpeg[pos + 1] == 0xE2
1798                && jpeg.len() > pos + 16
1799                && &jpeg[pos + 4..pos + 16] == b"ICC_PROFILE\0"
1800            {
1801                found_icc = true;
1802                break;
1803            }
1804            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
1805                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1806                pos += 2 + len;
1807            } else {
1808                pos += 1;
1809            }
1810        }
1811        assert!(found_icc, "ICC profile should be present in output");
1812    }
1813
1814    #[test]
1815    fn test_ycbcr_planar_encoder_odd_width() {
1816        use crate::encode::YCbCrPlanes;
1817
1818        // Non-8-aligned width (tests partial block handling)
1819        let width = 13usize;
1820        let height = 17usize;
1821
1822        let y_plane: Vec<f32> = (0..width * height).map(|i| (i % 256) as f32).collect();
1823        let cb_plane = vec![0.0f32; width * height];
1824        let cr_plane = vec![0.0f32; width * height];
1825
1826        let planes = YCbCrPlanes {
1827            y: &y_plane,
1828            y_stride: width,
1829            cb: &cb_plane,
1830            cb_stride: width,
1831            cr: &cr_plane,
1832            cr_stride: width,
1833        };
1834
1835        // Test with 4:4:4 (strip height 8)
1836        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::None);
1837        let mut enc = config
1838            .encode_from_ycbcr_planar(width as u32, height as u32)
1839            .unwrap();
1840
1841        enc.push(&planes, height, Unstoppable).unwrap();
1842
1843        let jpeg = enc.finish().unwrap();
1844        assert!(!jpeg.is_empty());
1845        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1846    }
1847
1848    #[test]
1849    fn test_ycbcr_planar_encoder_single_row_pushes() {
1850        use crate::encode::YCbCrPlanes;
1851
1852        // Push one row at a time - tests buffering
1853        let width = 16usize;
1854        let height = 24usize;
1855
1856        let y_plane: Vec<f32> = (0..width * height).map(|i| (i % 256) as f32).collect();
1857        let cb_plane = vec![0.0f32; width * height];
1858        let cr_plane = vec![0.0f32; width * height];
1859
1860        // Use 4:4:4 (strip height 8)
1861        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::None);
1862        let mut enc = config
1863            .encode_from_ycbcr_planar(width as u32, height as u32)
1864            .unwrap();
1865
1866        // Push one row at a time
1867        for row in 0..height {
1868            let start = row * width;
1869            let end = start + width;
1870            let planes = YCbCrPlanes {
1871                y: &y_plane[start..end],
1872                y_stride: width,
1873                cb: &cb_plane[start..end],
1874                cb_stride: width,
1875                cr: &cr_plane[start..end],
1876                cr_stride: width,
1877            };
1878            enc.push(&planes, 1, Unstoppable).unwrap();
1879        }
1880
1881        let jpeg = enc.finish().unwrap();
1882        assert!(!jpeg.is_empty());
1883        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1884    }
1885
1886    #[test]
1887    fn test_ycbcr_planar_encoder_420_partial_pushes() {
1888        use crate::encode::YCbCrPlanes;
1889
1890        // 4:2:0 has 16-row strip height - test partial push buffering
1891        let width = 16usize;
1892        let height = 32usize;
1893
1894        let y_plane: Vec<f32> = (0..width * height).map(|i| (i % 256) as f32).collect();
1895        let cb_plane = vec![0.0f32; width * height];
1896        let cr_plane = vec![0.0f32; width * height];
1897
1898        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter);
1899        let mut enc = config
1900            .encode_from_ycbcr_planar(width as u32, height as u32)
1901            .unwrap();
1902
1903        // Push in chunks smaller than strip height (16)
1904        // Push 5 + 5 + 6 = 16 rows (one strip)
1905        // Then 8 + 8 = 16 rows (one strip)
1906        let push_sizes = [5, 5, 6, 8, 8];
1907        let mut offset = 0;
1908        for &rows in &push_sizes {
1909            let start = offset * width;
1910            let end = start + rows * width;
1911            let planes = YCbCrPlanes {
1912                y: &y_plane[start..end],
1913                y_stride: width,
1914                cb: &cb_plane[start..end],
1915                cb_stride: width,
1916                cr: &cr_plane[start..end],
1917                cr_stride: width,
1918            };
1919            enc.push(&planes, rows, Unstoppable).unwrap();
1920            offset += rows;
1921        }
1922
1923        let jpeg = enc.finish().unwrap();
1924        assert!(!jpeg.is_empty());
1925        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1926    }
1927
1928    // =========================================================================
1929    // EncoderSegments integration tests
1930    // =========================================================================
1931
1932    #[test]
1933    fn test_encoder_segments_injection() {
1934        use crate::encode::extras::EncoderSegments;
1935
1936        // Create segments with EXIF and comment
1937        let segments = EncoderSegments::new()
1938            .set_exif(vec![0x49, 0x49, 0x2A, 0x00]) // Minimal TIFF header
1939            .add_comment("Test comment");
1940
1941        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).with_segments(segments);
1942
1943        let mut enc = config
1944            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
1945            .unwrap();
1946
1947        let pixels = vec![128u8; 8 * 8 * 3];
1948        enc.push_packed(&pixels, Unstoppable).unwrap();
1949
1950        let jpeg = enc.finish().unwrap();
1951
1952        // Verify SOI
1953        assert_eq!(&jpeg[0..2], &[0xFF, 0xD8]);
1954
1955        // Find EXIF (APP1 with "Exif\0\0" signature)
1956        let mut found_exif = false;
1957        let mut pos = 2;
1958        while pos + 4 < jpeg.len() {
1959            if jpeg[pos] == 0xFF
1960                && jpeg[pos + 1] == 0xE1
1961                && jpeg.len() > pos + 10
1962                && &jpeg[pos + 4..pos + 10] == b"Exif\0\0"
1963            {
1964                found_exif = true;
1965                break;
1966            }
1967            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
1968                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1969                pos += 2 + len;
1970            } else {
1971                pos += 1;
1972            }
1973        }
1974        assert!(found_exif, "EXIF segment not found in output");
1975
1976        // Find comment (COM marker 0xFE)
1977        let mut found_comment = false;
1978        pos = 2;
1979        while pos + 4 < jpeg.len() {
1980            if jpeg[pos] == 0xFF && jpeg[pos + 1] == 0xFE {
1981                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1982                let comment_data = &jpeg[pos + 4..pos + 2 + len];
1983                if comment_data == b"Test comment" {
1984                    found_comment = true;
1985                    break;
1986                }
1987            }
1988            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
1989                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
1990                pos += 2 + len;
1991            } else {
1992                pos += 1;
1993            }
1994        }
1995        assert!(found_comment, "Comment segment not found in output");
1996    }
1997
1998    #[test]
1999    fn test_encoder_segments_icc_chunking() {
2000        use crate::encode::extras::EncoderSegments;
2001
2002        // Large ICC profile that needs chunking
2003        let large_profile = vec![0xAB; 100_000];
2004
2005        let segments = EncoderSegments::new().set_icc(large_profile);
2006
2007        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).with_segments(segments);
2008
2009        let mut enc = config
2010            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
2011            .unwrap();
2012
2013        let pixels = vec![128u8; 8 * 8 * 3];
2014        enc.push_packed(&pixels, Unstoppable).unwrap();
2015
2016        let jpeg = enc.finish().unwrap();
2017
2018        // Count ICC chunks
2019        let mut chunk_count = 0;
2020        let mut pos = 2;
2021        while pos + 4 < jpeg.len() {
2022            if jpeg[pos] == 0xFF
2023                && jpeg[pos + 1] == 0xE2
2024                && jpeg.len() > pos + 16
2025                && &jpeg[pos + 4..pos + 16] == b"ICC_PROFILE\0"
2026            {
2027                chunk_count += 1;
2028            }
2029            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
2030                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
2031                pos += 2 + len;
2032            } else {
2033                pos += 1;
2034            }
2035        }
2036        assert_eq!(chunk_count, 2, "Expected 2 ICC chunks for 100KB profile");
2037    }
2038
2039    #[test]
2040    fn test_encoder_segments_xmp() {
2041        use crate::encode::extras::EncoderSegments;
2042
2043        let xmp = "<?xml version=\"1.0\"?><x:xmpmeta>test XMP data</x:xmpmeta>";
2044        let segments = EncoderSegments::new().set_xmp(xmp);
2045
2046        let config = EncoderConfig::ycbcr(85, ChromaSubsampling::Quarter).with_segments(segments);
2047
2048        let mut enc = config
2049            .encode_from_bytes(8, 8, PixelLayout::Rgb8Srgb)
2050            .unwrap();
2051
2052        let pixels = vec![128u8; 8 * 8 * 3];
2053        enc.push_packed(&pixels, Unstoppable).unwrap();
2054
2055        let jpeg = enc.finish().unwrap();
2056
2057        // Find XMP (APP1 with XMP namespace)
2058        let xmp_ns = b"http://ns.adobe.com/xap/1.0/\0";
2059        let mut found_xmp = false;
2060        let mut pos = 2;
2061        while pos + 4 < jpeg.len() {
2062            if jpeg[pos] == 0xFF && jpeg[pos + 1] == 0xE1 {
2063                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
2064                if jpeg.len() > pos + 4 + xmp_ns.len()
2065                    && &jpeg[pos + 4..pos + 4 + xmp_ns.len()] == xmp_ns
2066                {
2067                    found_xmp = true;
2068                    // Verify XMP content follows namespace
2069                    let xmp_start = pos + 4 + xmp_ns.len();
2070                    let xmp_end = pos + 2 + len;
2071                    if xmp_end <= jpeg.len() {
2072                        let xmp_data = &jpeg[xmp_start..xmp_end];
2073                        assert!(
2074                            xmp_data.starts_with(b"<?xml"),
2075                            "XMP data should start with XML declaration"
2076                        );
2077                    }
2078                    break;
2079                }
2080            }
2081            if jpeg[pos] == 0xFF && jpeg[pos + 1] != 0x00 && jpeg[pos + 1] != 0xFF {
2082                let len = ((jpeg[pos + 2] as usize) << 8) | (jpeg[pos + 3] as usize);
2083                pos += 2 + len;
2084            } else {
2085                pos += 1;
2086            }
2087        }
2088        assert!(found_xmp, "XMP segment not found in output");
2089    }
2090}