lesavka/server/src/video_sinks/hevc_mjpeg_guard.rs

use crate::video_support::env_u32;

const DEFAULT_HEVC_JPEG_QUALITY: u32 = 72;
const DEFAULT_HEVC_SIZE_DROP_PCT: u32 = 45;
const DEFAULT_HEVC_MIN_REFERENCE_BYTES: u32 = 64 * 1024;
const DEFAULT_HEVC_MIN_PAYLOAD_DISTINCT_BYTES: u32 = 12;
const DEFAULT_HEVC_DOMINANT_BYTE_PCT: u32 = 92;

/// Resolve the JPEG quality used after HEVC decode.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_JPEG_QUALITY`, clamped to 1..=100.
/// Output: the `jpegenc` quality value. Why: HEVC ingress must become MJPEG
/// for the existing UVC gadget path, and smaller JPEGs avoid UVC/browser
/// partial-frame smears without changing the calibrated A/V timing model.
pub(super) fn hevc_jpeg_quality() -> u32 {
    env_u32("LESAVKA_UVC_HEVC_JPEG_QUALITY", DEFAULT_HEVC_JPEG_QUALITY).clamp(1, 100)
}

/// Decide whether suspicious decoded-frame freezing is enabled.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_FREEZE_ON_SIZE_DROP`. Output: true unless
/// explicitly disabled. Why: when one damaged decoded MJPEG frame appears, a
/// short freeze is less disruptive than showing a grey slab or torn half-frame.
pub(super) fn freeze_on_size_drop_enabled() -> bool {
    std::env::var("LESAVKA_UVC_HEVC_FREEZE_ON_SIZE_DROP")
        .ok()
        .map(|value| {
            let trimmed = value.trim();
            !(trimmed.eq_ignore_ascii_case("0")
                || trimmed.eq_ignore_ascii_case("false")
                || trimmed.eq_ignore_ascii_case("no")
                || trimmed.eq_ignore_ascii_case("off"))
        })
        .unwrap_or(true)
}

/// Resolve the frame-size drop percentage that triggers a freeze.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_SIZE_DROP_PCT`, clamped to 1..=95.
/// Output: next-frame size percentage of the last good frame. Why: the guard
/// should catch sudden damaged-frame collapses while still allowing normal
/// bitrate variation from scene motion and encoder decisions.
pub(super) fn size_drop_pct() -> u32 {
    env_u32("LESAVKA_UVC_HEVC_SIZE_DROP_PCT", DEFAULT_HEVC_SIZE_DROP_PCT).clamp(1, 95)
}

/// Resolve the minimum reference frame size for collapse detection.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_MIN_REFERENCE_BYTES`. Output: byte count.
/// Why: tiny synthetic or blank frames should not establish a baseline that
/// causes later healthy low-detail frames to be frozen.
pub(super) fn min_reference_bytes() -> u32 {
    env_u32(
        "LESAVKA_UVC_HEVC_MIN_REFERENCE_BYTES",
        DEFAULT_HEVC_MIN_REFERENCE_BYTES,
    )
    .max(1)
}

/// Resolve the minimum compressed-payload byte variety for decoded MJPEG.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_MIN_PAYLOAD_DISTINCT_BYTES`. Output:
/// distinct byte count threshold. Why: grey/black smear failures can arrive as
/// complete JPEG buffers, so the guard needs a conservative flat-payload check
/// in addition to simple size-collapse detection.
pub(super) fn min_payload_distinct_bytes() -> u32 {
    env_u32(
        "LESAVKA_UVC_HEVC_MIN_PAYLOAD_DISTINCT_BYTES",
        DEFAULT_HEVC_MIN_PAYLOAD_DISTINCT_BYTES,
    )
    .clamp(1, 64)
}

/// Resolve the dominant-byte percentage that marks a payload as too flat.
///
/// Inputs: optional `LESAVKA_UVC_HEVC_DOMINANT_BYTE_PCT`, clamped to 50..=99.
/// Output: percentage threshold. Why: a single repeated byte occupying almost
/// all entropy-coded JPEG payload is more likely a damaged decode/transfer
/// artifact than useful webcam video, and freezing is the safer conference UX.
pub(super) fn dominant_byte_pct() -> u32 {
    env_u32(
        "LESAVKA_UVC_HEVC_DOMINANT_BYTE_PCT",
        DEFAULT_HEVC_DOMINANT_BYTE_PCT,
    )
    .clamp(50, 99)
}

