wl-webrtc/src/state_portal.rs

// 采集门户状态模块 —— 通过 PipeWire/DMA-BUF 进行屏幕采集并编码
use std::mem;
use std::os::fd::AsRawFd;
use std::path::PathBuf;

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

use crate::args::Args;
use crate::avhw::{self, EncState};
use crate::cap_portal::{CapPortal, PwDmaBufFrame, PwEvent};
use crate::fps_limit::FpsLimit;
use crate::transform::Transform;

/// 门户采集的阶段状态
/// - WaitingForFormat: 等待接收到第一帧 DMA-BUF 以确定视频格式参数
/// - Streaming: 已完成初始化，正在持续编码流
enum PortalStage {
    WaitingForFormat,
    Streaming,
}

/// 门户模式的主状态机
///
/// 负责管理从 PipeWire 采集屏幕帧、通过 VAAPI 硬件编码的完整生命周期。
/// 工作流程：等待第一帧 → 创建编码器 → 持续编码帧数据。
pub struct StatePortal {
    /// 当前采集阶段
    stage: PortalStage,
    /// 硬件编码器状态（第一帧到达后才初始化）
    enc: Option<EncState>,
    /// 帧率限制器
    fps_limit: FpsLimit<()>,
    /// PipeWire 屏幕采集端点
    cap: CapPortal,
    /// 命令行参数
    args: Args,
    /// 是否遇到错误
    errored: bool,
    /// 是否为第一帧（首帧跳过帧率限制）
    first_frame: bool,
    /// DRM 渲染设备路径（如 /dev/dri/renderD128）；None 表示首帧自动检测
    drm_device: Option<PathBuf>,
    /// 第一帧的时间戳（纳秒），用于计算相对 PTS
    first_pts_ns: Option<i64>,
}

impl StatePortal {
    /// 创建门户状态实例
    ///
    /// 初始化 DRM 设备路径和 PipeWire 采集端点，编码器延迟到第一帧到达时创建。
    pub fn new(args: Args) -> Result<Self> {
        let drm_device = resolve_drm_device(&args)?;
        if let Some(ref drm_device) = drm_device {
            tracing::info!("Using DRM device: {}", drm_device.display());
        } else {
            tracing::info!("DRM device auto-detection enabled");
        }

        let cap = CapPortal::new(&args)?;

        Ok(Self {
            stage: PortalStage::WaitingForFormat,
            enc: None,
            fps_limit: FpsLimit::new(args.fps),
            cap,
            args,
            errored: false,
            first_frame: true,
            drm_device,
            first_pts_ns: None,
        })
    }

    /// 轮询 PipeWire 事件并编码帧
    ///
    /// 尝试从采集端点接收一帧事件。返回 `Ok(true)` 表示已处理事件，
    /// `Ok(false)` 表示暂无数据。内部根据当前阶段（等待格式/流式）分发处理。
    pub fn poll_and_encode(&mut self) -> Result<bool> {
        let event = match self.cap.frame_receiver().try_recv() {
            Ok(event) => event,
            Err(_) => return Ok(false),
        };

        match event {
            PwEvent::Frame(frame) => {
                match self.stage {
                    PortalStage::WaitingForFormat => {
                        // 第一帧到达：记录格式信息并用该分辨率创建编码器
                        tracing::info!(
                            "First DMA-BUF frame: {}x{} format=0x{:08X} stride={} modifier=0x{:X}",
                            frame.width,
                            frame.height,
                            frame.format,
                            frame.stride,
                            frame.modifier
                        );

                        let drm_path = self.resolve_drm_device_for_frame(&frame)?;
                        let enc = avhw::create_encoder(
                            &drm_path,
                            self.args.output.as_ref(),
                            frame.width,
                            frame.height,
                            self.args.fps,
                            Transform::Normal,
                            self.args.bitrate,
                            self.args.gop_size,
                            None,
                        )?;

                        self.enc = Some(enc);
                        self.stage = PortalStage::Streaming;
                        drop(frame);
                    }
                    PortalStage::Streaming => {
                        // 流式阶段：处理每一帧 DMA-BUF 数据
                        self.handle_pw_frame(frame)?;
                    }
                }
            }
            PwEvent::StreamEnded => {
                // PipeWire 流结束（如用户停止了屏幕共享）
                tracing::warn!("PipeWire stream ended");
                self.errored = true;
            }
            PwEvent::Error(e) => {
                // PipeWire 返回错误
                tracing::error!("PipeWire error: {e}");
                self.errored = true;
            }
        }

        Ok(true)
    }

