guru-connect/agent/src/encoder/raw.rs

//! Raw frame encoder with Zstd compression
//!
//! Best for LAN connections where bandwidth is plentiful and latency is critical.
//! Compresses BGRA pixel data using Zstd for fast compression/decompression.

use super::{EncodedFrame, Encoder};
use crate::capture::{CapturedFrame, DirtyRect};
use crate::proto::{video_frame, DirtyRect as ProtoDirtyRect, RawFrame, VideoFrame};
use anyhow::Result;

/// Raw frame encoder with Zstd compression
pub struct RawEncoder {
    /// Compression level (1-22, default 3 for speed)
    compression_level: i32,

    /// Previous frame for delta detection
    previous_frame: Option<Vec<u8>>,

    /// Force keyframe on next encode
    force_keyframe: bool,

    /// Frame counter
    sequence: u32,
}

impl RawEncoder {
    /// Create a new raw encoder
    ///
    /// Quality 1-100 maps to Zstd compression level:
    /// - Low quality (1-33): Level 1-3 (fastest)
    /// - Medium quality (34-66): Level 4-9
    /// - High quality (67-100): Level 10-15 (best compression)
    pub fn new(quality: u32) -> Result<Self> {
        let compression_level = Self::quality_to_level(quality);

        Ok(Self {
            compression_level,
            previous_frame: None,
            force_keyframe: true, // Start with keyframe
            sequence: 0,
        })
    }

    /// Convert quality (1-100) to Zstd compression level
    fn quality_to_level(quality: u32) -> i32 {
        // Lower quality = faster compression (level 1-3)
        // Higher quality = better compression (level 10-15)
        // We optimize for speed, so cap at 6
        match quality {
            0..=33 => 1,
            34..=50 => 2,
            51..=66 => 3,
            67..=80 => 4,
            81..=90 => 5,
            _ => 6,
        }
    }

    /// Compress data using Zstd
    fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
        let compressed = zstd::encode_all(data, self.compression_level)?;
        Ok(compressed)
    }

    /// Detect dirty rectangles by comparing with previous frame
    fn detect_dirty_rects(
        &self,
        current: &[u8],
        previous: &[u8],
        width: u32,
        height: u32,
    ) -> Vec<DirtyRect> {
        // Simple block-based dirty detection
        // Divide screen into 64x64 blocks and check which changed
        const BLOCK_SIZE: u32 = 64;

        let mut dirty_rects = Vec::new();
        let stride = (width * 4) as usize;

        let blocks_x = (width + BLOCK_SIZE - 1) / BLOCK_SIZE;
        let blocks_y = (height + BLOCK_SIZE - 1) / BLOCK_SIZE;

        for by in 0..blocks_y {
            for bx in 0..blocks_x {
                let x = bx * BLOCK_SIZE;
                let y = by * BLOCK_SIZE;
                let block_w = (BLOCK_SIZE).min(width - x);
                let block_h = (BLOCK_SIZE).min(height - y);

                // Check if this block changed
                let mut changed = false;
                'block_check: for row in 0..block_h {
                    let row_start = ((y + row) as usize * stride) + (x as usize * 4);
                    let row_end = row_start + (block_w as usize * 4);

                    if row_end <= current.len() && row_end <= previous.len() {
                        if current[row_start..row_end] != previous[row_start..row_end] {
                            changed = true;
                            break 'block_check;
                        }
                    } else {
                        changed = true;
                        break 'block_check;
                    }
                }

                if changed {
                    dirty_rects.push(DirtyRect {
                        x,
                        y,
                        width: block_w,
                        height: block_h,
                    });
                }
            }
        }

        // Merge adjacent dirty rects (simple optimization)
        // TODO: Implement proper rectangle merging

        dirty_rects
    }

    /// Extract pixels for dirty rectangles only
    fn extract_dirty_pixels(
        &self,
        data: &[u8],
        width: u32,
        dirty_rects: &[DirtyRect],
    ) -> Vec<u8> {
        let stride = (width * 4) as usize;
        let mut pixels = Vec::new();

        for rect in dirty_rects {
            for row in 0..rect.height {
                let row_start = ((rect.y + row) as usize * stride) + (rect.x as usize * 4);
                let row_end = row_start + (rect.width as usize * 4);

                if row_end <= data.len() {
                    pixels.extend_from_slice(&data[row_start..row_end]);
                }
            }
        }

        pixels
    }
}

impl Encoder for RawEncoder {
    fn encode(&mut self, frame: &CapturedFrame) -> Result<EncodedFrame> {
        self.sequence = self.sequence.wrapping_add(1);

        let is_keyframe = self.force_keyframe || self.previous_frame.is_none();
        self.force_keyframe = false;

        let (data_to_compress, dirty_rects, full_frame) = if is_keyframe {
            // Keyframe: send full frame
            (frame.data.clone(), Vec::new(), true)
        } else if let Some(ref previous) = self.previous_frame {
            // Delta frame: detect and send only changed regions
            let dirty_rects =
                self.detect_dirty_rects(&frame.data, previous, frame.width, frame.height);

            if dirty_rects.is_empty() {
                // No changes, skip frame
                return Ok(EncodedFrame {
                    frame: VideoFrame::default(),
                    size: 0,
                    is_keyframe: false,
                });
            }

            // If too many dirty rects, just send full frame
            if dirty_rects.len() > 50 {
                (frame.data.clone(), Vec::new(), true)
            } else {
                let dirty_pixels = self.extract_dirty_pixels(&frame.data, frame.width, &dirty_rects);
                (dirty_pixels, dirty_rects, false)
            }
        } else {
            (frame.data.clone(), Vec::new(), true)
        };

        // Compress the data
        let compressed = self.compress(&data_to_compress)?;
        let size = compressed.len();

        // Build protobuf message
        let proto_dirty_rects: Vec<ProtoDirtyRect> = dirty_rects
            .iter()
            .map(|r| ProtoDirtyRect {
                x: r.x as i32,
                y: r.y as i32,
                width: r.width as i32,
                height: r.height as i32,
            })
            .collect();

        let raw_frame = RawFrame {
            width: frame.width as i32,
            height: frame.height as i32,
            data: compressed,
            compressed: true,
            dirty_rects: proto_dirty_rects,
            is_keyframe: full_frame,
        };

        let video_frame = VideoFrame {
            timestamp: frame.timestamp.elapsed().as_millis() as i64,
            display_id: frame.display_id as i32,
            sequence: self.sequence as i32,
            encoding: Some(video_frame::Encoding::Raw(raw_frame)),
        };

        // Save current frame for next comparison
        self.previous_frame = Some(frame.data.clone());

        Ok(EncodedFrame {
            frame: video_frame,
            size,
            is_keyframe: full_frame,
        })
    }

    fn request_keyframe(&mut self) {
        self.force_keyframe = true;
    }

    fn name(&self) -> &str {
        "raw+zstd"
    }
}