wl-webrtc/src/avhw/import.rs

use std::path::Path;
use std::slice;

use anyhow::{bail, Result};
use ffmpeg_next as ff;
use ffmpeg_next::ffi;

use super::filter::build_swenc_filter_graph;
use super::{ff_err, AvHwDevCtx, AvHwFrameCtx};
use super::{BitrateCommand, CpuNv12Frame, ResolutionChange};

pub struct SwEncImport {
    hw_dev: AvHwDevCtx,
    frames_rgb: AvHwFrameCtx,
    filter_graph: ff::filter::Graph,
    width: u32,
    height: u32,
    enc_width: u32,
    enc_height: u32,
    fps: u32,
    resolution_rx: Option<crossbeam_channel::Receiver<BitrateCommand>>,
    encoder_resolution_tx: Option<crossbeam_channel::Sender<ResolutionChange>>,
}

impl SwEncImport {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        drm_device: &Path,
        width: u32,
        height: u32,
        enc_width: u32,
        enc_height: u32,
        fps: u32,
    ) -> Result<Self> {
        let hw_dev = AvHwDevCtx::new_vaapi(drm_device)?;
        let frames_rgb =
            AvHwFrameCtx::for_capture(&hw_dev, width, height, ff::format::Pixel::BGRA)?;
        let filter_graph = build_swenc_filter_graph(
            &hw_dev,
            &frames_rgb,
            width,
            height,
            enc_width,
            enc_height,
            fps,
        )?;

        Ok(Self {
            hw_dev,
            frames_rgb,
            filter_graph,
            width,
            height,
            enc_width,
            enc_height,
            fps,
            resolution_rx: None,
            encoder_resolution_tx: None,
        })
    }

    #[allow(clippy::too_many_arguments)]
    pub fn new_with_resolution_control(
        drm_device: &Path,
        width: u32,
        height: u32,
        enc_width: u32,
        enc_height: u32,
        fps: u32,
        resolution_rx: crossbeam_channel::Receiver<BitrateCommand>,
        encoder_resolution_tx: crossbeam_channel::Sender<ResolutionChange>,
    ) -> Result<Self> {
        let mut this = Self::new(drm_device, width, height, enc_width, enc_height, fps)?;
        this.resolution_rx = Some(resolution_rx);
        this.encoder_resolution_tx = Some(encoder_resolution_tx);
        Ok(this)
    }

    pub fn frames_rgb(&self) -> &AvHwFrameCtx {
        let _ = self.hw_dev.as_ptr();
        &self.frames_rgb
    }

    pub fn import_and_scale(&mut self, hw_frame: &ff::frame::Video) -> Result<CpuNv12Frame> {
        self.poll_resolution_commands()?;

        let mut filter_src_ctx = self
            .filter_graph
            .get("in")
            .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?;
        let mut filter_src = filter_src_ctx.source();
        let mut filter_sink_ctx = self
            .filter_graph
            .get("out")
            .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?;
        let mut filter_sink = filter_sink_ctx.sink();

        filter_src
            .add(hw_frame)
            .map_err(|e| anyhow::anyhow!("software pipeline filter source add failed: {e}"))?;

        let mut first = None;
        let mut extra_count = 0usize;
        loop {
            let mut filtered = ff::frame::Video::empty();
            match filter_sink.frame(&mut filtered) {
                Ok(()) => {
                    if filtered.pts().is_none() {
                        filtered.set_pts(hw_frame.pts());
                    }
                    let cpu_frame = self.transfer_filtered_to_cpu(&filtered)?;
                    if first.is_none() {
                        first = Some(cpu_frame);
                    } else {
                        extra_count += 1;
                    }
                }
                Err(ff::Error::Other { errno }) if errno == ffi::EAGAIN => break,
                Err(e) => bail!("software pipeline filter sink get frame failed: {e}"),
            }
        }

        if extra_count > 0 {
            tracing::warn!(
                "software import filter produced {extra_count} extra frame(s); dropping extras"
            );
        }

        first.ok_or_else(|| anyhow::anyhow!("software pipeline produced no scaled frame"))
    }

    pub fn flush_import(&mut self) -> Result<Vec<CpuNv12Frame>> {
        let mut filter_src_ctx = self
            .filter_graph
            .get("in")
            .ok_or_else(|| anyhow::anyhow!("filter 'in' not found"))?;
        let mut filter_src = filter_src_ctx.source();
        if let Err(e) = filter_src.flush() {
            tracing::debug!("filter source flush error: {e}");
        }

        let mut filter_sink_ctx = self
            .filter_graph
            .get("out")
            .ok_or_else(|| anyhow::anyhow!("filter 'out' not found"))?;
        let mut filter_sink = filter_sink_ctx.sink();
        let mut frames = Vec::new();
        loop {
            let mut filtered = ff::frame::Video::empty();
            match filter_sink.frame(&mut filtered) {
                Ok(()) => frames.push(self.transfer_filtered_to_cpu(&filtered)?),
                Err(_) => break,
            }
        }

        Ok(frames)
    }

    fn poll_resolution_commands(&mut self) -> Result<()> {
        let Some(rx) = self.resolution_rx.as_ref().cloned() else {
            return Ok(());
        };

        let mut requested = None;
        while let Ok(cmd) = rx.try_recv() {
            match cmd {
                BitrateCommand::UpdateResolution { width, height } => {
                    requested = Some((width & !1, height & !1));
                }
                BitrateCommand::UpdateBitrate { .. } => {}
                BitrateCommand::ForceKeyframe => {}
            }
        }

        let Some((width, height)) = requested else {
            return Ok(());
        };
        if width == self.enc_width && height == self.enc_height {
            return Ok(());
        }

        tracing::info!(
            from = format_args!("{}x{}", self.enc_width, self.enc_height),
            to = format_args!("{}x{}", width, height),
            "rebuilding software import filter graph for resolution change"
        );
        let _ = self.flush_import();
        self.filter_graph = build_swenc_filter_graph(
            &self.hw_dev,
            &self.frames_rgb,
            self.width,
            self.height,
            width,
            height,
            self.fps,
        )?;
        self.enc_width = width;
        self.enc_height = height;

        if let Some(tx) = &self.encoder_resolution_tx {
            tx.send(ResolutionChange { width, height })
                .map_err(|_| anyhow::anyhow!("encoder resolution channel disconnected"))?;
        }

        Ok(())
    }

    fn transfer_filtered_to_cpu(&self, filtered: &ff::frame::Video) -> Result<CpuNv12Frame> {
        // SAFETY: av_frame_alloc returns a newly allocated AVFrame or null,
        // which is checked below.
        let mut sw_nv12 = unsafe { ffi::av_frame_alloc() };
        if sw_nv12.is_null() {
            bail!("av_frame_alloc failed for NV12 transfer frame");
        }

        // SAFETY: sw_nv12 is an allocated destination frame; filtered is a valid VAAPI NV12
        // surface produced by scale_vaapi at encoder dimensions.
        let transfer_ret = unsafe { ffi::av_hwframe_transfer_data(sw_nv12, filtered.as_ptr(), 0) };
        if transfer_ret < 0 {
            // SAFETY: sw_nv12 was allocated above and has not been freed yet.
            unsafe { ffi::av_frame_free(&mut sw_nv12) };
            bail!(
                "av_hwframe_transfer_data failed for GPU-downscaled frame: {}",
                ff_err(transfer_ret)
            );
        }

        // SAFETY: sw_nv12 was filled by av_hwframe_transfer_data. NV12 planes 0 and 1 are
        // initialized for enc_width x enc_height; linesize values define each row's byte span.
        let frame = unsafe {
            let y_ptr = (*sw_nv12).data[0];
            let uv_ptr = (*sw_nv12).data[1];
            if y_ptr.is_null() || uv_ptr.is_null() {
                ffi::av_frame_free(&mut sw_nv12);
                bail!("NV12 transfer frame missing Y/UV plane data");
            }
            let y_stride = (*sw_nv12).linesize[0] as usize;
            let uv_stride = (*sw_nv12).linesize[1] as usize;
            if (*sw_nv12).width != self.enc_width as i32
                || (*sw_nv12).height != self.enc_height as i32
            {
                ffi::av_frame_free(&mut sw_nv12);
                bail!("NV12 transfer frame has unexpected dimensions");
            }
            let y_len = y_stride * self.enc_height as usize;
            let uv_len = uv_stride * (self.enc_height as usize / 2);
            let y_data = slice::from_raw_parts(y_ptr, y_len).to_vec();
            let uv_data = slice::from_raw_parts(uv_ptr, uv_len).to_vec();
            let pts = filtered.pts().unwrap_or(0);
            ffi::av_frame_free(&mut sw_nv12);
            CpuNv12Frame {
                y_data,
                uv_data,
                y_stride,
                uv_stride,
                pts,
                capture_time: std::time::Instant::now(),
            }
        };

        Ok(frame)
    }
}