    fn resolve_drm_device_for_frame(&mut self, frame: &PwDmaBufFrame) -> Result<PathBuf> {
        if let Some(ref drm) = self.drm_device {
            return Ok(drm.clone());
        }

        let candidates = crate::state::find_drm_render_nodes();
        if candidates.is_empty() {
            bail!("No DRM render device found. Specify --drm-device.");
        }

        let mut failures = Vec::new();
        for candidate in &candidates {
            match crate::avhw::test_dma_buf_import(candidate, frame) {
                Ok(()) => {
                    tracing::info!(
                        "Auto-selected DRM device: {} (can import PipeWire DMA-BUF)",
                        candidate.display()
                    );
                    self.drm_device = Some(candidate.clone());
                    return Ok(candidate.clone());
                }
                Err(err) => {
                    tracing::debug!(
                        "DRM device {} cannot import frame: {err:#}",
                        candidate.display()
                    );
                    failures.push(format!("{}: {err:#}", candidate.display()));
                }
            }
        }

        bail!(
            "No DRM render device can import the PipeWire DMA-BUF frame. \
             Specify --drm-device. Tried: {}",
            failures.join("; ")
        )
    }

    /// 处理单帧 DMA-BUF 数据
    ///
    /// 完整的帧处理流水线：
    /// 1. 帧率限制（首帧跳过）
    /// 2. 构建 DRM 描述符
    /// 3. 分配 DRM_PRIME 源帧
    /// 4. 分配 VAAPI 硬件目标帧
    /// 5. 通过 DMA-BUF 导入将帧数据导入 VAAPI
    /// 6. 计算 PTS 时间戳
    /// 7. 回收 DRM 描述符内存
    /// 8. 编码输出
    fn handle_pw_frame(&mut self, frame: PwDmaBufFrame) -> Result<()> {
        // 1. FPS limiting (first frame bypasses)
        // 帧率限制（首帧跳过限制，确保立即编码）
        if self.first_frame {
            self.first_frame = false;
        } else {
            let now = std::time::Instant::now();
            if self.fps_limit.on_new_frame((), now).is_none() {
                return Ok(());
            }
        }

        // 2. Build DRM descriptor for DMA-BUF import
        // 根据 DMA-BUF 帧信息构建 FFmpeg DRM 描述符
        let desc = build_drm_descriptor(&frame);
        let desc_box = Box::new(desc);

        // 3. Allocate raw DRM_PRIME source frame using Video wrapper
        // 分配 DRM_PRIME 格式的源帧，将描述符指针挂载到 data[0]
        let mut raw_frame = ff::frame::Video::empty();
        unsafe {
            let raw_ptr = raw_frame.as_mut_ptr();
            (*raw_ptr).data[0] = Box::into_raw(desc_box) as *mut u8;
            (*raw_ptr).format = ffi::AVPixelFormat::AV_PIX_FMT_DRM_PRIME as i32;
            (*raw_ptr).width = frame.width as i32;
            (*raw_ptr).height = frame.height as i32;
        }

        // 4. Get encoder reference
        // 获取编码器引用
        let enc = match self.enc.as_mut() {
            Some(e) => e,
            None => {
                // Recover the Box to prevent memory leak of the descriptor
                // 编码器未初始化时回收描述符以防止内存泄漏
                unsafe {
                    let desc_ptr = (*raw_frame.as_ptr()).data[0] as *mut ffi::AVDRMFrameDescriptor;
                    if !desc_ptr.is_null() {
                        let _ = Box::from_raw(desc_ptr);
                    }
                }
                bail!("encoder not initialized");
            }
        };

        // 5. Allocate VAAPI hardware target frame
        // 分配 VAAPI 硬件帧缓冲区
        let mut hw_frame = ff::frame::Video::empty();
        let ret = unsafe {
            ffi::av_hwframe_get_buffer(enc.frames_rgb().as_ptr(), hw_frame.as_mut_ptr(), 0)
        };
        if ret < 0 {
            // Recover the Box to prevent memory leak of the descriptor
            // 分配失败时回收描述符防止内存泄漏
            unsafe {
                let desc_ptr = (*raw_frame.as_ptr()).data[0] as *mut ffi::AVDRMFrameDescriptor;
                if !desc_ptr.is_null() {
                    let _ = Box::from_raw(desc_ptr);
                }
            }
            bail!("av_hwframe_get_buffer failed: error {ret}");
        }

        // 6. Import DMA-BUF into VAAPI via transfer_data
        // 通过 DMA-BUF 导入将帧数据从 DRM 传输到 VAAPI 硬件表面
        let ret = unsafe {
            ffi::av_hwframe_transfer_data(hw_frame.as_mut_ptr(), raw_frame.as_ptr(), 0)
        };
        if ret < 0 {
            // 传输失败时回收描述符防止内存泄漏
            unsafe {
                let desc_ptr = (*raw_frame.as_ptr()).data[0] as *mut ffi::AVDRMFrameDescriptor;
                if !desc_ptr.is_null() {
                    let _ = Box::from_raw(desc_ptr);
                }
            }
            if ret == -(ffi::EINVAL as i32) {
                bail!(
                    "VAAPI does not support DMA-BUF modifier 0x{:X}",
                    frame.modifier
                );
            }
            bail!("av_hwframe_transfer_data failed: error {ret}");
        }

        // 7. Set PTS — convert PipeWire nanoseconds to encoder frame-number units
        //    PipeWire PTS is CLOCK_MONOTONIC in nanoseconds.
        //    Encoder time_base = 1/fps, so PTS must be in frame numbers.
        //    Use elapsed time since first frame to avoid i64 overflow on absolute timestamps.
        //
        // PTS 计算：将 PipeWire 的纳秒时间戳转换为编码器的帧号单位
        // PipeWire 使用 CLOCK_MONOTONIC 纳秒时间戳，编码器 time_base = 1/fps
        // 使用相对时间避免绝对时间戳导致的 i64 溢出
        let fps_i64 = self.args.fps as i64;
        let base_ns = *self.first_pts_ns.get_or_insert(frame.pts.max(0));
        let elapsed_ns = (frame.pts.max(0) - base_ns).max(0);
        let pts = elapsed_ns * fps_i64 / 1_000_000_000;
        unsafe {
            (*hw_frame.as_mut_ptr()).pts = pts;
        }

        // 8. Recover the Boxed descriptor from raw_frame *before* encoding.
        //    av_hwframe_transfer_data has already imported the DMA-BUF into the
        //    VAAPI surface, so FFmpeg no longer references the descriptor struct.
        //    Doing this before encode_frame ensures the descriptor is reclaimed
        //    even if encode_frame returns early via `?`.
        //
        // 在编码前回收描述符内存。
        // 此时 DMA-BUF 数据已导入 VAAPI 表面，FFmpeg 不再引用描述符结构体。
        // 在 encode_frame 之前回收确保即使编码返回错误也能正确释放内存。
        unsafe {
            let desc_ptr = (*raw_frame.as_ptr()).data[0] as *mut ffi::AVDRMFrameDescriptor;
            if !desc_ptr.is_null() {
                let _ = Box::from_raw(desc_ptr);
            }
        }

        // 9. Encode — safe to early-return via `?` now that descriptor is recovered.
        // 编码帧数据（此时描述符已回收，可安全通过 `?` 提前返回）
        enc.encode_frame(&hw_frame)?;
 
        // raw_frame and hw_frame drop here via Video::drop → av_frame_free
        // raw_frame 和 hw_frame 在此处通过 Video::drop → av_frame_free 释放
        Ok(())
    }

