scap-rs 0.1.0

/// Wayland 屏幕投射 Portal 实现
/// 通过 D-Bus 与 XDG Desktop Portal 通信创建屏幕捕获会话
use std::{
    sync::{Arc, Mutex, atomic::AtomicBool},
    time::Duration,
};

use dbus::{
    arg::{self, PropMap, RefArg, Variant},
    blocking::{Connection, Proxy},
    message::MatchRule,
    strings::{BusName, Interface},
};

// 此代码由 `dbus-codegen-rust -d org.freedesktop.portal.Desktop -p /org/freedesktop/portal/desktop -f org.freedesktop.portal.ScreenCast` 自动生成
// 参见 https://github.com/diwic/dbus-rs
// {
use dbus::blocking;

use crate::capturer::engine::linux::error::LinCapError;

/// XDG Desktop Portal ScreenCast 接口 trait
#[allow(unused)]
trait OrgFreedesktopPortalScreenCast {
    /// 创建屏幕投射会话
    fn create_session(&self, options: arg::PropMap) -> Result<dbus::Path<'static>, dbus::Error>;
    /// 选择投射源（显示器或窗口）
    fn select_sources(
        &self,
        session_handle: dbus::Path,
        options: arg::PropMap,
    ) -> Result<dbus::Path<'static>, dbus::Error>;
    /// 启动屏幕投射
    fn start(
        &self,
        session_handle: dbus::Path,
        parent_window: &str,
        options: arg::PropMap,
    ) -> Result<dbus::Path<'static>, dbus::Error>;
    /// 打开 PipeWire 远程连接
    fn open_pipe_wire_remote(
        &self,
        session_handle: dbus::Path,
        options: arg::PropMap,
    ) -> Result<arg::OwnedFd, dbus::Error>;
    /// 获取可用的源类型
    fn available_source_types(&self) -> Result<u32, dbus::Error>;
    /// 获取可用的光标模式
    fn available_cursor_modes(&self) -> Result<u32, dbus::Error>;
    /// 获取版本号
    fn version(&self) -> Result<u32, dbus::Error>;
}

