Skip to main content

blit_server/
surface_encoder.rs

1#![allow(clippy::too_many_arguments)]
2
3use blit_compositor::PixelData;
4use blit_remote::{
5    CODEC_SUPPORT_AV1, CODEC_SUPPORT_AV1_444, CODEC_SUPPORT_H264, CODEC_SUPPORT_H264_444,
6    SURFACE_FRAME_CODEC_AV1, SURFACE_FRAME_CODEC_H264,
7};
8use openh264::encoder::Encoder as OpenH264Encoder;
9use openh264::formats::YUVBuffer;
10
11#[derive(Clone, Copy, Debug, Eq, PartialEq)]
12pub enum SurfaceEncoderPreference {
13    VulkanVideoH264,
14    VulkanVideoAV1,
15    H264Software,
16    H264Vaapi,
17    AV1Vaapi,
18    NvencH264,
19    NvencAV1,
20    AV1Software,
21}
22
23// Type alias for backwards compatibility in tests.
24pub type SurfaceH264EncoderPreference = SurfaceEncoderPreference;
25
26/// openh264 hard limit: 3840x2160 horizontal or 2160x3840 vertical.
27const H264_MAX_WIDTH: u16 = 3840;
28const H264_MAX_HEIGHT: u16 = 2160;
29
30impl SurfaceEncoderPreference {
31    pub fn parse(value: &str) -> Option<Self> {
32        match value.trim() {
33            "h264-vulkan" => Some(Self::VulkanVideoH264),
34            "av1-vulkan" => Some(Self::VulkanVideoAV1),
35            "h264-software" | "software" => Some(Self::H264Software),
36            "h264-vaapi" | "vaapi" => Some(Self::H264Vaapi),
37            "av1-vaapi" => Some(Self::AV1Vaapi),
38            "h264-nvenc" => Some(Self::NvencH264),
39            "av1-nvenc" => Some(Self::NvencAV1),
40            "av1-software" => Some(Self::AV1Software),
41            _ => None,
42        }
43    }
44
45    /// Parse a comma-separated list of encoder preferences.
46    pub fn parse_list(value: &str) -> Result<Vec<Self>, String> {
47        let mut result = Vec::new();
48        for item in value.split(',') {
49            let item = item.trim();
50            if item.is_empty() {
51                continue;
52            }
53            result.push(Self::parse(item).ok_or_else(|| format!("unknown encoder: {item}"))?);
54        }
55        Ok(result)
56    }
57
58    /// Sensible default: hardware before software, NVENC preferred.
59    ///
60    /// Override at runtime with `BLIT_SURFACE_ENCODERS=h264-nvenc,h264-software`
61    /// (comma-separated list).
62    pub fn defaults() -> Vec<Self> {
63        if let Some(list) = std::env::var("BLIT_SURFACE_ENCODERS")
64            .ok()
65            .and_then(|v| Self::parse_list(&v).ok())
66        {
67            return list;
68        }
69        vec![
70            // Vulkan Video encoders are not yet stable — enable via
71            // BLIT_SURFACE_ENCODERS=av1-vulkan,h264-vulkan,...
72            // Self::VulkanVideoAV1,
73            // Self::VulkanVideoH264,
74            Self::NvencAV1,
75            Self::NvencH264,
76            Self::AV1Vaapi,
77            Self::H264Vaapi,
78            Self::H264Software,
79            Self::AV1Software,
80        ]
81    }
82
83    /// Returns true if the given codec_support bitmask allows this encoder.
84    /// A codec_support of 0 means "accept anything".
85    pub fn supported_by_client(self, codec_support: u8) -> bool {
86        if codec_support == 0 {
87            return true;
88        }
89        match self {
90            Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
91                codec_support & CODEC_SUPPORT_H264 != 0
92            }
93            Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
94                codec_support & CODEC_SUPPORT_AV1 != 0
95            }
96        }
97    }
98
99    /// Returns true if the client announced 4:4:4 chroma support for this
100    /// encoder's codec family.  Legacy clients (codec_support == 0) are
101    /// assumed to lack 4:4:4 support since the resulting Professional Profile
102    /// bitstreams are not universally decodable.
103    pub fn supports_444_by_client(self, codec_support: u8) -> bool {
104        if codec_support == 0 {
105            return false;
106        }
107        match self {
108            Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
109                codec_support & CODEC_SUPPORT_H264_444 != 0
110            }
111            Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
112                codec_support & CODEC_SUPPORT_AV1_444 != 0
113            }
114        }
115    }
116
117    /// Maximum surface dimensions the encoder can handle.
118    /// Returns `None` if there is no practical limit.
119    pub fn max_dimensions(self) -> Option<(u16, u16)> {
120        match self {
121            Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
122                Some((H264_MAX_WIDTH, H264_MAX_HEIGHT))
123            }
124            Self::AV1Vaapi | Self::NvencAV1 | Self::AV1Software | Self::VulkanVideoAV1 => None,
125        }
126    }
127
128    /// Tightest max dimensions across a list of preferences.
129    pub fn max_dimensions_for_list(prefs: &[Self]) -> Option<(u16, u16)> {
130        let mut result: Option<(u16, u16)> = None;
131        for p in prefs {
132            if let Some((w, h)) = p.max_dimensions() {
133                result = Some(match result {
134                    Some((rw, rh)) => (rw.min(w), rh.min(h)),
135                    None => (w, h),
136                });
137            }
138        }
139        result
140    }
141
142    /// Whether this encoder runs in the compositor via Vulkan Video.
143    pub fn is_vulkan_video(self) -> bool {
144        matches!(self, Self::VulkanVideoH264 | Self::VulkanVideoAV1)
145    }
146
147    /// Vulkan Video codec byte: 0x01 = H.264, 0x02 = AV1.
148    pub fn vulkan_codec(self) -> u8 {
149        match self {
150            Self::VulkanVideoAV1 => 0x02,
151            _ => 0x01,
152        }
153    }
154
155    /// Codec flag matching `SURFACE_FRAME_CODEC_*` constants.
156    pub fn codec_flag(self) -> u8 {
157        match self {
158            Self::H264Software | Self::H264Vaapi | Self::NvencH264 | Self::VulkanVideoH264 => {
159                SURFACE_FRAME_CODEC_H264
160            }
161            Self::AV1Vaapi | Self::AV1Software | Self::NvencAV1 | Self::VulkanVideoAV1 => {
162                SURFACE_FRAME_CODEC_AV1
163            }
164        }
165    }
166}
167
168/// Chroma subsampling mode.
169///
170/// - **Cs420** (default): 4:2:0 — U/V at half horizontal and half vertical
171///   resolution.  Universally supported, lower bandwidth.
172/// - **Cs444**: 4:4:4 — full-resolution chroma.  Eliminates colour fringing
173///   on sharp edges (ideal for text / UI), but requires encoder support.
174///
175/// Set via `BLIT_CHROMA` env var. Default: 444 (fall back to 420 if unsupported).
176/// Use `BLIT_CHROMA=420` to force 4:2:0.
177#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
178pub enum ChromaSubsampling {
179    Cs420,
180    #[default]
181    Cs444,
182}
183
184impl ChromaSubsampling {
185    pub fn label(self) -> &'static str {
186        match self {
187            Self::Cs420 => "4:2:0",
188            Self::Cs444 => "4:4:4",
189        }
190    }
191
192    pub fn from_env() -> Self {
193        match std::env::var("BLIT_CHROMA").ok().as_deref() {
194            Some("420") => Self::Cs420,
195            _ => Self::Cs444,
196        }
197    }
198
199    pub fn is_444(self) -> bool {
200        matches!(self, Self::Cs444)
201    }
202}
203
204/// Compute the AV1 level index string (e.g. "05") for the given dimensions,
205/// assuming 60 fps.  Mirrors the client-side `av1LevelString()`.
206pub fn av1_level_for(width: u32, height: u32) -> &'static str {
207    let sps = width as u64 * height as u64 * 60;
208    // (level_string, max_w, max_h, max_sample_rate)
209    const SPECS: &[(&str, u32, u32, u64)] = &[
210        ("00", 2048, 1152, 5_529_600),
211        ("01", 2816, 1152, 10_454_400),
212        ("04", 4352, 2448, 24_969_600),
213        ("05", 5504, 3096, 39_938_400),
214        ("08", 6144, 3456, 77_856_768),
215        ("09", 6144, 3456, 155_713_536),
216        ("12", 8192, 4352, 273_715_200),
217        ("13", 8192, 4352, 547_430_400),
218        ("16", 16384, 8704, 1_176_502_272),
219    ];
220    for &(level, max_w, max_h, max_rate) in SPECS {
221        if width <= max_w && height <= max_h && sps <= max_rate {
222            return level;
223        }
224    }
225    "16"
226}
227
228/// Video quality preset.  Higher quality uses more CPU / bandwidth.
229///
230/// - **Low**: speed 10, quantizer 180 — minimal CPU, visibly lossy
231/// - **Medium** (default): speed 10, quantizer 120 — good balance
232/// - **High**: speed 8, quantizer 80 — sharp, noticeable CPU use
233/// - **Ultra**: speed 6, quantizer 1 — near-lossless, heavy CPU
234/// - **Custom**: caller-specified AV1 quantizer (10–255)
235///
236/// Set via `BLIT_SURFACE_QUALITY=low|medium|high|ultra`.
237#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
238pub enum SurfaceQuality {
239    Low,
240    #[default]
241    Medium,
242    High,
243    Ultra,
244    /// Caller-specified AV1 quantizer (10–255).  H.264 QP, encoder speed,
245    /// and software-encoder bitrate are derived proportionally.
246    Custom {
247        quantizer: u8,
248    },
249}
250
251impl SurfaceQuality {
252    pub fn parse(value: &str) -> Option<Self> {
253        match value {
254            "low" => Some(Self::Low),
255            "medium" => Some(Self::Medium),
256            "high" => Some(Self::High),
257            "ultra" | "lossless" => Some(Self::Ultra),
258            _ => None,
259        }
260    }
261
262    /// Decode from the wire `quality` byte in C2S_SURFACE_SUBSCRIBE.
263    ///
264    /// - 0 → `None` (server default)
265    /// - 1–4 → named presets
266    /// - 10–255 → `Custom { quantizer: value }`
267    /// - 5–9 → reserved, treated as server default
268    pub fn from_wire(value: u8) -> Option<Self> {
269        match value {
270            1 => Some(Self::Low),
271            2 => Some(Self::Medium),
272            3 => Some(Self::High),
273            4 => Some(Self::Ultra),
274            v @ 10..=255 => Some(Self::Custom { quantizer: v }),
275            _ => None,
276        }
277    }
278
279    /// rav1e speed preset (0 = slowest/best, 10 = fastest/worst).
280    fn av1_speed(self) -> u8 {
281        match self {
282            Self::Low => 10,
283            Self::Medium => 10,
284            Self::High => 8,
285            Self::Ultra => 6,
286            Self::Custom { quantizer } => {
287                if quantizer <= 40 {
288                    6
289                } else if quantizer <= 80 {
290                    8
291                } else {
292                    10
293                }
294            }
295        }
296    }
297
298    /// AV1 quantizer (0 = lossless, 255 = worst).
299    /// Also used as VA-API `base_qindex` and NVENC AV1 QP.
300    fn av1_quantizer(self) -> usize {
301        match self {
302            Self::Low => 180,
303            Self::Medium => 120,
304            Self::High => 80,
305            Self::Ultra => 1,
306            Self::Custom { quantizer } => quantizer as usize,
307        }
308    }
309
310    /// rav1e min_quantizer — floor the encoder is allowed to improve to.
311    fn av1_min_quantizer(self) -> u8 {
312        match self {
313            Self::Low => 120,
314            Self::Medium => 80,
315            Self::High => 40,
316            Self::Ultra => 0,
317            Self::Custom { quantizer } => quantizer.saturating_sub(40),
318        }
319    }
320
321    /// H.264 QP for constant-quality mode (0 = best, 51 = worst).
322    /// Used by NVENC H.264 and VA-API H.264.
323    pub fn h264_qp(self) -> u8 {
324        match self {
325            Self::Low => 35,
326            Self::Medium => 28,
327            Self::High => 20,
328            Self::Ultra => 10,
329            Self::Custom { quantizer } => ((quantizer as u32) * 51 / 255).min(51) as u8,
330        }
331    }
332
333    /// NVENC AV1 QP for constant-quality mode (0 = best, 255 = worst).
334    /// Same scale as `av1_quantizer` / VA-API `base_qindex`.
335    pub fn nvenc_av1_qp(self) -> u32 {
336        self.av1_quantizer() as u32
337    }
338
339    /// AV1 QP for Vulkan Video encode (0 = best, 255 = worst).
340    /// Same base_qindex scale as VA-API / NVENC.
341    pub fn av1_qp_for_vulkan(self) -> u8 {
342        self.av1_quantizer().min(255) as u8
343    }
344
345    /// openh264 target bitrate in bits/sec.  Resolution-independent
346    /// approximation — openh264 adapts internally.
347    fn openh264_bitrate(self) -> u32 {
348        match self {
349            Self::Low => 500_000,
350            Self::Medium => 2_000_000,
351            Self::High => 8_000_000,
352            Self::Ultra => 20_000_000,
353            Self::Custom { quantizer } => {
354                // Linear interpolation: quantizer 0 → 20 Mbps, 255 → 500 kbps.
355                let q = quantizer as u32;
356                20_000_000 - q * (20_000_000 - 500_000) / 255
357            }
358        }
359    }
360}
361
362pub struct SurfaceEncoder {
363    /// Dimensions the encoder actually operates at (may be padded to even for H.264).
364    width: u32,
365    height: u32,
366    /// Original surface dimensions before any padding.
367    source_width: u32,
368    source_height: u32,
369    kind: SurfaceEncoderKind,
370    /// Negotiated chroma subsampling (may differ from requested if backend
371    /// does not support 4:4:4).
372    chroma: ChromaSubsampling,
373}
374
375enum SurfaceEncoderKind {
376    H264Software(Box<SoftwareH264Encoder>),
377    NvencH264(Box<crate::nvenc_encode::NvencDirectEncoder>),
378    NvencAV1(Box<crate::nvenc_encode::NvencDirectEncoder>),
379    #[cfg(target_os = "linux")]
380    H264Vaapi(Box<crate::vaapi_encode::VaapiDirectEncoder>),
381    #[cfg(target_os = "linux")]
382    AV1Vaapi(Box<crate::vaapi_encode::VaapiAv1Encoder>),
383    AV1Software(Box<SoftwareAV1Encoder>),
384}
385
386impl SurfaceEncoder {
387    /// Try each preference in order; return the first that succeeds and
388    /// the client can decode.  `codec_support` is a bitmask of
389    /// `CODEC_SUPPORT_*` (0 = accept anything).
390    pub fn new(
391        preferences: &[SurfaceEncoderPreference],
392        width: u32,
393        height: u32,
394        vaapi_device: &str,
395        quality: SurfaceQuality,
396        verbose: bool,
397        codec_support: u8,
398        chroma: ChromaSubsampling,
399    ) -> Result<Self, String> {
400        let source_width = width;
401        let source_height = height;
402        let mut last_err = String::from("no encoders configured");
403
404        // Single pass: for each encoder preference, try 4:4:4 first
405        // (if requested and client-supported), then fall back to 4:2:0,
406        // before moving to the next encoder.  This ensures e.g.
407        // h264-software 4:2:0 beats av1-software 4:4:4.
408        let try_444 = chroma.is_444();
409        if try_444 && verbose {
410            eprintln!(
411                "[surface-encoder] 4:4:4 eligible: codec_support={codec_support:#04x} for {source_width}x{source_height}",
412            );
413        }
414
415        for &pref in preferences {
416            if pref.is_vulkan_video() {
417                continue;
418            }
419            if !pref.supported_by_client(codec_support) {
420                continue;
421            }
422
423            // Try 4:4:4 first for this encoder if the client supports it.
424            if try_444 && pref.supports_444_by_client(codec_support) {
425                match Self::try_one(
426                    pref,
427                    width,
428                    height,
429                    source_width,
430                    source_height,
431                    vaapi_device,
432                    quality,
433                    verbose,
434                    ChromaSubsampling::Cs444,
435                ) {
436                    Ok(enc) => {
437                        if verbose {
438                            eprintln!(
439                                "[surface-encoder] using {:?} 4:4:4 for {source_width}x{source_height}",
440                                pref
441                            );
442                        }
443                        return Ok(enc);
444                    }
445                    Err(err) => {
446                        if verbose {
447                            eprintln!(
448                                "[surface-encoder] {:?} 4:4:4 unavailable for {source_width}x{source_height}: {err}",
449                                pref
450                            );
451                        }
452                        // The 4:2:0 fallback below will overwrite last_err
453                        // on failure; no need to record this one.
454                    }
455                }
456            }
457
458            // Fall back to 4:2:0 for this encoder.
459            match Self::try_one(
460                pref,
461                width,
462                height,
463                source_width,
464                source_height,
465                vaapi_device,
466                quality,
467                verbose,
468                ChromaSubsampling::Cs420,
469            ) {
470                Ok(enc) => {
471                    if verbose {
472                        eprintln!(
473                            "[surface-encoder] using {:?} 4:2:0 for {source_width}x{source_height}",
474                            pref
475                        );
476                    }
477                    return Ok(enc);
478                }
479                Err(err) => {
480                    if verbose {
481                        eprintln!(
482                            "[surface-encoder] {:?} 4:2:0 unavailable for {source_width}x{source_height}: {err}",
483                            pref
484                        );
485                    }
486                    last_err = err;
487                }
488            }
489        }
490        Err(last_err)
491    }
492
493    fn try_one(
494        pref: SurfaceEncoderPreference,
495        width: u32,
496        height: u32,
497        source_width: u32,
498        source_height: u32,
499        vaapi_device: &str,
500        quality: SurfaceQuality,
501        verbose: bool,
502        chroma: ChromaSubsampling,
503    ) -> Result<Self, String> {
504        let _ = vaapi_device;
505        validate_surface_dimensions(width, height, pref)?;
506
507        match pref {
508            SurfaceEncoderPreference::VulkanVideoH264
509            | SurfaceEncoderPreference::VulkanVideoAV1 => {
510                Err("Vulkan Video encoders are managed by the compositor".into())
511            }
512            SurfaceEncoderPreference::NvencH264 => {
513                let (width, height) = ((width + 1) & !1, (height + 1) & !1);
514                let qp = quality.h264_qp() as u32;
515                Ok(Self {
516                    width,
517                    height,
518                    source_width,
519                    source_height,
520                    kind: SurfaceEncoderKind::NvencH264(Box::new(
521                        crate::nvenc_encode::NvencDirectEncoder::try_new(
522                            "h264", width, height, qp, verbose, chroma,
523                        )?,
524                    )),
525                    chroma,
526                })
527            }
528            SurfaceEncoderPreference::NvencAV1 => {
529                // AV1 superblocks are 64x64; NVENC requires even dimensions
530                // at minimum.  Round up to a multiple of 2 (matching H.264)
531                // so chroma planes stay aligned.
532                let (width, height) = ((width + 1) & !1, (height + 1) & !1);
533                let qp = quality.nvenc_av1_qp();
534                Ok(Self {
535                    width,
536                    height,
537                    source_width,
538                    source_height,
539                    kind: SurfaceEncoderKind::NvencAV1(Box::new(
540                        crate::nvenc_encode::NvencDirectEncoder::try_new(
541                            "av1", width, height, qp, verbose, chroma,
542                        )?,
543                    )),
544                    chroma,
545                })
546            }
547            #[cfg(target_os = "linux")]
548            SurfaceEncoderPreference::H264Vaapi => {
549                let (width, height) = ((width + 1) & !1, (height + 1) & !1);
550                Ok(Self {
551                    width,
552                    height,
553                    source_width,
554                    source_height,
555                    kind: SurfaceEncoderKind::H264Vaapi(Box::new(
556                        crate::vaapi_encode::VaapiDirectEncoder::try_new(
557                            width,
558                            height,
559                            vaapi_device,
560                            quality.h264_qp(),
561                            verbose,
562                            chroma,
563                        )?,
564                    )),
565                    chroma,
566                })
567            }
568            #[cfg(not(target_os = "linux"))]
569            SurfaceEncoderPreference::H264Vaapi => Err("VA-API is only available on Unix".into()),
570            #[cfg(target_os = "linux")]
571            SurfaceEncoderPreference::AV1Vaapi => {
572                // Round up to 64-pixel superblocks.  AMD's AV1 backend
573                // rejects `vaCreateContext` with VA_STATUS_ERROR_
574                // RESOLUTION_NOT_SUPPORTED (0x13) below a codec-
575                // specific minimum; 256 isn't enough on many chips, so
576                // we default to 512.  AV1's `render_width/height` in
577                // the frame header still carries the actual
578                // `source_width/source_height`, so the client's
579                // WebCodecs decoder crops back to the requested size.
580                //
581                // BUT: the encoder encodes every pixel including the
582                // padding.  For small thumbnails, the padded area can
583                // exceed the source by >10×, turning a bandwidth-
584                // saving thumbnail into a bandwidth *amplifier*.  Bail
585                // out when padding would waste more than the content
586                // area — the fallback encoder chain picks a smaller-
587                // friendly backend (e.g. H264Software) instead.
588                const VAAPI_AV1_MIN: u32 = 512;
589                let enc_w = width.div_ceil(64) * 64;
590                let enc_h = height.div_ceil(64) * 64;
591                let (width, height) = (enc_w.max(VAAPI_AV1_MIN), enc_h.max(VAAPI_AV1_MIN));
592                let source_area = (source_width as u64) * (source_height as u64);
593                let padded_area = (width as u64) * (height as u64);
594                if padded_area > source_area.saturating_mul(2) {
595                    return Err(format!(
596                        "AV1Vaapi padding {width}x{height} > 2× source \
597                         {source_width}x{source_height} — falling back",
598                    ));
599                }
600                Ok(Self {
601                    width,
602                    height,
603                    source_width,
604                    source_height,
605                    kind: SurfaceEncoderKind::AV1Vaapi(Box::new(
606                        crate::vaapi_encode::VaapiAv1Encoder::try_new(
607                            width,
608                            height,
609                            source_width,
610                            source_height,
611                            vaapi_device,
612                            quality.av1_quantizer() as u8,
613                            verbose,
614                            chroma,
615                        )?,
616                    )),
617                    chroma,
618                })
619            }
620            #[cfg(not(target_os = "linux"))]
621            SurfaceEncoderPreference::AV1Vaapi => Err("VA-API is only available on Linux".into()),
622            SurfaceEncoderPreference::AV1Software => Ok(Self {
623                width,
624                height,
625                source_width,
626                source_height,
627                kind: SurfaceEncoderKind::AV1Software(Box::new(SoftwareAV1Encoder::new(
628                    width, height, quality, chroma,
629                )?)),
630                chroma,
631            }),
632            SurfaceEncoderPreference::H264Software => {
633                if chroma.is_444() {
634                    return Err("openh264 does not support 4:4:4".into());
635                }
636                let (width, height) = ((width + 1) & !1, (height + 1) & !1);
637                Ok(Self {
638                    width,
639                    height,
640                    source_width,
641                    source_height,
642                    kind: SurfaceEncoderKind::H264Software(Box::new(SoftwareH264Encoder::new(
643                        quality,
644                    )?)),
645                    chroma,
646                })
647            }
648        }
649    }
650
651    /// The original surface dimensions before any encoder padding.
652    pub fn source_dimensions(&self) -> (u32, u32) {
653        (self.source_width, self.source_height)
654    }
655
656    /// The encoder's padded dimensions (may be larger than source due to
657    /// alignment requirements, e.g. AV1 64-pixel superblock alignment).
658    pub fn encoder_dimensions(&self) -> (u32, u32) {
659        (self.width, self.height)
660    }
661
662    /// Human-readable name of the active encoder backend, sent to clients
663    /// for display in debug panels.  Includes chroma subsampling when 4:4:4.
664    pub fn encoder_name(&self) -> &'static str {
665        match (&self.kind, self.chroma) {
666            (SurfaceEncoderKind::H264Software(_), _) => "h264-software",
667            (SurfaceEncoderKind::NvencH264(_), ChromaSubsampling::Cs444) => "h264-nvenc 4:4:4",
668            (SurfaceEncoderKind::NvencH264(_), _) => "h264-nvenc",
669            (SurfaceEncoderKind::NvencAV1(_), ChromaSubsampling::Cs444) => "av1-nvenc 4:4:4",
670            (SurfaceEncoderKind::NvencAV1(_), _) => "av1-nvenc",
671            #[cfg(target_os = "linux")]
672            (SurfaceEncoderKind::H264Vaapi(_), ChromaSubsampling::Cs444) => "h264-vaapi 4:4:4",
673            #[cfg(target_os = "linux")]
674            (SurfaceEncoderKind::H264Vaapi(_), _) => "h264-vaapi",
675            #[cfg(target_os = "linux")]
676            (SurfaceEncoderKind::AV1Vaapi(_), ChromaSubsampling::Cs444) => "av1-vaapi 4:4:4",
677            #[cfg(target_os = "linux")]
678            (SurfaceEncoderKind::AV1Vaapi(_), _) => "av1-vaapi",
679            (SurfaceEncoderKind::AV1Software(_), ChromaSubsampling::Cs444) => "av1-software 4:4:4",
680            (SurfaceEncoderKind::AV1Software(_), _) => "av1-software",
681        }
682    }
683
684    /// WebCodecs codec string for the active encoder.  Sent to the client
685    /// so it can configure `VideoDecoder` with the correct profile/level.
686    pub fn webcodecs_codec_string(&self) -> String {
687        match &self.kind {
688            SurfaceEncoderKind::H264Software(_) => {
689                if self.chroma.is_444() {
690                    "avc1.F4001f".to_string()
691                } else {
692                    "avc1.42001f".to_string()
693                }
694            }
695            #[cfg(target_os = "linux")]
696            SurfaceEncoderKind::H264Vaapi(_) => {
697                if self.chroma.is_444() {
698                    "avc1.F4001f".to_string()
699                } else {
700                    "avc1.640034".to_string()
701                }
702            }
703            SurfaceEncoderKind::NvencH264(_) => "avc1.640034".to_string(),
704            SurfaceEncoderKind::NvencAV1(_) | SurfaceEncoderKind::AV1Software(_) => {
705                let profile = if self.chroma.is_444() { 2 } else { 0 };
706                let level = av1_level_for(self.source_width, self.source_height);
707                format!("av01.{profile}.{level}M.08")
708            }
709            #[cfg(target_os = "linux")]
710            SurfaceEncoderKind::AV1Vaapi(_) => {
711                let profile = if self.chroma.is_444() { 2 } else { 0 };
712                let level = av1_level_for(self.source_width, self.source_height);
713                format!("av01.{profile}.{level}M.08")
714            }
715        }
716    }
717
718    pub fn codec_flag(&self) -> u8 {
719        match &self.kind {
720            SurfaceEncoderKind::H264Software(_) => SURFACE_FRAME_CODEC_H264,
721            #[cfg(target_os = "linux")]
722            SurfaceEncoderKind::H264Vaapi(_) => SURFACE_FRAME_CODEC_H264,
723            SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
724                enc.codec_flag()
725            }
726            #[cfg(target_os = "linux")]
727            SurfaceEncoderKind::AV1Vaapi(_) => SURFACE_FRAME_CODEC_AV1,
728            SurfaceEncoderKind::AV1Software(_) => SURFACE_FRAME_CODEC_AV1,
729        }
730    }
731
732    pub fn request_keyframe(&mut self) {
733        match &mut self.kind {
734            SurfaceEncoderKind::H264Software(enc) => enc.request_keyframe(),
735            SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
736                enc.request_keyframe()
737            }
738            #[cfg(target_os = "linux")]
739            SurfaceEncoderKind::H264Vaapi(enc) => enc.request_keyframe(),
740            #[cfg(target_os = "linux")]
741            SurfaceEncoderKind::AV1Vaapi(enc) => enc.request_keyframe(),
742            SurfaceEncoderKind::AV1Software(enc) => enc.request_keyframe(),
743        }
744    }
745
746    /// Get GBM-allocated LINEAR BGRA buffers for zero-copy compositor→encoder.
747    #[cfg(target_os = "linux")]
748    pub fn gbm_buffers(&self) -> &[crate::vaapi_encode::GbmExportedBuffer] {
749        match &self.kind {
750            SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_buffers(),
751            SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_buffers(),
752            _ => &[],
753        }
754    }
755
756    #[cfg(target_os = "linux")]
757    pub fn gbm_nv12_buffers(&self) -> &[crate::vaapi_encode::GbmNv12Buffer] {
758        match &self.kind {
759            SurfaceEncoderKind::H264Vaapi(enc) => enc.gbm_nv12_buffers(),
760            SurfaceEncoderKind::AV1Vaapi(enc) => enc.gbm_nv12_buffers(),
761            _ => &[],
762        }
763    }
764
765    #[cfg(target_os = "linux")]
766    pub fn allocate_nv12_buffers(&mut self, drm_fd: std::os::fd::RawFd, count: usize) {
767        match &mut self.kind {
768            SurfaceEncoderKind::H264Vaapi(enc) => {
769                if let Some(vpp) = &mut enc.vpp {
770                    vpp.allocate_nv12_buffers(drm_fd, count);
771                }
772            }
773            SurfaceEncoderKind::AV1Vaapi(enc) => {
774                if let Some(vpp) = &mut enc.vpp {
775                    vpp.allocate_nv12_buffers(drm_fd, count);
776                }
777            }
778            _ => {}
779        }
780    }
781
782    #[cfg(target_os = "linux")]
783    pub fn drm_fd_raw(&self) -> std::os::fd::RawFd {
784        use std::os::fd::AsRawFd;
785        match &self.kind {
786            SurfaceEncoderKind::H264Vaapi(enc) => enc._drm_fd.as_raw_fd(),
787            SurfaceEncoderKind::AV1Vaapi(enc) => enc._drm_fd.as_raw_fd(),
788            _ => -1,
789        }
790    }
791
792    /// Get VA display pointer (as usize).
793    #[cfg(target_os = "linux")]
794    #[allow(dead_code)]
795    pub fn va_display_usize(&self) -> usize {
796        match &self.kind {
797            SurfaceEncoderKind::H264Vaapi(enc) => enc.va_display_usize(),
798            SurfaceEncoderKind::AV1Vaapi(enc) => enc.va_display_usize(),
799            _ => 0,
800        }
801    }
802
803    pub fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
804        // NVENC handles RGBA→encoder-size padding internally in pinned
805        // GPU memory, so pass the original un-padded buffer with source
806        // dimensions.  The generic padding below produces enc_w stride
807        // which would cause a diagonal-skew artefact when
808        // encode_rgba_padded re-interprets it at src_w stride.
809        if let SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) =
810            &mut self.kind
811        {
812            let (sw, sh) = (self.source_width as usize, self.source_height as usize);
813            let mut result = enc.encode_rgba_padded(rgba, sw, sh);
814            self.fixup_keyframe(&mut result);
815            return result;
816        }
817
818        let enc_len = expected_rgba_len(self.width, self.height);
819        let enc_len = match enc_len {
820            Some(v) => v,
821            None => {
822                eprintln!(
823                    "[surface-encoder] expected_rgba_len overflow {}x{}",
824                    self.width, self.height
825                );
826                return None;
827            }
828        };
829        let rgba = if rgba.len() == enc_len {
830            std::borrow::Cow::Borrowed(rgba)
831        } else {
832            // The source buffer may be smaller when the original surface had
833            // odd dimensions (H.264 rounds up to even).  Pad with edge-pixel
834            // duplication.
835            let total_px = rgba.len() / 4;
836            if total_px == 0 {
837                return None;
838            }
839            // Infer source width: try self.width, then self.width - 1
840            let src_w = [self.width as usize, (self.width - 1) as usize]
841                .into_iter()
842                .find(|&w| w > 0 && total_px.is_multiple_of(w))?;
843            let src_h = total_px / src_w;
844            if src_h == 0 {
845                return None;
846            }
847            let dst_w = self.width as usize;
848            let dst_h = self.height as usize;
849            let mut padded = vec![0u8; enc_len];
850            for row in 0..dst_h {
851                let src_row = row.min(src_h - 1);
852                for col in 0..dst_w {
853                    let src_col = col.min(src_w - 1);
854                    let si = (src_row * src_w + src_col) * 4;
855                    let di = (row * dst_w + col) * 4;
856                    padded[di..di + 4].copy_from_slice(&rgba[si..si + 4]);
857                }
858            }
859            std::borrow::Cow::Owned(padded)
860        };
861
862        match &mut self.kind {
863            SurfaceEncoderKind::H264Software(encoder) => {
864                encoder.encode(&rgba, self.width, self.height)
865            }
866            // NVENC early-returned above.
867            SurfaceEncoderKind::NvencH264(_) | SurfaceEncoderKind::NvencAV1(_) => unreachable!(),
868            #[cfg(target_os = "linux")]
869            SurfaceEncoderKind::H264Vaapi(enc) => {
870                let mut bgra = rgba.into_owned();
871                for px in bgra.chunks_exact_mut(4) {
872                    px.swap(0, 2);
873                }
874                let (sw, sh) = (self.source_width as usize, self.source_height as usize);
875                enc.encode_bgra_padded(&bgra, sw, sh)
876            }
877            #[cfg(target_os = "linux")]
878            SurfaceEncoderKind::AV1Vaapi(enc) => {
879                let mut bgra = rgba.into_owned();
880                for px in bgra.chunks_exact_mut(4) {
881                    px.swap(0, 2);
882                }
883                let (sw, sh) = (self.source_width as usize, self.source_height as usize);
884                enc.encode_bgra_padded(&bgra, sw, sh)
885            }
886            SurfaceEncoderKind::AV1Software(encoder) => encoder.encode(&rgba),
887        }
888    }
889
890    /// Encode a frame from native pixel data (BGRA, NV12, RGBA, or DMA-BUF).
891    /// Dispatches to the most efficient path for each format.
892    pub fn encode_pixels(&mut self, pixels: &PixelData) -> Option<(Vec<u8>, bool)> {
893        match pixels {
894            PixelData::Nv12 {
895                data,
896                y_stride,
897                uv_stride,
898            } => self.encode_nv12(data, *y_stride, *uv_stride),
899            PixelData::Bgra(bgra) => self.encode_bgra(bgra),
900            PixelData::Rgba(rgba) => self.encode(rgba),
901            #[cfg(target_os = "linux")]
902            PixelData::DmaBuf {
903                fd,
904                fourcc,
905                modifier,
906                stride,
907                offset,
908                ..
909            } => self
910                .encode_dmabuf(fd, *fourcc, *modifier, *stride, *offset)
911                .or_else(|| {
912                    // DMA-BUF import failed (e.g. VAAPI can't import Vulkan
913                    // stride).  Fall back to CPU mmap + BGRA encode.
914                    let w = self.width;
915                    let h = self.height;
916                    let rgba = pixels.to_rgba(w, h);
917                    if !rgba.is_empty() {
918                        self.encode(&rgba)
919                    } else {
920                        None
921                    }
922                }),
923            #[cfg(not(target_os = "linux"))]
924            PixelData::DmaBuf { .. } => None,
925            #[cfg(target_os = "linux")]
926            PixelData::Nv12DmaBuf {
927                fd,
928                stride,
929                uv_offset,
930                width,
931                height,
932                sync_fd,
933            } => {
934                // If the compositor exported a sync_fd (tiled NV12 on radv),
935                // wait for the GPU to finish the BGRA→NV12 compute before
936                // reading.  This runs in spawn_blocking so blocking is fine.
937                if let Some(sfd) = sync_fd {
938                    use std::os::fd::AsRawFd;
939                    let mut pfd = libc::pollfd {
940                        fd: sfd.as_raw_fd(),
941                        events: libc::POLLIN,
942                        revents: 0,
943                    };
944                    unsafe { libc::poll(&mut pfd, 1, 5000) };
945                }
946                self.encode_nv12_dmabuf(fd, *stride, *uv_offset, *width, *height)
947            }
948            .or_else(|| {
949                // VA surface lookup failed — mmap the DMA-BUF and
950                // fall back to encode_nv12 (upload path).
951                use std::os::fd::AsRawFd;
952                let h = *height as usize;
953                let s = *stride as usize;
954                let uv_off = *uv_offset as usize;
955                let raw = fd.as_raw_fd();
956                let map_size = uv_off + s * h.div_ceil(2);
957                let ptr = unsafe {
958                    libc::mmap(
959                        std::ptr::null_mut(),
960                        map_size,
961                        libc::PROT_READ,
962                        libc::MAP_SHARED,
963                        raw,
964                        0,
965                    )
966                };
967                if ptr == libc::MAP_FAILED || ptr.is_null() {
968                    return None;
969                }
970                let data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_size) };
971                let result = self.encode_nv12(data, s, s);
972                unsafe { libc::munmap(ptr, map_size) };
973                result
974            }),
975            #[cfg(not(target_os = "linux"))]
976            PixelData::Nv12DmaBuf { .. } => None,
977            PixelData::VaSurface { .. } => None,
978            // Vulkan Video pre-encoded — should be handled before reaching
979            // SurfaceEncoder.  If it gets here, we can't re-encode.
980            PixelData::Encoded { .. } => None,
981        }
982    }
983
984    /// Encode from a VA-API-allocated NV12 surface (zero-copy).
985    /// The compute shader wrote NV12 into the exported DMA-BUF; we look up
986    /// the owning VA surface by inode and encode directly — no PRIME import.
987    #[cfg(target_os = "linux")]
988    fn encode_nv12_dmabuf(
989        &mut self,
990        fd: &std::sync::Arc<std::os::fd::OwnedFd>,
991        _stride: u32,
992        _uv_offset: u32,
993        _width: u32,
994        _height: u32,
995    ) -> Option<(Vec<u8>, bool)> {
996        use std::os::fd::AsRawFd;
997        let raw_fd = fd.as_raw_fd();
998        let find_surface = |nv12s: &[crate::vaapi_encode::GbmNv12Buffer]| -> Option<u32> {
999            let buf = nv12s.iter().find(|n| n.fd.as_raw_fd() == raw_fd)?;
1000            // va_surface==0 means GBM fallback — no direct encode, use mmap.
1001            if buf.va_surface == 0 {
1002                return None;
1003            }
1004            Some(buf.va_surface)
1005        };
1006        let mut result = match &mut self.kind {
1007            SurfaceEncoderKind::AV1Vaapi(enc) => {
1008                let surf = find_surface(enc.gbm_nv12_buffers())?;
1009                enc.encode_surface(surf)
1010            }
1011            SurfaceEncoderKind::H264Vaapi(enc) => {
1012                let surf = find_surface(enc.gbm_nv12_buffers())?;
1013                enc.encode_surface(surf)
1014            }
1015            _ => None,
1016        };
1017        self.fixup_keyframe(&mut result);
1018        result
1019    }
1020
1021    /// Encode from a DMA-BUF fd — tries zero-copy GPU import first,
1022    /// falls back to CPU mmap readback if no GPU path is available.
1023    #[cfg(target_os = "linux")]
1024    fn encode_dmabuf(
1025        &mut self,
1026        fd: &std::os::fd::OwnedFd,
1027        fourcc: u32,
1028        modifier: u64,
1029        stride: u32,
1030        offset: u32,
1031    ) -> Option<(Vec<u8>, bool)> {
1032        use std::os::fd::AsRawFd;
1033
1034        // The encoder's source dimensions match the DMA-BUF dimensions
1035        // (both come from last_pixels).
1036        let src_w = self.source_width;
1037        let src_h = self.source_height;
1038
1039        // --- Zero-copy GPU path (NVENC CUDA import) ---
1040        // VA-API encode uses the Nv12DmaBuf path instead (compute shader
1041        // writes NV12 into VA-API-exported surfaces, no PRIME import).
1042        let mut gpu_result = match &mut self.kind {
1043            SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => enc
1044                .encode_dmabuf_fd(
1045                    fd.as_raw_fd(),
1046                    fourcc,
1047                    modifier,
1048                    stride,
1049                    offset,
1050                    src_w,
1051                    src_h,
1052                ),
1053            _ => None,
1054        };
1055        if gpu_result.is_some() {
1056            self.fixup_keyframe(&mut gpu_result);
1057            return gpu_result;
1058        }
1059
1060        // --- CPU readback fallback ---
1061        // Only reached if zero-copy failed (VPP unavailable, or non-VA-API encoder).
1062        // The GBM BO is created with GBM_BO_USE_LINEAR so mmap reads
1063        // pixels in the correct linear layout.
1064        self.encode_dmabuf_cpu_fallback(fd, fourcc, stride, offset)
1065    }
1066
1067    /// CPU-side fallback for DMA-BUF encoding: mmap the fd, read pixels,
1068    /// and encode through the normal BGRA/NV12 path.
1069    #[cfg(target_os = "linux")]
1070    fn encode_dmabuf_cpu_fallback(
1071        &mut self,
1072        fd: &std::os::fd::OwnedFd,
1073        fourcc: u32,
1074        stride: u32,
1075        _offset: u32,
1076    ) -> Option<(Vec<u8>, bool)> {
1077        use std::os::fd::AsRawFd;
1078
1079        let w = self.source_width as usize;
1080        let h = self.source_height as usize;
1081        let stride = stride as usize;
1082        let raw_fd = fd.as_raw_fd();
1083
1084        // Determine total mmap size from fd (seek to end).
1085        let file_size = unsafe { libc::lseek(raw_fd, 0, libc::SEEK_END) };
1086        if file_size <= 0 {
1087            return None;
1088        }
1089        let map_len = file_size as usize;
1090
1091        // DMA-BUF sync: start read
1092        #[repr(C)]
1093        struct DmaBufSync {
1094            flags: u64,
1095        }
1096        const DMA_BUF_SYNC_READ: u64 = 1;
1097        const DMA_BUF_SYNC_START: u64 = 0;
1098        const DMA_BUF_SYNC_END: u64 = 4;
1099        // ioctl number for DMA_BUF_IOCTL_SYNC — use c_ulong and cast at
1100        // call sites so this works on both x86_64 (ioctl takes c_ulong)
1101        // and aarch64 (ioctl takes c_int).
1102        const DMA_BUF_IOCTL_SYNC: libc::c_ulong = 0x40086200;
1103
1104        // Use poll() to check if the DMA-BUF fence is ready before
1105        // attempting sync.  Anonymous /dmabuf: fds from Vulkan WSI may
1106        // have implicit GPU fences that block indefinitely on SYNC_START.
1107        {
1108            let mut pfd = libc::pollfd {
1109                fd: raw_fd,
1110                events: libc::POLLIN,
1111                revents: 0,
1112            };
1113            let ready = unsafe { libc::poll(&mut pfd, 1, 0) };
1114            if ready <= 0 {
1115                // Not ready — skip sync, accept possible tearing.
1116            } else {
1117                let sync_start = DmaBufSync {
1118                    flags: DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ,
1119                };
1120                unsafe {
1121                    libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_start);
1122                }
1123            }
1124        }
1125
1126        // mmap the DMA-BUF for reading.
1127        let ptr = unsafe {
1128            libc::mmap(
1129                std::ptr::null_mut(),
1130                map_len,
1131                libc::PROT_READ,
1132                libc::MAP_SHARED,
1133                raw_fd,
1134                0,
1135            )
1136        };
1137        if ptr == libc::MAP_FAILED {
1138            let sync_end = DmaBufSync {
1139                flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1140            };
1141            unsafe {
1142                libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1143            }
1144            return None;
1145        }
1146        let plane_data = unsafe { std::slice::from_raw_parts(ptr as *const u8, map_len) };
1147
1148        // Detect OpenGL FBO-backed DMA-BUFs (anonymous, not /dev/dri/).
1149        // These have bottom-up row order and must be flipped.
1150        let is_gl_fbo = {
1151            let mut link = [0u8; 128];
1152            let path = format!("/proc/self/fd/{raw_fd}\0");
1153            let n = unsafe {
1154                libc::readlink(path.as_ptr() as *const _, link.as_mut_ptr() as *mut _, 127)
1155            };
1156            !(n > 0 && link[..n as usize].starts_with(b"/dev/dri/"))
1157        };
1158
1159        let result = if fourcc == blit_compositor::drm_fourcc::ARGB8888
1160            || fourcc == blit_compositor::drm_fourcc::XRGB8888
1161        {
1162            // BGRA in memory.
1163            let mut packed = Vec::with_capacity(w * h * 4);
1164            for i in 0..h {
1165                // Flip row order for GL FBO buffers.
1166                let row = if is_gl_fbo { h - 1 - i } else { i };
1167                let start = row * stride;
1168                let end = start + w * 4;
1169                if end <= plane_data.len() {
1170                    packed.extend_from_slice(&plane_data[start..end]);
1171                }
1172            }
1173            self.encode_bgra(&packed)
1174        } else if fourcc == blit_compositor::drm_fourcc::ABGR8888
1175            || fourcc == blit_compositor::drm_fourcc::XBGR8888
1176        {
1177            // RGBA in memory.
1178            let mut packed = Vec::with_capacity(w * h * 4);
1179            for i in 0..h {
1180                let row = if is_gl_fbo { h - 1 - i } else { i };
1181                let start = row * stride;
1182                let end = start + w * 4;
1183                if end <= plane_data.len() {
1184                    packed.extend_from_slice(&plane_data[start..end]);
1185                }
1186            }
1187            self.encode(&packed)
1188        } else if fourcc == blit_compositor::drm_fourcc::NV12 {
1189            // NV12: Y plane at offset 0 with `stride` pitch, UV plane
1190            // immediately following at y_size offset with the same pitch.
1191            // For linear single-fd NV12 DMA-BUFs both planes are contiguous.
1192            let uv_stride = stride; // UV stride matches Y stride for linear NV12
1193            let y_size = stride * h;
1194            let uv_h = h.div_ceil(2);
1195            let uv_size = uv_stride * uv_h;
1196            if map_len >= y_size + uv_size {
1197                // Pack Y rows then UV rows tightly (strip stride padding).
1198                let out_stride = w;
1199                let mut data = vec![0u8; out_stride * h + out_stride * uv_h];
1200                for row in 0..h {
1201                    let src = row * stride;
1202                    let dst = row * out_stride;
1203                    if src + w <= plane_data.len() {
1204                        data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1205                    }
1206                }
1207                let uv_dst_base = out_stride * h;
1208                for row in 0..uv_h {
1209                    let src = y_size + row * uv_stride;
1210                    let dst = uv_dst_base + row * out_stride;
1211                    if src + w <= plane_data.len() {
1212                        data[dst..dst + w].copy_from_slice(&plane_data[src..src + w]);
1213                    }
1214                }
1215                self.encode_nv12(&data, out_stride, out_stride)
1216            } else {
1217                None
1218            }
1219        } else {
1220            None
1221        };
1222
1223        // Unmap and end sync.
1224        unsafe {
1225            libc::munmap(ptr, map_len);
1226        }
1227        // Only sync end if we did sync start (non-blocking check).
1228        let sync_end = DmaBufSync {
1229            flags: DMA_BUF_SYNC_END | DMA_BUF_SYNC_READ,
1230        };
1231        unsafe {
1232            libc::ioctl(raw_fd, DMA_BUF_IOCTL_SYNC as _, &sync_end);
1233        }
1234
1235        result
1236    }
1237
1238    /// Hardware encoders (NVENC, VA-API) may report the wrong picture type
1239    /// due to struct layout mismatches.  Re-detect from the bitstream as a
1240    /// cheap safety net.  This is applied to every encode path so that RGBA,
1241    /// BGRA, NV12, and DMA-BUF frames all get the same keyframe fixup.
1242    fn fixup_keyframe(&self, result: &mut Option<(Vec<u8>, bool)>) {
1243        if let Some((data, is_key)) = result.as_mut()
1244            && !*is_key
1245        {
1246            *is_key = match &self.kind {
1247                SurfaceEncoderKind::NvencH264(_) => h264_stream_contains_idr(data),
1248                SurfaceEncoderKind::NvencAV1(_) => av1_stream_contains_keyframe(data),
1249                #[cfg(target_os = "linux")]
1250                SurfaceEncoderKind::H264Vaapi(_) => h264_stream_contains_idr(data),
1251                #[cfg(target_os = "linux")]
1252                SurfaceEncoderKind::AV1Vaapi(_) => av1_stream_contains_keyframe(data),
1253                _ => false,
1254            };
1255        }
1256    }
1257
1258    /// Encode from BGRA pixels — converts directly to YUV, skipping RGBA.
1259    fn encode_bgra(&mut self, bgra: &[u8]) -> Option<(Vec<u8>, bool)> {
1260        let enc_w = self.width as usize;
1261        let enc_h = self.height as usize;
1262        let src_w = self.source_width as usize;
1263        let src_h = self.source_height as usize;
1264
1265        let mut result = match &mut self.kind {
1266            SurfaceEncoderKind::H264Software(encoder) => {
1267                let yuv = bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h);
1268                let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1269                encoder.encode_yuv(&yuv_buf, self.width, self.height)
1270            }
1271            SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1272                enc.encode_bgra_padded(bgra, src_w, src_h)
1273            }
1274            #[cfg(target_os = "linux")]
1275            SurfaceEncoderKind::H264Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1276            #[cfg(target_os = "linux")]
1277            SurfaceEncoderKind::AV1Vaapi(enc) => enc.encode_bgra_padded(bgra, src_w, src_h),
1278            SurfaceEncoderKind::AV1Software(encoder) => {
1279                let yuv = if self.chroma.is_444() {
1280                    bgra_to_yuv444_padded(bgra, src_w, src_h, enc_w, enc_h)
1281                } else {
1282                    bgra_to_yuv420_padded(bgra, src_w, src_h, enc_w, enc_h)
1283                };
1284                encoder.encode_yuv_planes(&yuv)
1285            }
1286        };
1287        self.fixup_keyframe(&mut result);
1288        result
1289    }
1290
1291    /// Encode from NV12 data — zero colorspace conversion for VA-API/NVENC,
1292    /// and only a deinterleave for software encoders.
1293    fn encode_nv12(
1294        &mut self,
1295        data: &[u8],
1296        y_stride: usize,
1297        uv_stride: usize,
1298    ) -> Option<(Vec<u8>, bool)> {
1299        // NV12 data was captured at source dimensions.
1300        let src_w = self.source_width as usize;
1301        let src_h = self.source_height as usize;
1302
1303        let mut result = match &mut self.kind {
1304            SurfaceEncoderKind::H264Software(encoder) => {
1305                let enc_w = self.width as usize;
1306                let enc_h = self.height as usize;
1307                if enc_w == src_w && enc_h == src_h {
1308                    let yuv = nv12_to_yuv420(data, y_stride, uv_stride, src_w, src_h);
1309                    let yuv_buf = YUVBuffer::from_vec(yuv, enc_w, enc_h);
1310                    encoder.encode_yuv(&yuv_buf, self.width, self.height)
1311                } else {
1312                    let pd = PixelData::Nv12 {
1313                        data: std::sync::Arc::new(data.to_vec()),
1314                        y_stride,
1315                        uv_stride,
1316                    };
1317                    let rgba = pd.to_rgba(self.source_width, self.source_height);
1318                    return self.encode(&rgba);
1319                }
1320            }
1321            SurfaceEncoderKind::NvencH264(enc) | SurfaceEncoderKind::NvencAV1(enc) => {
1322                // NVENC accepts NV12 natively — upload directly, no conversion.
1323                enc.encode_nv12(data, y_stride, uv_stride, src_h)
1324            }
1325            #[cfg(target_os = "linux")]
1326            SurfaceEncoderKind::H264Vaapi(enc) => {
1327                let uv_offset = y_stride * src_h;
1328                let y_data = &data[..uv_offset];
1329                let uv_data = &data[uv_offset..];
1330                enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1331            }
1332            #[cfg(target_os = "linux")]
1333            SurfaceEncoderKind::AV1Vaapi(enc) => {
1334                let uv_offset = y_stride * src_h;
1335                let y_data = &data[..uv_offset];
1336                let uv_data = &data[uv_offset..];
1337                enc.encode_nv12(y_data, uv_data, y_stride, uv_stride)
1338            }
1339            SurfaceEncoderKind::AV1Software(encoder) => {
1340                encoder.encode_nv12(data, y_stride, uv_stride, src_w, src_h)
1341            }
1342        };
1343        self.fixup_keyframe(&mut result);
1344        result
1345    }
1346}
1347
1348fn validate_surface_dimensions(
1349    width: u32,
1350    height: u32,
1351    _preference: SurfaceEncoderPreference,
1352) -> Result<(), String> {
1353    if width == 0 || height == 0 {
1354        return Err("surface encoder requires non-zero dimensions".into());
1355    }
1356    // Odd dimensions are fine — H.264 constructors pad to even internally,
1357    // and AV1/rav1e handles odd dimensions natively.
1358    let _ = expected_rgba_len(width, height)
1359        .ok_or_else(|| format!("surface encoder dimensions overflow for {width}x{height}"))?;
1360    Ok(())
1361}
1362
1363fn expected_rgba_len(width: u32, height: u32) -> Option<usize> {
1364    (width as usize)
1365        .checked_mul(height as usize)?
1366        .checked_mul(4)
1367}
1368
1369// ---------------------------------------------------------------------------
1370// Per-pixel math — #[inline(always)] so LLVM sees through the call in the
1371// hot loop and auto-vectorises the surrounding code.
1372// ---------------------------------------------------------------------------
1373
1374#[inline(always)]
1375fn rgb_to_y(r: i32, g: i32, b: i32) -> u8 {
1376    ((66 * r + 129 * g + 25 * b + 128) >> 8)
1377        .wrapping_add(16)
1378        .clamp(0, 255) as u8
1379}
1380
1381#[inline(always)]
1382fn rgb_to_u(r: i32, g: i32, b: i32) -> u8 {
1383    ((-38 * r - 74 * g + 112 * b + 128) >> 8)
1384        .wrapping_add(128)
1385        .clamp(0, 255) as u8
1386}
1387
1388#[inline(always)]
1389fn rgb_to_v(r: i32, g: i32, b: i32) -> u8 {
1390    ((112 * r - 94 * g - 18 * b + 128) >> 8)
1391        .wrapping_add(128)
1392        .clamp(0, 255) as u8
1393}
1394
1395// ---------------------------------------------------------------------------
1396// Bulk colorspace helpers — written for auto-vectorisation: flat pre-allocated
1397// output, direct indexing, no branches, no extend_from_slice.
1398// ---------------------------------------------------------------------------
1399
1400/// Flat Y-plane pass over packed 4-byte pixels.  `pixel_r/g/b` closures
1401/// extract R, G, B from the pixel at byte offset `i` (always a multiple of 4).
1402/// This is shared between RGBA, BGRA, and any other 4-byte packed format.
1403#[inline(always)]
1404fn compute_y_plane(
1405    src: &[u8],
1406    width: usize,
1407    height: usize,
1408    y_plane: &mut [u8],
1409    r_off: usize,
1410    g_off: usize,
1411    b_off: usize,
1412) {
1413    let total = width * height;
1414    for (px, y_out) in y_plane[..total].iter_mut().enumerate() {
1415        let i = px * 4;
1416        let r = src[i + r_off] as i32;
1417        let g = src[i + g_off] as i32;
1418        let b = src[i + b_off] as i32;
1419        *y_out = rgb_to_y(r, g, b);
1420    }
1421}
1422
1423/// Flat chroma pass (2x2 subsampling) over packed 4-byte pixels.
1424#[inline(always)]
1425fn compute_uv_planes(
1426    src: &[u8],
1427    width: usize,
1428    height: usize,
1429    u_plane: &mut [u8],
1430    v_plane: &mut [u8],
1431    r_off: usize,
1432    g_off: usize,
1433    b_off: usize,
1434) {
1435    let chroma_w = width.div_ceil(2);
1436    let chroma_h = height.div_ceil(2);
1437    for cy in 0..chroma_h {
1438        for cx in 0..chroma_w {
1439            let row = cy * 2;
1440            let col = cx * 2;
1441            // Average 2x2 block, clamping to source bounds for odd dims.
1442            let mut u_sum = 0i32;
1443            let mut v_sum = 0i32;
1444            for dy in 0..2u32 {
1445                for dx in 0..2u32 {
1446                    let sr = (row + dy as usize).min(height - 1);
1447                    let sc = (col + dx as usize).min(width - 1);
1448                    let i = (sr * width + sc) * 4;
1449                    let r = src[i + r_off] as i32;
1450                    let g = src[i + g_off] as i32;
1451                    let b = src[i + b_off] as i32;
1452                    u_sum += rgb_to_u(r, g, b) as i32;
1453                    v_sum += rgb_to_v(r, g, b) as i32;
1454                }
1455            }
1456            let idx = cy * chroma_w + cx;
1457            u_plane[idx] = (u_sum / 4) as u8;
1458            v_plane[idx] = (v_sum / 4) as u8;
1459        }
1460    }
1461}
1462
1463/// Padded Y-plane: produces `enc_w × enc_h` luma samples from a
1464/// `src_w × src_h` packed-pixel source, clamping coordinates to source bounds.
1465#[inline(always)]
1466fn compute_y_plane_padded(
1467    src: &[u8],
1468    src_w: usize,
1469    src_h: usize,
1470    enc_w: usize,
1471    enc_h: usize,
1472    y_plane: &mut [u8],
1473    r_off: usize,
1474    g_off: usize,
1475    b_off: usize,
1476) {
1477    for row in 0..enc_h {
1478        let sr = row.min(src_h - 1);
1479        for col in 0..enc_w {
1480            let sc = col.min(src_w - 1);
1481            let i = (sr * src_w + sc) * 4;
1482            let r = src[i + r_off] as i32;
1483            let g = src[i + g_off] as i32;
1484            let b = src[i + b_off] as i32;
1485            y_plane[row * enc_w + col] = rgb_to_y(r, g, b);
1486        }
1487    }
1488}
1489
1490/// Padded chroma planes: produces `ceil(enc_w/2) × ceil(enc_h/2)` chroma
1491/// samples with edge-pixel duplication for pixels beyond `src_w × src_h`.
1492#[inline(always)]
1493fn compute_uv_planes_padded(
1494    src: &[u8],
1495    src_w: usize,
1496    src_h: usize,
1497    enc_w: usize,
1498    enc_h: usize,
1499    u_plane: &mut [u8],
1500    v_plane: &mut [u8],
1501    r_off: usize,
1502    g_off: usize,
1503    b_off: usize,
1504) {
1505    let chroma_w = enc_w.div_ceil(2);
1506    let chroma_h = enc_h.div_ceil(2);
1507    for cy in 0..chroma_h {
1508        for cx in 0..chroma_w {
1509            let row = cy * 2;
1510            let col = cx * 2;
1511            let mut u_sum = 0i32;
1512            let mut v_sum = 0i32;
1513            for dy in 0..2u32 {
1514                for dx in 0..2u32 {
1515                    let sr = (row + dy as usize).min(src_h - 1);
1516                    let sc = (col + dx as usize).min(src_w - 1);
1517                    let i = (sr * src_w + sc) * 4;
1518                    let r = src[i + r_off] as i32;
1519                    let g = src[i + g_off] as i32;
1520                    let b = src[i + b_off] as i32;
1521                    u_sum += rgb_to_u(r, g, b) as i32;
1522                    v_sum += rgb_to_v(r, g, b) as i32;
1523                }
1524            }
1525            let idx = cy * chroma_w + cx;
1526            u_plane[idx] = (u_sum / 4) as u8;
1527            v_plane[idx] = (v_sum / 4) as u8;
1528        }
1529    }
1530}
1531
1532/// Compute full-resolution chroma planes (4:4:4) from packed 4-byte pixels
1533/// with edge-pixel padding to encoder dimensions.
1534#[inline(always)]
1535fn compute_uv_planes_444_padded(
1536    src: &[u8],
1537    src_w: usize,
1538    src_h: usize,
1539    enc_w: usize,
1540    enc_h: usize,
1541    u_plane: &mut [u8],
1542    v_plane: &mut [u8],
1543    r_off: usize,
1544    g_off: usize,
1545    b_off: usize,
1546) {
1547    for row in 0..enc_h {
1548        let sr = row.min(src_h - 1);
1549        for col in 0..enc_w {
1550            let sc = col.min(src_w - 1);
1551            let i = (sr * src_w + sc) * 4;
1552            let r = src[i + r_off] as i32;
1553            let g = src[i + g_off] as i32;
1554            let b = src[i + b_off] as i32;
1555            let idx = row * enc_w + col;
1556            u_plane[idx] = rgb_to_u(r, g, b);
1557            v_plane[idx] = rgb_to_v(r, g, b);
1558        }
1559    }
1560}
1561
1562/// BGRA -> I444 (YUV 4:4:4) with edge-pixel padding to encoder dimensions.
1563fn bgra_to_yuv444_padded(
1564    bgra: &[u8],
1565    src_w: usize,
1566    src_h: usize,
1567    enc_w: usize,
1568    enc_h: usize,
1569) -> Vec<u8> {
1570    let plane_size = enc_w * enc_h;
1571    let mut yuv = vec![0u8; plane_size * 3];
1572    let (y_plane, uv) = yuv.split_at_mut(plane_size);
1573    let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1574    // BGRA offsets: B=0, G=1, R=2, A=3
1575    compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1576    compute_uv_planes_444_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1577    yuv
1578}
1579
1580/// RGBA -> I444 (YUV 4:4:4).
1581fn rgba_to_yuv444(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1582    let plane_size = width * height;
1583    let mut yuv = vec![0u8; plane_size * 3];
1584    let (y_plane, uv) = yuv.split_at_mut(plane_size);
1585    let (u_plane, v_plane) = uv.split_at_mut(plane_size);
1586    // RGBA offsets: R=0, G=1, B=2, A=3
1587    compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1588    compute_uv_planes_444_padded(
1589        rgba, width, height, width, height, u_plane, v_plane, 0, 1, 2,
1590    );
1591    yuv
1592}
1593
1594/// BGRA -> I420 with edge-pixel padding to encoder dimensions.
1595/// `src_w × src_h` is the actual pixel count in `bgra`.
1596/// `enc_w × enc_h` is the encoder output dimensions (>= src).
1597fn bgra_to_yuv420_padded(
1598    bgra: &[u8],
1599    src_w: usize,
1600    src_h: usize,
1601    enc_w: usize,
1602    enc_h: usize,
1603) -> Vec<u8> {
1604    let y_size = enc_w * enc_h;
1605    // Use div_ceil to match encode_yuv_planes (rav1e) which expects
1606    // ceil(w/2) × ceil(h/2) chroma planes.  Truncating division produces
1607    // a short buffer when enc_w or enc_h is odd (AV1Software doesn't pad),
1608    // causing a panic in encode_yuv_planes's slice indexing.
1609    let uv_w = enc_w.div_ceil(2);
1610    let uv_size = uv_w * enc_h.div_ceil(2);
1611    let mut yuv = vec![0u8; y_size + uv_size * 2];
1612    let (y_plane, uv) = yuv.split_at_mut(y_size);
1613    let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1614    // BGRA offsets: B=0, G=1, R=2, A=3
1615    compute_y_plane_padded(bgra, src_w, src_h, enc_w, enc_h, y_plane, 2, 1, 0);
1616    compute_uv_planes_padded(bgra, src_w, src_h, enc_w, enc_h, u_plane, v_plane, 2, 1, 0);
1617    yuv
1618}
1619
1620/// RGBA -> I420 (Y + U + V planar).
1621fn rgba_to_yuv420(rgba: &[u8], width: usize, height: usize) -> Vec<u8> {
1622    let y_size = width * height;
1623    let uv_w = width.div_ceil(2);
1624    let uv_size = uv_w * height.div_ceil(2);
1625    let mut yuv = vec![0u8; y_size + uv_size * 2];
1626    let (y_plane, uv) = yuv.split_at_mut(y_size);
1627    let (u_plane, v_plane) = uv.split_at_mut(uv_size);
1628    // RGBA offsets: R=0, G=1, B=2, A=3
1629    compute_y_plane(rgba, width, height, y_plane, 0, 1, 2);
1630    compute_uv_planes(rgba, width, height, u_plane, v_plane, 0, 1, 2);
1631    yuv
1632}
1633
1634/// NV12 -> I420: Y plane memcpy + UV deinterleave.
1635/// Input: contiguous buffer with Y at data[..y_stride*height],
1636///        UV at data[y_stride*height..].
1637fn nv12_to_yuv420(
1638    data: &[u8],
1639    y_stride: usize,
1640    uv_stride: usize,
1641    width: usize,
1642    height: usize,
1643) -> Vec<u8> {
1644    let y_size = width * height;
1645    let uv_w = width.div_ceil(2);
1646    let uv_h = height.div_ceil(2);
1647    let uv_size = uv_w * uv_h;
1648    let mut yuv = vec![0u8; y_size + uv_size * 2];
1649    let (y_out, uv_out) = yuv.split_at_mut(y_size);
1650    let (u_out, v_out) = uv_out.split_at_mut(uv_size);
1651
1652    let uv_offset = y_stride * height;
1653
1654    // Copy Y plane (strip stride padding)
1655    for row in 0..height {
1656        let src = row * y_stride;
1657        let dst = row * width;
1658        y_out[dst..dst + width].copy_from_slice(&data[src..src + width]);
1659    }
1660
1661    // Deinterleave UV -> separate U, V.
1662    // uv_w may be one more than the source has (odd width), so clamp
1663    // to the number of pairs actually present in each source row.
1664    let src_uv_pairs = width / 2;
1665    for row in 0..uv_h {
1666        let src_start = uv_offset + row.min(height / 2 - 1) * uv_stride;
1667        let dst_start = row * uv_w;
1668        for col in 0..uv_w {
1669            let sc = col.min(src_uv_pairs.saturating_sub(1));
1670            u_out[dst_start + col] = data[src_start + sc * 2];
1671            v_out[dst_start + col] = data[src_start + sc * 2 + 1];
1672        }
1673    }
1674
1675    yuv
1676}
1677
1678/// Scan an Annex B H.264 bitstream for an IDR NAL unit (type 5).
1679fn h264_stream_contains_idr(data: &[u8]) -> bool {
1680    annex_b_contains_nal(data, |byte| (byte & 0x1f) == 5)
1681}
1682
1683/// Walk Annex B start codes and return true if any NAL's first byte satisfies `pred`.
1684fn annex_b_contains_nal(data: &[u8], pred: impl Fn(u8) -> bool) -> bool {
1685    let mut i = 0usize;
1686    while i < data.len() {
1687        let start_code_len = if data[i..].starts_with(&[0, 0, 0, 1]) {
1688            4
1689        } else if data[i..].starts_with(&[0, 0, 1]) {
1690            3
1691        } else {
1692            i += 1;
1693            continue;
1694        };
1695
1696        let nal_header = i + start_code_len;
1697        if let Some(&byte) = data.get(nal_header)
1698            && pred(byte)
1699        {
1700            return true;
1701        }
1702
1703        i = nal_header.saturating_add(1);
1704    }
1705
1706    false
1707}
1708
1709/// Check whether an AV1 OBU bitstream contains a sequence header, which
1710/// NVENC emits only for key frames.  This mirrors `h264_stream_contains_idr`
1711/// as a cheap bitstream-level safety net.
1712///
1713/// NVENC typically prepends a temporal delimiter OBU (type 2) before the
1714/// sequence header, so we must walk the OBU chain rather than only checking
1715/// the first byte.
1716fn av1_stream_contains_keyframe(data: &[u8]) -> bool {
1717    // OBU header byte: forbidden(1) | obu_type(4) | extension(1) | has_size(1) | reserved(1)
1718    // OBU types: 1 = SEQUENCE_HEADER, 2 = TEMPORAL_DELIMITER, 3 = FRAME_HEADER,
1719    //            6 = FRAME (header + tile data).
1720    let mut pos = 0;
1721    while pos < data.len() {
1722        let header = data[pos];
1723        let obu_type = (header >> 3) & 0xF;
1724        let has_extension = (header >> 2) & 1;
1725        let has_size = (header >> 1) & 1;
1726        pos += 1;
1727
1728        // Skip optional extension byte.
1729        if has_extension != 0 {
1730            if pos >= data.len() {
1731                break;
1732            }
1733            pos += 1;
1734        }
1735
1736        // OBU_SEQUENCE_HEADER → this is a key frame.
1737        if obu_type == 1 {
1738            return true;
1739        }
1740
1741        // If has_size is set, read the LEB128-encoded payload size and
1742        // skip past the OBU payload to inspect the next OBU.
1743        if has_size != 0 {
1744            let mut size: u64 = 0;
1745            let mut shift = 0u32;
1746            while pos < data.len() {
1747                let byte = data[pos];
1748                pos += 1;
1749                size |= ((byte & 0x7F) as u64) << shift;
1750                if byte & 0x80 == 0 {
1751                    break;
1752                }
1753                shift += 7;
1754                if shift >= 56 {
1755                    return false; // malformed LEB128
1756                }
1757            }
1758            pos = pos.saturating_add(size as usize);
1759        } else {
1760            // No size field — the rest of the buffer is this OBU's payload;
1761            // we can't skip past it to find subsequent OBUs.
1762            break;
1763        }
1764    }
1765    false
1766}
1767
1768struct SoftwareH264Encoder {
1769    encoder: OpenH264Encoder,
1770}
1771
1772impl SoftwareH264Encoder {
1773    fn new(quality: SurfaceQuality) -> Result<Self, String> {
1774        use openh264::encoder::{EncoderConfig, RateControlMode};
1775        let config = EncoderConfig::new()
1776            .set_bitrate_bps(quality.openh264_bitrate())
1777            .rate_control_mode(RateControlMode::Bitrate);
1778        let encoder =
1779            OpenH264Encoder::with_api_config(openh264::OpenH264API::from_source(), config)
1780                .map_err(|err| format!("failed to create OpenH264 encoder: {err:?}"))?;
1781        Ok(Self { encoder })
1782    }
1783
1784    fn request_keyframe(&mut self) {
1785        self.encoder.force_intra_frame();
1786    }
1787
1788    fn encode(&mut self, rgba: &[u8], width: u32, height: u32) -> Option<(Vec<u8>, bool)> {
1789        let yuv = rgba_to_yuv420(rgba, width as usize, height as usize);
1790        let yuv_buf = YUVBuffer::from_vec(yuv, width as usize, height as usize);
1791        self.encode_yuv(&yuv_buf, width, height)
1792    }
1793
1794    /// Encode from a pre-built YUV buffer (avoids redundant conversion).
1795    fn encode_yuv(
1796        &mut self,
1797        yuv_buf: &YUVBuffer,
1798        width: u32,
1799        height: u32,
1800    ) -> Option<(Vec<u8>, bool)> {
1801        let bitstream = match self.encoder.encode(yuv_buf) {
1802            Ok(bs) => bs,
1803            Err(e) => {
1804                eprintln!("[surface-encoder] openh264 encode failed {width}x{height}: {e:?}");
1805                return None;
1806            }
1807        };
1808        let nal_data = bitstream.to_vec();
1809        if nal_data.is_empty() {
1810            eprintln!("[surface-encoder] openh264 produced empty NAL {width}x{height}");
1811            return None;
1812        }
1813        let is_keyframe = h264_stream_contains_idr(&nal_data);
1814        Some((nal_data, is_keyframe))
1815    }
1816}
1817
1818// ---------------------------------------------------------------------------
1819// AV1 (rav1e)
1820// ---------------------------------------------------------------------------
1821
1822struct SoftwareAV1Encoder {
1823    ctx: rav1e::Context<u8>,
1824    width: usize,
1825    height: usize,
1826    force_keyframe: bool,
1827    chroma: ChromaSubsampling,
1828}
1829
1830impl SoftwareAV1Encoder {
1831    fn new(
1832        width: u32,
1833        height: u32,
1834        quality: SurfaceQuality,
1835        chroma: ChromaSubsampling,
1836    ) -> Result<Self, String> {
1837        use rav1e::prelude::*;
1838
1839        let chroma_sampling = if chroma.is_444() {
1840            ChromaSampling::Cs444
1841        } else {
1842            ChromaSampling::Cs420
1843        };
1844        let mut speed = SpeedSettings::from_preset(quality.av1_speed());
1845        speed.rdo_lookahead_frames = 1;
1846        let enc = EncoderConfig {
1847            width: width as usize,
1848            height: height as usize,
1849            chroma_sampling,
1850            chroma_sample_position: ChromaSamplePosition::Unknown,
1851            speed_settings: speed,
1852            low_latency: true,
1853            min_key_frame_interval: 0,
1854            max_key_frame_interval: 60,
1855            quantizer: quality.av1_quantizer(),
1856            min_quantizer: quality.av1_min_quantizer(),
1857            bitrate: 0,
1858            ..Default::default()
1859        };
1860        let cfg = Config::new().with_encoder_config(enc);
1861        let ctx = cfg
1862            .new_context()
1863            .map_err(|e| format!("rav1e context creation failed: {e}"))?;
1864        Ok(Self {
1865            ctx,
1866            width: width as usize,
1867            height: height as usize,
1868            force_keyframe: false,
1869            chroma,
1870        })
1871    }
1872
1873    fn request_keyframe(&mut self) {
1874        self.force_keyframe = true;
1875    }
1876
1877    fn encode(&mut self, rgba: &[u8]) -> Option<(Vec<u8>, bool)> {
1878        let yuv = if self.chroma.is_444() {
1879            rgba_to_yuv444(rgba, self.width, self.height)
1880        } else {
1881            rgba_to_yuv420(rgba, self.width, self.height)
1882        };
1883        self.encode_yuv_planes(&yuv)
1884    }
1885
1886    fn encode_nv12(
1887        &mut self,
1888        data: &[u8],
1889        y_stride: usize,
1890        uv_stride: usize,
1891        width: usize,
1892        height: usize,
1893    ) -> Option<(Vec<u8>, bool)> {
1894        let yuv = nv12_to_yuv420(data, y_stride, uv_stride, width, height);
1895        self.encode_yuv_planes(&yuv)
1896    }
1897
1898    /// Encode from pre-converted planar YUV data (Y + U + V contiguous).
1899    /// Layout depends on chroma: I420 (half-res UV) or I444 (full-res UV).
1900    fn encode_yuv_planes(&mut self, yuv: &[u8]) -> Option<(Vec<u8>, bool)> {
1901        let width = self.width;
1902        let height = self.height;
1903        let y_size = width * height;
1904        let (uv_w, uv_size) = if self.chroma.is_444() {
1905            (width, width * height)
1906        } else {
1907            let uv_w = width.div_ceil(2);
1908            let uv_h = height.div_ceil(2);
1909            (uv_w, uv_w * uv_h)
1910        };
1911
1912        let y_plane = &yuv[..y_size];
1913        let u_plane = &yuv[y_size..y_size + uv_size];
1914        let v_plane = &yuv[y_size + uv_size..];
1915
1916        let mut frame = self.ctx.new_frame();
1917        frame.planes[0].copy_from_raw_u8(y_plane, width, 1);
1918        frame.planes[1].copy_from_raw_u8(u_plane, uv_w, 1);
1919        frame.planes[2].copy_from_raw_u8(v_plane, uv_w, 1);
1920
1921        self.send_and_receive(frame)
1922    }
1923
1924    fn send_and_receive(&mut self, frame: rav1e::Frame<u8>) -> Option<(Vec<u8>, bool)> {
1925        use rav1e::prelude::*;
1926
1927        if self.force_keyframe {
1928            let params = FrameParameters {
1929                frame_type_override: FrameTypeOverride::Key,
1930                ..Default::default()
1931            };
1932            if self.ctx.send_frame((frame, params)).is_ok() {
1933                self.force_keyframe = false;
1934            }
1935        } else {
1936            let _ = self.ctx.send_frame(frame);
1937        }
1938
1939        match self.ctx.receive_packet() {
1940            Ok(packet) => {
1941                let is_key = packet.frame_type == rav1e::prelude::FrameType::KEY;
1942                Some((packet.data, is_key))
1943            }
1944            Err(rav1e::EncoderStatus::Encoded) | Err(rav1e::EncoderStatus::NeedMoreData) => None,
1945            Err(_) => None,
1946        }
1947    }
1948}
1949
1950#[cfg(test)]
1951mod tests {
1952    use super::*;
1953
1954    /// Build a minimal AV1 OBU with the given type, has_size=1.
1955    fn make_obu(obu_type: u8, payload: &[u8]) -> Vec<u8> {
1956        // header: forbidden=0, obu_type(4), extension=0, has_size=1, reserved=0
1957        let header = (obu_type & 0xF) << 3 | 0b10; // has_size=1
1958        let mut obu = vec![header];
1959        // LEB128-encode the payload length.
1960        let mut size = payload.len();
1961        loop {
1962            let mut byte = (size & 0x7F) as u8;
1963            size >>= 7;
1964            if size > 0 {
1965                byte |= 0x80;
1966            }
1967            obu.push(byte);
1968            if size == 0 {
1969                break;
1970            }
1971        }
1972        obu.extend_from_slice(payload);
1973        obu
1974    }
1975
1976    #[test]
1977    fn av1_keyframe_with_sequence_header_only() {
1978        // Sequence header OBU (type 1) as the only OBU — keyframe.
1979        let data = make_obu(1, &[0xAA; 10]);
1980        assert!(av1_stream_contains_keyframe(&data));
1981    }
1982
1983    #[test]
1984    fn av1_keyframe_with_temporal_delimiter_prefix() {
1985        // Temporal delimiter (type 2) + sequence header (type 1) — keyframe.
1986        // This is the typical NVENC output for a keyframe.
1987        let mut data = make_obu(2, &[]); // temporal delimiter, empty payload
1988        data.extend(make_obu(1, &[0xBB; 8])); // sequence header
1989        data.extend(make_obu(6, &[0xCC; 20])); // frame OBU
1990        assert!(av1_stream_contains_keyframe(&data));
1991    }
1992
1993    #[test]
1994    fn av1_non_keyframe_with_temporal_delimiter() {
1995        // Temporal delimiter (type 2) + frame (type 6) — not a keyframe.
1996        let mut data = make_obu(2, &[]);
1997        data.extend(make_obu(6, &[0xDD; 15]));
1998        assert!(!av1_stream_contains_keyframe(&data));
1999    }
2000
2001    #[test]
2002    fn av1_non_keyframe_frame_header_only() {
2003        // Frame header (type 3) — not a keyframe.
2004        let data = make_obu(3, &[0xEE; 5]);
2005        assert!(!av1_stream_contains_keyframe(&data));
2006    }
2007
2008    #[test]
2009    fn av1_empty_stream() {
2010        assert!(!av1_stream_contains_keyframe(&[]));
2011    }
2012
2013    #[test]
2014    fn av1_keyframe_large_leb128_size() {
2015        // Temporal delimiter with a larger payload needing multi-byte LEB128,
2016        // followed by a sequence header.
2017        let mut data = make_obu(2, &[0x00; 200]);
2018        data.extend(make_obu(1, &[0xFF; 4]));
2019        assert!(av1_stream_contains_keyframe(&data));
2020    }
2021}