webtransport-video/backend/src/engine.rs

use std::num::NonZero;
use std::thread;
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use tokio::sync::broadcast;
use tracing::{error, info, warn};
use vpx_rs::enc::{CodecId, EncoderFrameFlags, EncodingDeadline};
use vpx_rs::{Encoder, EncoderConfig, RateControl, Timebase};
use scap::capturer::{Area, Capturer, Options, Point, Resolution, Size};
use scap::frame::FrameType;

use crate::config::Args;
use crate::pixelutil;

/// Represents a video packet to be sent to clients.
#[derive(Clone, Debug)]
pub struct VideoPacket {
    pub data: Vec<u8>,
    pub is_key: bool,
}

/// Manages video capture and encoding.
pub struct VideoEngine {
    args: Args,
    packet_tx: broadcast::Sender<VideoPacket>,
    shutdown_tx: broadcast::Sender<()>,
}

impl VideoEngine {
    pub fn new(args: Args) -> Self {
        let (packet_tx, _) = broadcast::channel(100);
        let (shutdown_tx, _) = broadcast::channel(1);
        Self {
            args,
            packet_tx,
            shutdown_tx,
        }
    }

    pub fn subscribe(&self) -> broadcast::Receiver<VideoPacket> {
        self.packet_tx.subscribe()
    }

    pub fn start(&self) {
        let args = self.args.clone();
        let packet_tx = self.packet_tx.clone();
        let mut shutdown_rx = self.shutdown_tx.subscribe();

        // Spawn the capture/encode loop in a separate dedicated thread
        // because it involves blocking heavy computation.
        thread::spawn(move || {
            if let Err(e) = run_capture_loop(args, packet_tx, &mut shutdown_rx) {
                error!("Capture loop failed: {:?}", e);
            }
        });
    }

    #[allow(dead_code)]
    pub fn stop(&self) {
        let _ = self.shutdown_tx.send(());
    }
}

fn run_capture_loop(
    args: Args,
    packet_tx: broadcast::Sender<VideoPacket>,
    shutdown_rx: &mut broadcast::Receiver<()>,
) -> Result<()> {
    // 1. Setup Capturer
    let options = Options {
        fps: args.fps,
        show_cursor: true,
        show_highlight: true,
        output_type: FrameType::YUVFrame,
        output_resolution: Resolution::_1080p,
        crop_area: Some(Area {
            origin: Point { x: 0.0, y: 0.0 },
            size: Size {
                width: args.width as f64,
                height: args.height as f64,
            },
        }),
        ..Default::default()
    };

    let mut capturer = Capturer::build(options).context("Failed to build capturer")?;
    capturer.start_capture();

    // 2. Setup Encoder
    let mut config = EncoderConfig::<u8>::new(
        CodecId::VP9,
        args.width,
        args.height,
        Timebase {
            num: NonZero::new(1).unwrap(),
            den: NonZero::new(args.fps).unwrap(),
        },
        RateControl::ConstantBitRate(args.bitrate),
    )
    .context("Failed to create encoder config")?;
    
    config.threads = 4; // Use reasonable number of threads
    config.lag_in_frames = 0; // Low latency

    let mut encoder = Encoder::new(config).context("Failed to create encoder")?;

    let frame_duration = Duration::from_secs_f64(1.0 / args.fps as f64);
    let mut next_frame_time = Instant::now();
    let mut pts: i64 = 0;

    info!("Starting capture loop: {}x{} @ {}fps", args.width, args.height, args.fps);

    loop {
        // Check for shutdown signal
        if shutdown_rx.try_recv().is_ok() {
            info!("Shutting down capture loop");
            break;
        }

        let now = Instant::now();
        if now < next_frame_time {
            thread::sleep(next_frame_time - now);
        }
        next_frame_time += frame_duration;

        match capturer.get_next_frame() {
            Ok(captured_frame) => {
                // Convert frame
                let yuv_buffer = pixelutil::frame_to_yuv(captured_frame);
                let yuv_image = pixelutil::apply_frame(&yuv_buffer, args.width as usize, args.height as usize);

                // Encode
                // For real-time streaming, we force keyframes periodically could be an improvement,
                // but for now relying on default behavior or client requests (not impl here).
                let flags = EncoderFrameFlags::empty(); 
                
                match encoder.encode(
                    pts,
                    1, // Duration (in timebase units)
                    yuv_image,
                    EncodingDeadline::Realtime,
                    flags,
                ) {
                    Ok(packets) => {
                         for packet in packets {
                            if let vpx_rs::Packet::CompressedFrame(frame) = packet {
                                let video_packet = VideoPacket {
                                    data: frame.data.to_vec(),
                                    is_key: frame.flags.is_key,
                                };
                                
                                // Broadcast to all connected clients
                                // It's okay if no one is listening
                                let _ = packet_tx.send(video_packet);
                            }
                        }
                        pts += 1;
                    },
                    Err(e) => {
                        error!("Encoding error: {:?}", e);
                        // Continue loop, maybe transient
                    }
                }
            }
            Err(e) => {
                // Scap might return error if no frame is available yet or other issues
                // We just log debug/warn and continue
                 warn!("Capture error (might be transient): {:?}", e);
                 thread::sleep(Duration::from_millis(10));
            }
        }
    }

    Ok(())
}