    /// 刷新编码器缓冲区，输出所有剩余帧
    pub fn flush(&mut self) -> Result<()> {
        if let Some(enc) = &mut self.enc {
            enc.flush()?;
        }
        Ok(())
    }

    /// 关闭状态：刷新编码器并清理资源
    pub fn shutdown(&mut self) {
        if let Err(e) = self.flush() {
            tracing::error!("Flush error during shutdown: {e}");
        }
        tracing::info!("StatePortal shutdown complete");
    }

    /// 返回是否遇到不可恢复的错误
    pub fn is_errored(&self) -> bool {
        self.errored
    }
}

/// 根据 DMA-BUF 帧信息构建 FFmpeg DRM 帧描述符
///
/// 将 PipeWire 提供的 DMA-BUF 参数（fd、偏移量、步长、修饰符等）
/// 转换为 FFmpeg 的 AVDRMFrameDescriptor 结构体，用于零拷贝硬件导入。
fn build_drm_descriptor(frame: &PwDmaBufFrame) -> ffi::AVDRMFrameDescriptor {
    let mut desc: ffi::AVDRMFrameDescriptor = unsafe { mem::zeroed() };

    // DMA-BUF 对象层：一个 fd 对应一个内存对象
    desc.nb_objects = 1;
    desc.objects[0].fd = frame.fd.as_raw_fd();
    desc.objects[0].size = 0; // 大小为 0 表示整个 fd
    desc.objects[0].format_modifier = frame.modifier;

    // 像素格式层：单层单平面布局（如 XR24 格式）
    desc.nb_layers = 1;
    desc.layers[0].format = frame.format;
    desc.layers[0].nb_planes = 1;
    desc.layers[0].planes[0].object_index = 0;
    desc.layers[0].planes[0].offset = frame.offset as isize;
    desc.layers[0].planes[0].pitch = frame.stride as isize;

    desc
}