/// 实现 XDG Desktop Portal ScreenCast 接口
impl<'a, T: blocking::BlockingSender, C: ::std::ops::Deref<Target = T>>
    OrgFreedesktopPortalScreenCast for blocking::Proxy<'a, C>
{
    /// 创建屏幕投射会话
    fn create_session(&self, options: arg::PropMap) -> Result<dbus::Path<'static>, dbus::Error> {
        self.method_call(
            "org.freedesktop.portal.ScreenCast",
            "CreateSession",
            (options,),
        ).map(|r: (dbus::Path<'static>,)| r.0)
    }

    /// 选择投射源
    fn select_sources(
        &self,
        session_handle: dbus::Path,
        options: arg::PropMap,
    ) -> Result<dbus::Path<'static>, dbus::Error> {
        self.method_call(
            "org.freedesktop.portal.ScreenCast",
            "SelectSources",
            (session_handle, options),
        ).map(|r: (dbus::Path<'static>,)| r.0)
    }

    /// 启动屏幕投射
    fn start(
        &self,
        session_handle: dbus::Path,
        parent_window: &str,
        options: arg::PropMap,
    ) -> Result<dbus::Path<'static>, dbus::Error> {
        self.method_call(
            "org.freedesktop.portal.ScreenCast",
            "Start",
            (session_handle, parent_window, options),
        ).map(|r: (dbus::Path<'static>,)| r.0)
    }

    /// 打开 PipeWire 远程连接
    fn open_pipe_wire_remote(
        &self,
        session_handle: dbus::Path,
        options: arg::PropMap,
    ) -> Result<arg::OwnedFd, dbus::Error> {
        self.method_call(
            "org.freedesktop.portal.ScreenCast",
            "OpenPipeWireRemote",
            (session_handle, options),
        ).map(|r: (arg::OwnedFd,)| r.0)
    }

    /// 获取可用的源类型
    fn available_source_types(&self) -> Result<u32, dbus::Error> {
        <Self as blocking::stdintf::org_freedesktop_dbus::Properties>::get(
            self,
            "org.freedesktop.portal.ScreenCast",
            "AvailableSourceTypes",
        )
    }

    /// 获取可用的光标模式
    fn available_cursor_modes(&self) -> Result<u32, dbus::Error> {
        <Self as blocking::stdintf::org_freedesktop_dbus::Properties>::get(
            self,
            "org.freedesktop.portal.ScreenCast",
            "AvailableCursorModes",
        )
    }

    /// 获取版本号
    fn version(&self) -> Result<u32, dbus::Error> {
        <Self as blocking::stdintf::org_freedesktop_dbus::Properties>::get(
            self,
            "org.freedesktop.portal.ScreenCast",
            "version",
        )
    }
}
// }

// This code was autogenerated with `dbus-codegen-rust --file org.freedesktop.portal.Request.xml`, see https://github.com/diwic/dbus-rs
// {
/// XDG Desktop Portal Request 接口 trait
// 此代码由 `dbus-codegen-rust --file org.freedesktop.portal.Request.xml` 自动生成
// 参见 https://github.com/diwic/dbus-rs
// {
#[allow(unused)]
trait OrgFreedesktopPortalRequest {
    /// 关闭请求
    fn close(&self) -> Result<(), dbus::Error>;
}

/// 请求响应数据结构
#[derive(Debug)]
pub struct OrgFreedesktopPortalRequestResponse {
    pub response: u32,         // 响应代码
    pub results: arg::PropMap, // 结果属性映射
}

/// 实现 D-Bus 参数追加
impl arg::AppendAll for OrgFreedesktopPortalRequestResponse {
    fn append(&self, i: &mut arg::IterAppend) {
        arg::RefArg::append(&self.response, i);
        arg::RefArg::append(&self.results, i);
    }
}

/// 实现 D-Bus 参数读取
impl arg::ReadAll for OrgFreedesktopPortalRequestResponse {
    fn read(i: &mut arg::Iter) -> Result<Self, arg::TypeMismatchError> {
        Ok(OrgFreedesktopPortalRequestResponse {
            response: i.read()?,
            results: i.read()?,
        })
    }
}

/// 实现 D-Bus 信号参数
impl dbus::message::SignalArgs for OrgFreedesktopPortalRequestResponse {
    const NAME: &'static str = "Response";
    const INTERFACE: &'static str = "org.freedesktop.portal.Request";
}

/// 实现 XDG Desktop Portal Request 接口
impl<'a, T: blocking::BlockingSender, C: ::std::ops::Deref<Target = T>> OrgFreedesktopPortalRequest
    for blocking::Proxy<'a, C>
{
    /// 关闭请求
    fn close(&self) -> Result<(), dbus::Error> {
        self.method_call("org.freedesktop.portal.Request", "Close", ())
    }
}
// }

/// 请求响应类型别名
type Response = Option<OrgFreedesktopPortalRequestResponse>;

/// 流信息字典结构
#[derive(Debug)]
#[allow(dead_code)]
pub struct StreamVardict {
    id: Option<String>,           // 流 ID
    position: Option<(i32, i32)>, // 流位置
    size: Option<(i32, i32)>,     // 流尺寸
    source_type: Option<u32>,     // 源类型
    mapping_id: Option<String>,   // 映射 ID
}

/// 屏幕投射流结构
#[derive(Debug)]
#[allow(unused)]
pub struct Stream(u32, StreamVardict);

impl Stream {
    /// 获取 PipeWire 节点 ID
    pub fn pw_node_id(&self) -> u32 {
        self.0
    }

    /// 从 D-Bus 变体解析流信息
    pub fn from_dbus(stream: &Variant<Box<dyn RefArg>>) -> Option<Self> {
        let mut stream = stream.as_iter()?.next()?.as_iter()?;
        let pipewire_node_id = stream.next()?.as_iter()?.next()?.as_u64()?;

        // TODO: 获取其余属性

        Some(Self(
            pipewire_node_id as u32,
            StreamVardict {
                id: None,
                position: None,
                size: None,
                source_type: None,
                mapping_id: None,
            },
        ))
    }
}

/// 匹配响应代码宏
macro_rules! match_response {
    ( $code:expr ) => {
        match $code {
            0 => {} // 成功
            1 => {
                return Err(LinCapError::new(String::from("用户取消了交互")));
            }
            2 => {
                return Err(LinCapError::new(String::from("用户交互以其他方式结束")));
            }
            _ => unreachable!(),
        }
    };
}

/// 屏幕投射 Portal 结构
/// 管理与 XDG Desktop Portal 的 D-Bus 通信
pub struct ScreenCastPortal<'a> {
    proxy: Proxy<'a, &'a Connection>, // D-Bus 代理
    token: String,                    // 请求令牌
    cursor_mode: u32,                 // 光标模式
}

impl<'a> ScreenCastPortal<'a> {
    /// 创建新的 ScreenCastPortal 实例
    pub fn new(connection: &'a Connection) -> Self {
        // 创建 D-Bus 代理连接
        let proxy = connection.with_proxy(
            "org.freedesktop.portal.Desktop",
            "/org/freedesktop/portal/desktop",
            Duration::from_secs(4),
        );

        // 生成随机令牌
        let token = format!("scap_{}", rand::random::<u16>());

        Self {
            proxy,
            token,
            cursor_mode: 1, // 默认光标模式
        }
    }

    /// 创建会话参数
    fn create_session_args(&self) -> arg::PropMap {
        let mut map = arg::PropMap::new();
        map.insert(
            String::from("handle_token"),
            Variant(Box::new(self.token.clone())),
        );
        map.insert(
            String::from("session_handle_token"),
            Variant(Box::new(self.token.clone())),
        );
        map
    }

    /// 选择源参数
    fn select_sources_args(&self) -> Result<arg::PropMap, dbus::Error> {
        let mut map = arg::PropMap::new();
        map.insert(
            String::from("handle_token"),
            Variant(Box::new(self.token.clone())),
        );
        map.insert(
            String::from("types"),
            Variant(Box::new(self.proxy.available_source_types()?)),
        );
        map.insert(String::from("multiple"), Variant(Box::new(false)));
        map.insert(
            String::from("cursor_mode"),
            Variant(Box::new(self.cursor_mode)),
        );
        Ok(map)
    }

    /// 处理请求响应
    fn handle_req_response(
        connection: &Connection,
        path: dbus::Path<'static>,
        iterations: usize,
        timeout: Duration,
        response: Arc<Mutex<Response>>,
    ) -> Result<(), dbus::Error> {
        let got_response = Arc::new(AtomicBool::new(false));
        let got_response_clone = Arc::clone(&got_response);

        // 设置 D-Bus 信号匹配规则
        let mut rule = MatchRule::new();
        rule.path = Some(path);
        rule.msg_type = Some(dbus::MessageType::Signal);
        rule.sender = Some(BusName::from("org.freedesktop.portal.Desktop"));
        rule.interface = Some(Interface::from("org.freedesktop.portal.Request"));
        connection.add_match(
            rule,
            move |res: OrgFreedesktopPortalRequestResponse, _chuh, _msg| {
                let mut response = response.lock().expect("锁定响应互斥锁失败");
                *response = Some(res);
                got_response_clone.store(true, std::sync::atomic::Ordering::Relaxed);
                false
            },
        )?;

        // 等待响应
        for _ in 0..iterations {
            connection.process(timeout)?;

            if got_response.load(std::sync::atomic::Ordering::Relaxed) {
                break;
            }
        }

        Ok(())
    }

    /// 创建屏幕投射会话
    fn create_session(&self) -> Result<dbus::Path<'_>, LinCapError> {
        let request_handle = self.proxy.create_session(self.create_session_args())?;

        let response = Arc::new(Mutex::new(None));
        let response_clone = Arc::clone(&response);
        Self::handle_req_response(
            self.proxy.connection,
            request_handle,
            100,
            Duration::from_millis(100),
            response_clone,
        )?;

        if let Some(res) = response.lock()?.take() {
            match_response!(res.response);
            match res
                .results
                .get("session_handle")
                .map(|h| h.0.as_str().map(String::from))
            {
                Some(h) => {
                    let p = dbus::Path::from(match h {
                        Some(p) => p,
                        None => {
                            return Err(LinCapError::new(String::from(
                                "收到无效的 session_handle",
                            )));
                        }
                    });

                    return Ok(p);
                }
                None => return Err(LinCapError::new(String::from("未获取到 session handle"))),
            }
        }

        Err(LinCapError::new(String::from("未获取到响应")))
    }

    /// 选择投射源
    fn select_sources(&self, session_handle: dbus::Path) -> Result<(), LinCapError> {
        let request_handle = self
            .proxy
            .select_sources(session_handle, self.select_sources_args()?)?;

        let response = Arc::new(Mutex::new(None));
        let response_clone = Arc::clone(&response);
        Self::handle_req_response(
            self.proxy.connection,
            request_handle,
            1200, // 等待 2 分钟
            Duration::from_millis(100),
            response_clone,
        )?;

        if let Some(res) = response.lock()?.take() {
            match_response!(res.response);
            return Ok(());
        }

        Err(LinCapError::new(String::from("未获取到响应")))
    }

    /// 启动屏幕投射
    fn start(&self, session_handle: dbus::Path) -> Result<Stream, LinCapError> {
        let request_handle = self.proxy.start(session_handle, "", PropMap::new())?;

        let response = Arc::new(Mutex::new(None));
        let response_clone = Arc::clone(&response);
        Self::handle_req_response(
            self.proxy.connection,
            request_handle,
            100, // 等待 10 秒
            Duration::from_millis(100),
            response_clone,
        )?;

        if let Some(res) = response.lock()?.take() {
            match_response!(res.response);
            match res.results.get("streams") {
                Some(s) => match Stream::from_dbus(s) {
                    Some(s) => return Ok(s),
                    None => return Err(LinCapError::new(String::from("提取流属性失败"))),
                },
                None => return Err(LinCapError::new(String::from("未获取到任何流"))),
            }
        }

        Err(LinCapError::new(String::from("未获取到响应")))
    }

    /// 创建完整的屏幕投射流
    /// 依次执行：创建会话 -> 选择源 -> 启动投射
    pub fn create_stream(&self) -> Result<Stream, LinCapError> {
        let session_handle = self.create_session()?;
        self.select_sources(session_handle.clone())?;
        self.start(session_handle)
    }

    /// 设置是否显示光标
    pub fn show_cursor(mut self, mode: bool) -> Result<Self, LinCapError> {
        let available_modes = self.proxy.available_cursor_modes()?;
        if mode && available_modes & 2 == 2 {
            self.cursor_mode = 2; // 显示光标
            return Ok(self);
        }
        if !mode && available_modes & 1 == 1 {
            self.cursor_mode = 1; // 隐藏光标
            return Ok(self);
        }

        Err(LinCapError::new("不支持的光标模式".to_string()))
    }
}