/// Return whether a decoded buffer looks like one complete JPEG image.
///
/// Inputs: decoded MJPEG bytes. Output: true when SOI, SOS, and EOI markers
/// are present. Why: incomplete JPEGs can still be non-empty buffers; freezing
/// them is safer than letting the UVC helper expose a partially decoded frame.
fn looks_like_complete_jpeg(bytes: &[u8]) -> bool {
    bytes.len() >= 4
        && bytes.starts_with(&[0xff, 0xd8])
        && bytes.ends_with(&[0xff, 0xd9])
        && bytes.windows(2).any(|pair| pair == [0xff, 0xda])
}

/// Return whether one complete JPEG has an implausibly flat payload.
///
/// Inputs: decoded MJPEG bytes. Output: true for dominant-byte or very low
/// variety payloads. Why: the visible failure mode is often a grey/black slab
/// that is syntactically present but not meaningful webcam video.
fn suspiciously_flat_payload(bytes: &[u8]) -> bool {
    if bytes.len() < min_reference_bytes() as usize / 4 {
        return false;
    }

    let start = bytes
        .windows(2)
        .position(|pair| pair == [0xff, 0xda])
        .map(|idx| (idx + 2).min(bytes.len()))
        .unwrap_or_else(|| bytes.len().min(256));
    let end = bytes.len().saturating_sub(2);
    if end <= start || end - start < 512 {
        return false;
    }

    let payload = &bytes[start..end];
    let mut counts = [0u32; 256];
    for byte in payload {
        counts[*byte as usize] += 1;
    }

    let distinct = counts.iter().filter(|count| **count > 0).count() as u32;
    let dominant = counts.iter().copied().max().unwrap_or(0) as u64;
    let total = payload.len() as u64;

    distinct < min_payload_distinct_bytes()
        || dominant.saturating_mul(100) >= total.saturating_mul(u64::from(dominant_byte_pct()))
}

/// Decide whether a decoded HEVC-to-MJPEG frame should be frozen out.
///
/// Inputs: byte length of the last successfully spooled decoded MJPEG and the
/// next decoded MJPEG. Output: true when the next frame looks like a damaged
/// collapse. Why: keeping the last good frame preserves freshness and sync
/// better than forwarding a syntactically valid but visually corrupted JPEG.
pub(super) fn should_freeze_decoded_mjpeg(previous_bytes: u64, next_bytes: usize) -> bool {
    if !freeze_on_size_drop_enabled() || previous_bytes < u64::from(min_reference_bytes()) {
        return false;
    }

    let next_bytes = next_bytes as u64;
    let threshold_bytes = previous_bytes.saturating_mul(u64::from(size_drop_pct())) / 100;
    next_bytes < threshold_bytes
}

/// Decide whether a decoded MJPEG payload should be frozen before UVC spool.
///
/// Inputs: byte length of the last successfully spooled decoded MJPEG and the
/// decoded MJPEG bytes. Output: true when the payload is incomplete, collapsed,
/// or suspiciously flat. Why: users prefer a short freeze over grey slabs,
/// mostly black frames, or torn images that conferencing apps may otherwise
/// display as if they were valid webcam frames.
pub(super) fn should_freeze_decoded_mjpeg_frame(previous_bytes: u64, decoded_mjpeg: &[u8]) -> bool {
    if !freeze_on_size_drop_enabled() || previous_bytes < u64::from(min_reference_bytes()) {
        return false;
    }
    !looks_like_complete_jpeg(decoded_mjpeg)
        || should_freeze_decoded_mjpeg(previous_bytes, decoded_mjpeg.len())
        || suspiciously_flat_payload(decoded_mjpeg)
}