/// 解析 DRM 渲染设备路径
///
/// 仅使用命令行指定的设备路径；未指定则在首帧到达时自动检测。
fn resolve_drm_device(args: &Args) -> Result<Option<PathBuf>> {
    if let Some(ref drm) = args.drm_device {
        return Ok(Some(PathBuf::from(drm)));
    }
    Ok(None)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::os::fd::{FromRawFd, OwnedFd};

    /// 创建测试用的 DMA-BUF 帧数据（使用 stderr fd 的副本作为占位）
    fn make_test_frame() -> PwDmaBufFrame {
        // Create a dummy fd from stderr (always valid fd 2)
        let fd = unsafe { OwnedFd::from_raw_fd(libc::dup(2)) };
        PwDmaBufFrame {
            fd,
            offset: 0,
            stride: 1920 * 4,
            modifier: 0, // DRM_FORMAT_MOD_LINEAR
            width: 1920,
            height: 1080,
            format: 0x34325258, // XR24 little-endian
            pts: 12345,
        }
    }

    /// 测试 DRM 描述符构建（单平面情况）
    #[test]
    fn build_drm_descriptor_single_plane() {
        let frame = make_test_frame();
        let desc = build_drm_descriptor(&frame);

        assert_eq!(desc.nb_objects, 1);
        assert_eq!(desc.objects[0].format_modifier, 0);
        assert_eq!(desc.nb_layers, 1);
        assert_eq!(desc.layers[0].format, 0x34325258);
        assert_eq!(desc.layers[0].nb_planes, 1);
        assert_eq!(desc.layers[0].planes[0].object_index, 0);
        assert_eq!(desc.layers[0].planes[0].offset, 0);
        assert_eq!(desc.layers[0].planes[0].pitch, 1920 * 4);
    }

    /// 测试显式指定 DRM 设备时的解析
    #[test]
    fn resolve_drm_device_explicit() {
        let args = Args {
            output: "test.mp4".to_string(),
            output_name: None,
            fps: 30,
            codec: "h264".to_string(),
            hw_accel: "vaapi".to_string(),
            drm_device: Some("/dev/dri/renderD128".to_string()),
            bitrate: None,
            gop_size: None,
            verbose: false,
            backend: None,
            port: 0,
        };
        let result = resolve_drm_device(&args);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), std::path::PathBuf::from("/dev/dri/renderD128"));
    }
}