Haocode/ui/views/system_tools/portal_capture.py

# -*- coding: utf-8 -*-
"""
xdg-desktop-portal 截图（P1-04，Linux Wayland 专用窄适配器）

Wayland compositor 不允许应用直接抓取屏幕画面（安全模型），截图必须走
xdg-desktop-portal 的 org.freedesktop.portal.Screenshot：
  1. gdbus 调 Screenshot(handle, parent_window=/, {}) → 返回 request 对象路径；
  2. gdbus monitor 监听该 request 对象：
       FilePicked(file_uri, ...) → 用户已授权并保存成功 → 返回文件路径；
       Finished（无 FilePicked）   → 用户取消/授权被拒 → 返回 denied；
  3. 全程有显式时间预算；超时 → timeout。

实现只用系统自带 CLI（gdbus，GLib 的一部分）+ 标准库，不新增 Python 依赖、
不绕过 compositor、不伪造成功。portal/桌面环境不支持时返回明确的不可用原因。
"""
import json
import os
import shutil
import subprocess
import sys
import time
import urllib.parse

from PyQt6 import QtCore

PORTAL_DEST = "org.freedesktop.portal.Desktop"
PORTAL_PATH = "/org/freedesktop/portal/desktop"
SCREENSHOT_IFACE = "org.freedesktop.portal.Screenshot"
REQUEST_IFACE = "org.freedesktop.portal.Request"


def detect_portal():
    """探测 portal 截图能力 → (ok: bool, reason: str | None)。"""
    if sys.platform != "linux":
        return False, "portal 截图仅用于 Linux"
    if not (os.environ.get("XDG_RUNTIME_DIR", "").strip()
            or os.environ.get("DBUS_SESSION_BUS_ADDRESS", "").strip()):
        return False, ("未检测到 D-Bus 会话（XDG_RUNTIME_DIR/DBUS_SESSION_BUS_ADDRESS 均缺失）"
                       " → portal 截图不可用")
    if not shutil.which("gdbus"):
        return False, ("未找到 gdbus CLI（GLib 组件）→ portal 截图不可用；"
                       "聊天与其他功能不受影响")
    return True, None


def _gdbus_screenshot_request(handle: str, timeout_s: float):
    """调 Screenshot 方法，返回 (request_path, stderr)；失败抛 RuntimeError。"""
    cmd = ["gdbus", "call", "--session", "--dest", PORTAL_DEST,
           "--object-path", PORTAL_PATH,
           f"--method={SCREENSHOT_IFACE}.Screenshot",
           handle, "/", "{}"]
    try:
        p = subprocess.run(cmd, capture_output=True, text=True,
                           encoding="utf-8", errors="replace",
                           timeout=timeout_s, check=False)
    except subprocess.TimeoutExpired:
        raise RuntimeError(f"portal Screenshot 调用超时（{timeout_s:.0f}s 预算）")
    if p.returncode != 0:
        err = (p.stderr or p.stdout or "").strip().splitlines()
        raise RuntimeError("portal 不可用: " + (err[-1] if err else "未知错误"))
    out = (p.stdout or "").strip()
    # gdbus 输出形如：('/org/freedesktop/portal/desktop/request/1000/haocode/42', <>)
    start = out.find("('")
    end = out.find("',") if start != -1 else -1
    if start == -1 or end == -1:
        raise RuntimeError(f"无法解析 Screenshot 返回: {out[:120]!r}")
    return out[start + 2:end], (p.stderr or "").strip()


def _gdbus_monitor_request(request_path: str, budget_s: float):
    """监听 request 对象直到 FilePicked / Finished / 预算耗尽。

    返回 (status, path_or_reason)：status ∈ {"ok", "denied", "timeout"}。
    """
    cmd = ["gdbus", "monitor", "--session", "--object-path", request_path]
    p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                         text=True, encoding="utf-8", errors="replace")
    deadline = time.monotonic() + budget_s
    file_path = None
    finished = False
    try:
        while time.monotonic() < deadline:
            line = p.stdout.readline()
            if not line:
                if p.poll() is not None:
                    break
                continue
            line = line.strip()
            if not line.startswith("{"):
                continue
            try:
                ev = json.loads(line)
            except Exception:
                continue
            member = ev.get("member", "")
            body = ev.get("body", [])
            if member == "FilePicked" and body:
                uri = str(body[0])
                # 剥前缀 + unquote（比 urlparse 稳健：兼容 file:///tmp/x 与 file://D%3A%5Cx）
                file_path = urllib.parse.unquote(uri[len("file://"):]) \
                    if uri.startswith("file://") else uri
            elif member == "Finished" and ev.get("interface", "") == REQUEST_IFACE:
                finished = True
                break
        if file_path:
            return "ok", file_path
        if finished and time.monotonic() <= deadline:
            return "denied", "用户在 portal 授权窗口取消/拒绝"
        return "timeout", f"等待 portal 授权超时（{budget_s:.0f}s 预算）"
    finally:
        try:
            p.kill()
        except Exception:
            pass
        try:
            p.wait(timeout=3)
        except Exception:
            pass


def portal_screenshot_sync(handle: str = "haocode-shot",
                           request_timeout_s: float = 10.0,
                           wait_budget_s: float = 120.0):
    """同步执行 portal 截图（调用方须保证不在 GUI 线程阻塞，或用 PortalScreenshotWorker）。

    返回 (ok: bool, path_or_reason: str)。
    """
    ok, reason = detect_portal()
    if not ok:
        return False, reason
    try:
        request_path, _err = _gdbus_screenshot_request(handle, request_timeout_s)
    except RuntimeError as e:
        return False, str(e)
    status, result = _gdbus_monitor_request(request_path, wait_budget_s)
    if status == "ok":
        return True, result
    if status == "denied":
        return False, result
    return False, result


class PortalScreenshotWorker(QtCore.QThread):
    """GUI 线程安全的 portal 截图：done(ok, path_or_reason) 信号回主线程。"""

    done = QtCore.pyqtSignal(bool, str)

    def __init__(self, handle: str = "haocode-shot",
                 request_timeout_s: float = 10.0,
                 wait_budget_s: float = 120.0, parent=None):
        super().__init__(parent)
        self._handle = handle
        self._request_timeout_s = request_timeout_s
        self._wait_budget_s = wait_budget_s

    def run(self):
        try:
            ok, result = portal_screenshot_sync(self._handle,
                                                self._request_timeout_s,
                                                self._wait_budget_s)
        except Exception as e:
            ok, result = False, f"portal 截图异常（{type(e).__name__}: {e}）"
        self.done.emit(ok, result)