wl-webrtc/src/bin/vaapi_import_bench/pipeline_cpu.rs

use std::path::Path;
use std::time::Instant;

use anyhow::{bail, Result};

use ffmpeg_next::ffi;

use wl_webrtc::avhw::{av_err_to_string, AvHwFrameCtx};
use wl_webrtc::cap_portal::{CapPortal, PwCtrlEvent};

use crate::pipeline_gpu::import_frame;
use crate::software::{
    create_software_encoder, create_sws_context, encode_yuv_frame, finish_encoder,
};
use crate::stats::FrameStats;

#[allow(clippy::too_many_arguments)]
pub(crate) fn run_cpu_pipeline(
    cap: &CapPortal,
    frames_ctx: &AvHwFrameCtx,
    output: &str,
    frames: u32,
    src_width: u32,
    src_height: u32,
    enc_width: u32,
    enc_height: u32,
) -> Result<FrameStats> {
    let mut encoder = create_software_encoder(Path::new(output), enc_width, enc_height)?;
    let sws_ctx = create_sws_context(
        src_width,
        src_height,
        ffi::AVPixelFormat::AV_PIX_FMT_BGRA,
        enc_width,
        enc_height,
    )?;

    println!(
        "       Encoder: {}, {}x{} YUV420P",
        encoder.codec_name, enc_width, enc_height
    );
    println!("       Output: {output}");
    println!("       CPU Pipeline: DMA-BUF 4K BGRA -> av_hwframe_map -> av_hwframe_transfer_data -> sws_scale -> YUV420P 2K -> encode\n");

    let mut stats = FrameStats {
        codec_name: encoder.codec_name.clone(),
        output_path: output.to_string(),
        ..FrameStats::default()
    };
    let total_start = Instant::now();
    let mut pts: i64 = 0;

    while stats.frames_encoded < frames {
        if let Ok(ctrl) = cap.event_receiver().try_recv() {
            match ctrl {
                PwCtrlEvent::StreamEnded => break,
                PwCtrlEvent::Error(e) => bail!(
                    "PipeWire error after {} CPU frames: {e}",
                    stats.frames_encoded
                ),
                PwCtrlEvent::FormatChanged { .. } => {}
            }
        }

        let frame = match cap
            .frame_receiver()
            .recv_timeout(std::time::Duration::from_secs(5))
        {
            Ok(f) => f,
            Err(_) => break,
        };

        let frame_start = Instant::now();
        let t_import = Instant::now();
        let vaapi_frame = match import_frame(frames_ctx, &frame) {
            Ok(f) => f,
            Err(e) => {
                stats.import_failures += 1;
                if stats.import_failures <= 3 {
                    eprintln!("CPU frame {}: import failed: {e}", stats.frames_encoded);
                }
                continue;
            }
        };
        let import_us = t_import.elapsed().as_micros() as u64;

        let t_transfer = Instant::now();
        // SAFETY: sw_frame is allocated by FFmpeg and freed on all paths below.
        let mut sw_frame = unsafe { ffi::av_frame_alloc() };
        if sw_frame.is_null() {
            bail!("CPU frame {}: av_frame_alloc failed", stats.frames_encoded);
        }
        // SAFETY: sw_frame is an allocated destination; vaapi_frame is a valid VAAPI source frame.
        let transfer_ret =
            unsafe { ffi::av_hwframe_transfer_data(sw_frame, vaapi_frame.as_ptr(), 0) };
        if transfer_ret < 0 {
            // SAFETY: sw_frame was allocated above and has not been freed yet.
            unsafe { ffi::av_frame_free(&mut sw_frame) };
            bail!(
                "CPU frame {}: av_hwframe_transfer_data failed: {} ({})",
                stats.frames_encoded,
                transfer_ret,
                av_err_to_string(transfer_ret)
            );
        }
        let transfer_us = t_transfer.elapsed().as_micros() as u64;

        let t_scale = Instant::now();
        // SAFETY: sw_frame contains transferred BGRA data; encoder.yuv_frame is writable YUV420P
        // at the configured output dimensions; sws_ctx converts and downscales between them.
        unsafe {
            ffi::av_frame_make_writable(encoder.yuv_frame);
            ffi::sws_scale(
                sws_ctx.0,
                (*sw_frame).data.as_ptr() as *const *const u8,
                (*sw_frame).linesize.as_ptr(),
                0,
                (*sw_frame).height,
                (*encoder.yuv_frame).data.as_ptr() as *mut *mut u8,
                (*encoder.yuv_frame).linesize.as_ptr(),
            );
        }
        let scale_us = t_scale.elapsed().as_micros() as u64;
        // SAFETY: sw_frame was allocated above and is no longer needed after scaling.
        unsafe { ffi::av_frame_free(&mut sw_frame) };

        let encode_us = encode_yuv_frame(&mut encoder, &mut pts)?;
        let total_us = frame_start.elapsed().as_micros() as u64;

        stats.import_us.push(import_us);
        stats.transfer_us.push(transfer_us);
        stats.scale_us.push(scale_us);
        stats.encode_us.push(encode_us);
        stats.total_us.push(total_us);
        stats.frames_encoded += 1;

        if stats.frames_encoded <= 3 || stats.frames_encoded.is_multiple_of(30) {
            println!(
                "       CPU frame {:>4}/{frames}: import={:.2}ms transfer={:.2}ms scale={:.2}ms encode={:.2}ms total={:.2}ms",
                stats.frames_encoded,
                import_us as f64 / 1000.0,
                transfer_us as f64 / 1000.0,
                scale_us as f64 / 1000.0,
                encode_us as f64 / 1000.0,
                total_us as f64 / 1000.0,
            );
        }
    }

    finish_encoder(encoder)?;
    stats.elapsed_secs = total_start.elapsed().as_secs_f64();
    Ok(stats)
}