auto_control/core/account_probe.py

"""实时探测"这台设备现在登的是哪个抖音号"（**按名字找设备**用）。

**谁用它**：人工入口 —— 「发布计划」页的「一键推送（按名字找设备）」要先知道
"哪台在线设备现在登着这个名字"，才能把素材推过去。
**自动任务不靠它**：任务里的「读取本机账号名字」永远是现场读屏幕（缓存旧了也不能影响发布正确性）。

它只做**只读/幂等**的动作：连设备 → 需要时把抖音拉到前台 → 点「我」→ 读昵称控件 → 清洗。
**绝不登录、绝不滑块、绝不改设备设置**；连不上/读不到只报那一台，不影响别的设备。
也绝不 `kill-server` / `disconnect`（技术约束红线）。

昵称控件在 `config.DOUYIN_NAME_SELECTOR_VALUE`（**抓一次填 .env**，与任务骨架同一处配置）；
没填就明确报"还没抓"，不去猜。
"""
import time
from concurrent.futures import ThreadPoolExecutor

from config import DOUYIN_NAME_SELECTOR_TYPE, DOUYIN_NAME_SELECTOR_VALUE
from core import u2_helper
from core.logger import get_logger

_log = get_logger("core.account_probe")

DOUYIN_PKG = "com.ss.android.ugc.aweme"
_CONNECT_TIMEOUT = 20        # u2.connect 超时（秒）；设备端 u2 服务没响应时别把请求挂死
_STEP_WAIT = 2.5             # 进「我」页面后等它渲染完


def _connect(serial, timeout=_CONNECT_TIMEOUT):
    """带超时的 u2.connect。

    ⚠ 用**手动 shutdown(wait=False)** —— `with ThreadPoolExecutor(...)` 退出时会
    `shutdown(wait=True)`，超时线程不死它就跟着等，30 秒的"超时"会变成真卡住
    （`core/device_worker` 里踩过这个坑）。
    """
    import uiautomator2 as u2
    pool = ThreadPoolExecutor(max_workers=1)
    try:
        return pool.submit(u2.connect, serial).result(timeout=timeout)
    finally:
        pool.shutdown(wait=False)


def probe_one(serial, click_me=True):
    """读一台设备当前登录的抖音名字。返回 `{serial, name, error}`。

    `name` 非空 = 读到了；`error` 非空 = 没读到（人话原因，给界面直接显示）。
    """
    if not DOUYIN_NAME_SELECTOR_VALUE:
        return {"serial": serial, "name": "",
                "error": "还没抓抖音昵称控件：在 .env 里填 DOUYIN_NAME_SELECTOR_VALUE"
                         "（抓取元素抓一下「我」页面上的昵称）"}
    d = None
    try:
        d = _connect(serial)
    except Exception as e:
        return {"serial": serial, "name": "",
                "error": f"连不上设备（u2 {_CONNECT_TIMEOUT}s 超时）：{type(e).__name__}"}
    try:
        pkg = u2_helper.current_package(d)
        if pkg != DOUYIN_PKG:
            # 不在抖音 → 拉到前台（幂等：已经在就什么都不做）
            try:
                d.app_start(DOUYIN_PKG, wait=True)
            except Exception as e:
                return {"serial": serial, "name": "",
                        "error": f"打不开抖音：{type(e).__name__}: {str(e)[:60]}"}
            time.sleep(4)
        if click_me:
            u2_helper.find_and_click(d, timeout=3,
                                     **{"descriptionContains": "我"})
        time.sleep(_STEP_WAIT)
        from core import video_plan as vp
        if DOUYIN_NAME_SELECTOR_TYPE == "xpath":
            el = d.xpath(DOUYIN_NAME_SELECTOR_VALUE)
            raw = el.get_text() if el.exists else ""
        else:
            el = d(**{DOUYIN_NAME_SELECTOR_TYPE: DOUYIN_NAME_SELECTOR_VALUE})
            raw = el.get_text(timeout=3) if el.exists(timeout=3) else ""
        name = vp.clean_account_name(raw)
        if not name:
            return {"serial": serial, "name": "",
                    "error": f"读到空（元素原文 {raw!r}）—— 昵称控件的选择器可能抓错了"}
        return {"serial": serial, "name": name, "error": ""}
    except Exception as e:
        return {"serial": serial, "name": "",
                "error": f"读取失败：{type(e).__name__}: {str(e)[:80]}"}
    finally:
        # **不 disconnect、不 kill-server**（红线）：u2 的连接自己会超时回收
        pass


def busy_serials():
    """正在跑任务/正在连接的设备（探测要跳过它们：抢屏会打断任务）。"""
    from core import device_worker
    out = set()
    try:
        for s in device_worker.get_all_worker_status():
            if (s.get("status") or "") in ("running", "connecting"):
                out.add(s.get("serial") or "")
    except Exception as e:
        _log.warning(f"读设备运行状态失败（按「都不忙」处理）: {e}")
    return {x for x in out if x}


def probe_online(serials=None, skip_busy=True, workers=4):
    """并行探测多台设备。返回 `{serial: {serial, name, error}}`。

    · `serials=None` → 取设备池里**在线且启用**的设备（`device_pool.list_ready()`）
    · `skip_busy=True`（默认）→ 跳过正在跑任务/连接中的设备（别抢它们的屏）
    · 单台失败只影响那一台；顺序按传入/池里的顺序
    """
    from core import device_pool
    if serials is None:
        try:
            serials = list(device_pool.list_ready())
        except Exception as e:
            _log.warning(f"取在线设备失败: {e}")
            return {}
    serials = [s for s in (serials or []) if s]
    if skip_busy:
        busy = busy_serials()
        if busy:
            _log.info(f"探测跳过正在跑任务的设备：{'、'.join(sorted(busy))}")
        serials = [s for s in serials if s not in busy]
    out = {}
    if not serials:
        return out
    with ThreadPoolExecutor(max_workers=max(1, min(workers, len(serials)))) as pool:
        for r in pool.map(probe_one, serials):
            out[r["serial"]] = r
    return out


def cache_names(result):
    """把探测结果写进 `device.cur_account_name/cur_name_at`（**只用于人工入口的设备建议**）。

    ⚠ 必须在 **Flask app context** 里调用（web 请求线程本身就有；脚本里自己推）。
    探测失败的那台**不清空旧值**（旧值好歹是个线索，界面上带时间戳显示）。
    """
    from core.models import Device, db
    try:
        now = time.strftime("%Y-%m-%d %H:%M:%S")
        for serial, r in (result or {}).items():
            name = (r or {}).get("name") or ""
            if not name:
                continue
            row = Device.query.get(serial)
            if not row:
                continue
            row.cur_account_name = name[:80]
            row.cur_name_at = now
        db.session.commit()
    except Exception as e:
        _log.warning(f"写设备当前账号缓存失败（不影响探测结果）: {e}")
        try:
            db.session.rollback()
        except Exception:
            pass