project/tbsg/snapshot.py

"""
网页快照模块 — 模仿美团 MTScreenshot：
详情页滚动截图拼接长图 → 压缩 → 上传阿里云OSS → 删除本地临时文件 → 返回OSS URL
任何失败返回空字符串，不阻塞主流程
"""
import io
import json
import os
import random
import re
import subprocess
import time
from pathlib import Path

import oss2
from PIL import Image

from human_touch import wrist_arc_pts, motionevent_drag, touchpipe_drag
from commons import err_log

CONFIG_PATH = Path(__file__).parent / "config.json"
SCREENSHOT_DIR = Path(__file__).parent / "scrape_data"

RESIZE_RATIO = 0.8     # 长图缩放比例（美团同款）
JPEG_QUALITY = 70      # 压缩质量
DEFAULT_SCROLLS = 1    # 首屏后上滑1次（共2屏），慢速滑动

_bucket = None


def _load_oss_config() -> dict:
    cfg = {}
    try:
        with open(CONFIG_PATH, encoding="utf-8") as f:
            cfg = json.load(f)
    except Exception:
        pass
    return cfg.get("oss", {}) or {}


def _get_bucket():
    global _bucket
    if _bucket is not None:
        return _bucket
    c = _load_oss_config()
    if not all([c.get("access_key_id"), c.get("access_key_secret"), c.get("endpoint"), c.get("bucket_name")]):
        print("[snapshot] OSS配置不完整，跳过快照")
        _bucket = None
        return None
    try:
        auth = oss2.Auth(c["access_key_id"], c["access_key_secret"])
        _bucket = oss2.Bucket(auth, c["endpoint"], c["bucket_name"])
        print(f"[snapshot] OSS连接成功: {c['bucket_name']}")
    except Exception as e:
        print(f"[snapshot] OSS连接失败: {e}")
        _bucket = None
    return _bucket


def _upload_to_oss(local_path: str):
    """上传OSS，成功返回完整URL，失败返回None（美团同款）"""
    bucket = _get_bucket()
    if not bucket or not os.path.exists(local_path):
        return None
    c = _load_oss_config()
    file_name = os.path.basename(local_path)
    safe_name = re.sub(r'[^\w\.\-]', '_', file_name)
    oss_key = f"{c.get('oss_prefix', 'scrape_data/')}{safe_name}"
    try:
        oss2.resumable_upload(bucket, oss_key, local_path)
        return f"https://{c['bucket_name']}.{c['endpoint']}/{oss_key}"
    except Exception as e:
        print(f"[snapshot] OSS上传失败: {e}")
        return None


def _merge_screenshots(screen_bytes_list):
    """竖接长图（美团同款思路，PIL实现）"""
    imgs = []
    for b in screen_bytes_list:
        try:
            imgs.append(Image.open(io.BytesIO(b)).convert("RGB"))
        except Exception:
            pass
    if not imgs:
        return None
    w = max(i.width for i in imgs)
    total_h = sum(i.height for i in imgs)
    merged = Image.new("RGB", (w, total_h), (255, 255, 255))
    y = 0
    for i in imgs:
        merged.paste(i, (0, y))
        y += i.height
        i.close()
    return merged


def collect_snapshot(driver, title: str = "", device_id: str = "", scroll_times: int = DEFAULT_SCROLLS) -> str:
    """
    详情页滚动截图拼接 → 压缩 → 上传OSS → 删本地 → 返回OSS URL。
    返回 (url, reason)：url为空时 reason 是失败原因。
    """
    try:
        import cv2
        SCREENSHOT_DIR.mkdir(parents=True, exist_ok=True)   # 中间帧截图落盘前先建目录
        w, h = driver.window_size()
        # 1. 滚动截图：首屏 + 1次上滑（共2屏，慢速、间隔充足）
        #    每屏落盘并用cv2验证PNG完整性（adb流式传输可能损坏，损坏则重试）
        screen_bytes_list = []
        for i in range(scroll_times + 1):
            shot = str(SCREENSHOT_DIR / f"_snap_{device_id}_{i}.png")
            ok = False
            for _try in range(3):
                driver.screenshot(shot)
                if cv2.imread(shot) is not None:
                    ok = True
                    break
                time.sleep(0.5)
            if not ok:
                print(f"[snapshot] 第{i+1}屏截图损坏，放弃快照")
                return "", f"第{i+1}屏截图损坏(重试3次)"
            with open(shot, "rb") as f:
                screen_bytes_list.append(f.read())
            os.remove(shot)  # 中间帧用完即删，不占空间
            if i < scroll_times:
                # 拟人弧线滑动(手腕枢轴模型); TouchPipe→motionevent→直线 三级降级
                try:
                    pts = wrist_arc_pts(w, h, int(h * 0.5), n=60, drift_range=(50, 110))
                    if not touchpipe_drag(driver, pts):
                        print("[snapshot] TouchPipe未生效,降级motionevent")
                        pts10 = wrist_arc_pts(w, h, int(h * 0.5), n=10, drift_range=(50, 110))
                        motionevent_drag(device_id, pts10)
                except Exception as e:
                    print(f"[snapshot] 拟人滑动异常({e}),退回直线滑动")
                    subprocess.run(
                        ["adb", "-s", device_id, "shell", "input", "swipe",
                         str(w // 2), str(int(h * 0.75)), str(w // 2), str(int(h * 0.25)), "600"],
                        capture_output=True, timeout=10,
                    )
                time.sleep(random.uniform(1.8, 2.4))
        print(f"[snapshot] 已截 {len(screen_bytes_list)} 屏，拼接中...")

        # 2. 拼接 + 压缩
        merged = _merge_screenshots(screen_bytes_list)
        if merged is None:
            print("[snapshot] 截图拼接失败")
            return "", "截图拼接失败"
        if 0 < RESIZE_RATIO < 1.0:
            merged = merged.resize(
                (int(merged.width * RESIZE_RATIO), int(merged.height * RESIZE_RATIO)),
                Image.LANCZOS,
            )

        # 3. 保存本地临时文件（命名与美团一致：时间戳_平台标识_设备ID_标题）
        SCREENSHOT_DIR.mkdir(exist_ok=True)
        ts = time.strftime("%Y%m%d_%H%M%S")
        safe_title = re.sub(r'[\\/*?:"<>|]', '_', str(title or "snapshot"))[:40]
        local_path = str(SCREENSHOT_DIR / f"{ts}_tbsg_{device_id}_{safe_title}.jpg")
        merged.save(local_path, format="JPEG", quality=JPEG_QUALITY)
        merged.close()

        # 4. 上传OSS
        url = _upload_to_oss(local_path)

        # 5. 上传成功后删除本地文件（美团同款：不占本地空间）
        try:
            os.remove(local_path)
        except Exception:
            pass
        if not url:
            print(f"[snapshot] 上传失败，本地文件保留: {local_path}")
            return "", f"OSS上传失败({local_path})"
        print(f"[snapshot] 上传成功: {url}")
        return url, ""
    except Exception as e:
        print(f"[snapshot] 快照采集失败: {e}")
        err_log.log_error(device_id, "snapshot_error", exc=e, extra={"title": title})
        return "", f"异常: {e}"