#[cfg(test)]
mod tests {
    #[test]
    fn hevc_jpeg_quality_defaults_to_moderate_transport_pressure() {
        temp_env::with_var_unset("LESAVKA_UVC_HEVC_JPEG_QUALITY", || {
            assert_eq!(super::hevc_jpeg_quality(), 72);
        });

        temp_env::with_var("LESAVKA_UVC_HEVC_JPEG_QUALITY", Some("101"), || {
            assert_eq!(super::hevc_jpeg_quality(), 100);
        });

        temp_env::with_var("LESAVKA_UVC_HEVC_JPEG_QUALITY", Some("0"), || {
            assert_eq!(super::hevc_jpeg_quality(), 1);
        });
    }

    #[test]
    fn freeze_guard_catches_large_decoded_frame_collapses() {
        temp_env::with_vars(
            [
                ("LESAVKA_UVC_HEVC_FREEZE_ON_SIZE_DROP", Some("1")),
                ("LESAVKA_UVC_HEVC_SIZE_DROP_PCT", Some("45")),
                ("LESAVKA_UVC_HEVC_MIN_REFERENCE_BYTES", Some("65536")),
            ],
            || {
                assert!(super::should_freeze_decoded_mjpeg(200_000, 80_000));
                assert!(!super::should_freeze_decoded_mjpeg(200_000, 110_000));
                assert!(!super::should_freeze_decoded_mjpeg(20_000, 1_000));
            },
        );
    }

    #[test]
    fn freeze_guard_can_be_disabled_for_diagnostics() {
        temp_env::with_vars(
            [
                ("LESAVKA_UVC_HEVC_FREEZE_ON_SIZE_DROP", Some("0")),
                ("LESAVKA_UVC_HEVC_SIZE_DROP_PCT", Some("95")),
                ("LESAVKA_UVC_HEVC_MIN_REFERENCE_BYTES", Some("1")),
            ],
            || {
                assert!(!super::should_freeze_decoded_mjpeg(200_000, 1_000));
            },
        );
    }

    #[test]
    /// Verifies flat or incomplete HEVC-decoded MJPEG frames freeze out.
    ///
    /// Inputs: synthetic JPEG-like payloads. Output: assertions only. Why:
    /// the source-level guard should keep the same behavior as the root
    /// taxonomy tests even when this helper is compiled inside the server crate.
    fn freeze_guard_rejects_incomplete_or_flat_decoded_payloads() {
        /// Build a minimal JPEG-like buffer for guard heuristic tests.
        ///
        /// Inputs: entropy-coded payload bytes. Output: synthetic SOI/SOS/EOI
        /// wrapper. Why: the guard tests need deterministic payload shape
        /// without depending on slow image encoding.
        fn jpeg_with_payload(payload: &[u8]) -> Vec<u8> {
            let mut bytes = vec![0xff, 0xd8, 0xff, 0xda];
            bytes.extend_from_slice(payload);
            bytes.extend_from_slice(&[0xff, 0xd9]);
            bytes
        }

        temp_env::with_vars(
            [
                ("LESAVKA_UVC_HEVC_FREEZE_ON_SIZE_DROP", Some("1")),
                ("LESAVKA_UVC_HEVC_SIZE_DROP_PCT", Some("45")),
                ("LESAVKA_UVC_HEVC_MIN_REFERENCE_BYTES", Some("65536")),
                (
                    "LESAVKA_UVC_HEVC_MIN_PAYLOAD_DISTINCT_BYTES",
                    Some("12"),
                ),
                ("LESAVKA_UVC_HEVC_DOMINANT_BYTE_PCT", Some("92")),
            ],
            || {
                let healthy_payload: Vec<u8> =
                    (0..120_000).map(|idx| (idx % 251) as u8).collect();
                let flat_payload = vec![0x80; 120_000];
                let mut truncated = jpeg_with_payload(&healthy_payload);
                truncated.pop();

                assert!(!super::should_freeze_decoded_mjpeg_frame(
                    200_000,
                    &jpeg_with_payload(&healthy_payload),
                ));
                assert!(super::should_freeze_decoded_mjpeg_frame(
                    200_000,
                    &jpeg_with_payload(&flat_payload),
                ));
                assert!(super::should_freeze_decoded_mjpeg_frame(
                    200_000, &truncated,
                ));
            },
        );
    }
}