auto_control/core/device_pool.py

"""设备池：SQLite devices 表 = 设备清单，本地 adb = 在线状态（STF 替代数据源）。

设计：
  - list_configured() — devices 表里 enabled 的设备（清单，管理页维护）
  - list_online()     — 本机 `adb devices` 里 state=device 的设备（实时；
                        池内有 USB 设备时合并 220 远程 adb server 状态）
  - list_ready()      — configured ∩ online（任务调度用）
  - CRUD              — add/update/remove/set_enabled（设备池管理页用）

红线：全模块不 connect / 不 kill-server / 不 disconnect，遵守既有技术约束。
后台线程（task_manager 等）调用时由本模块自行推 app context，调用方无需关心。
"""
import time
import threading
import subprocess

from config import ADB_PATH, USB_ADB_HOST, USB_ADB_PORT
from core.adb_helper import _adb, _adb_remote
from core.logger import get_logger
from core.models import db, Device

_log = get_logger("core.device_pool")

_app = None
# 设备换地址（认领/迁址）后的回调：(old_serial, new_serial) -> None
# 由 web_server 装配时注册为 TaskManager.sync_device_serial——
# 分组/任务的内存副本在 TaskManager 里，只改库不改内存不生效（见 _move_device_row）
_move_hook = None


def set_move_hook(fn):
    """注册"设备换地址"回调（web_server 创建完 TaskManager 后调用）。"""
    global _move_hook
    _move_hook = fn


def init_app(app):
    """web_server 启动时调用：绑定 Flask app（供后台线程推 db context），
    并后台刷新一次设备型号（首次启动/设备更换后型号可能变化）。"""
    global _app
    _app = app
    refresh_names()          # 同步刷一次名称快照，让通知从第一条起就有名字
    try:
        t = threading.Thread(target=_refresh_models_bg, daemon=True)
        t.start()
        t2 = threading.Thread(target=_names_loop, name="device-names", daemon=True)
        t2.start()
    except Exception:
        pass


# ================== 名称内存快照（给通知用） ==================
# 为什么单独存一份：webhook 通知里要显示"是哪台设备"（A08）而不是地址（IP），
# 但 `notifier.notify()` 有一条硬红线——**零 DB 访问**（不能为了取个名字去查库，
# 更不能在业务线程里阻塞）。所以由这里维护一份 serial→名称 的内存快照：
# 启动刷一次、池子有变动时刷一次、再兜底每 60s 刷一次（覆盖整库恢复等外部改动）。
_names = {}
_names_lock = threading.Lock()
_NAMES_INTERVAL = 60


def refresh_names():
    """把设备池的 serial→名称刷进内存快照（只查库，不碰 adb）。返回条数。"""
    try:
        rows = list_devices()
    except Exception as e:
        _log.debug(f"刷新设备名快照失败: {e}")
        return 0
    m = {}
    for d in rows:
        s = d.get("serial") or ""
        if s:
            m[s] = d.get("name") or ""
    with _names_lock:
        _names.clear()
        _names.update(m)
    return len(m)


def name_of(serial):
    """serial 对应的设备名称（**纯内存，可在通知路径上调用**）。没有则返回 ""。"""
    if not serial:
        return ""
    with _names_lock:
        return _names.get(str(serial), "")


def _names_loop():
    """兜底定时刷新名称快照（覆盖整库恢复这类进程外改动）。"""
    while True:
        time.sleep(_NAMES_INTERVAL)
        refresh_names()


def _refresh_models_bg():
    """后台批量采集在线设备型号（启动时/手动触发）。失败静默，不影响启动。"""
    time.sleep(3)  # 等服务起来再查
    try:
        refresh_all_models()
    except Exception as e:
        _log.warning(f"启动型号采集失败（可稍后手动采集）: {e}")


def _ctx():
    """后台线程访问 db 需自行推 app context（与 task_manager._db 同思路）。"""
    if _app is None:
        raise RuntimeError("device_pool 未关联 Flask app（web_server 启动时调用 init_app）")
    return _app.app_context()


# ================== 清单 ==================
def list_configured(enabled_only=True):
    """devices 表里的设备清单（serial 列表）。"""
    with _ctx():
        q = Device.query
        if enabled_only:
            q = q.filter_by(enabled=True)
        return [d.serial for d in q.order_by(Device.serial).all()]


def list_devices():
    """全量设备行（设备池管理页用）。"""
    with _ctx():
        return [d.to_dict() for d in Device.query.order_by(Device.serial).all()]


def serial_of(device_name):
    """设备名 → **当前**地址（拿不到返回空串）。

    从已删掉的 `core/ledger.py` 迁过来的（账号台账停用后，那边只剩这一个函数还有用 ——
    它跟账号无关，只做"设备名 ↔ 地址"）。

    ⚠ 为什么不能直接用计划行里的 `serial` 快照：设备地址会变（换 IP、重连）。
    快照是"录入时"的值，久了就是错的 —— 要下发操作（推视频等）时必须现查。
    优先设备池的**内存快照**（零 DB、已经是刷新过的），再退回 DB。
    """
    name = (device_name or "").strip()
    if not name:
        return ""
    try:
        for d in list_devices():
            if (d.get("name") or "") == name:
                return d.get("serial") or ""
    except Exception:
        pass
    try:
        with _ctx():
            dev = Device.query.filter(Device.name == name).first()
            return (dev.serial if dev else "") or ""
    except Exception:
        return ""


# ================== 在线状态 ==================
def list_online():
    """本机 adb 中 state=device 的设备 + 池内 USB 设备（经 220 远程 adb server）。

    池内无 USB 设备（serial 无冒号）时不查远程 server，零额外开销。
    """
    out = _adb("devices") or ""
    serials = []
    for line in out.splitlines()[1:]:
        parts = line.split()
        if len(parts) >= 2 and parts[0] and parts[1] == "device":
            serials.append(parts[0])
    # 池内有 USB 设备时，补查远程 adb server（220，host 网络 5037 全接口监听）
    try:
        usb = [s for s in list_configured() if ":" not in s]
    except Exception:
        usb = []
    if usb:
        try:
            out = _adb_remote(USB_ADB_HOST, USB_ADB_PORT, "devices") or ""
            for line in out.splitlines()[1:]:
                parts = line.split()
                if len(parts) >= 2 and parts[0] and parts[1] == "device":
                    if parts[0] not in serials:
                        serials.append(parts[0])
        except Exception:
            pass
    return serials


def is_online(serial):
    """设备当前是否在线（本机 adb 可达且状态正常）。"""
    return serial in list_online()


def list_ready():
    """可调度设备：清单 ∩ 在线（取代 STF list_free_devices）。"""
    online = set(list_online())
    return [s for s in list_configured() if s in online]


# ================== 设备指纹（识别"同一台物理设备"） ==================
def read_fingerprint(serial, timeout=6):
    """读取设备指纹（ro.serialno）——设备换 IP 后据此认领回原记录。

    只对网络设备（serial 含 ":"）读取：USB 设备的 serial 本身就是稳定序列号，
    不存在"换地址"问题。采集失败返回空串（静默，不影响主流程）。
    """
    serial = (serial or "").strip()
    if not serial or ":" not in serial:
        return ""
    # 部分机型 ro.serialno 为空，退回 ro.boot.serialno
    for prop in ("ro.serialno", "ro.boot.serialno"):
        try:
            r = subprocess.run([ADB_PATH, "-s", serial, "shell", "getprop", prop],
                               capture_output=True, timeout=timeout)
            fp = (r.stdout or b"").decode("utf-8", errors="replace").strip()
        except Exception:
            continue
        if fp and fp.isprintable():
            return fp
    return ""


def find_by_fingerprint(fingerprint, exclude_serial=""):
    """按指纹查池中设备（返回 dict 或 None），可排除指定 serial。"""
    fingerprint = (fingerprint or "").strip()
    if not fingerprint:
        return None
    with _ctx():
        row = Device.query.filter(Device.fingerprint == fingerprint).first()
        if not row or row.serial == exclude_serial:
            return None
        return row.to_dict()


def name_taken(name, exclude_serial=""):
    """名称是否已被其它设备占用（唯一约束的应用层检查，给友好提示用）。"""
    name = (name or "").strip()
    if not name:
        return False
    with _ctx():
        q = Device.query.filter(Device.name == name)
        if exclude_serial:
            q = q.filter(Device.serial != exclude_serial)
        return q.first() is not None


def _move_device_row(old, new_serial, fingerprint=""):
    """把 old 记录迁到 new_serial（调用方持 app context）。

    只动 device 表本身；**分组/任务引用的同步交给上层**（见 set_move_hook）：
    分组与任务在 TaskManager 里还有一份内存副本，调度用的是内存对象——
    只改库不改内存，不重启不生效。而 device_pool 不能反向依赖 task_manager
    （会造成循环 import），所以用回调把这件事交给装配层。
    """
    old_serial, name = old.serial, (old.name or "")
    # 新地址上若已有记录（重复添加等），以"被认领的旧记录"为准，删掉它
    dup = Device.query.get(new_serial)
    if dup is not None and dup.serial != old.serial:
        db.session.delete(dup)
    old.serial = new_serial
    if fingerprint:
        old.fingerprint = fingerprint
    db.session.commit()
    refresh_names()          # 键变了（serial → 名称的映射也跟着变）
    if _move_hook is not None:
        try:
            _move_hook(old_serial, new_serial)
        except Exception as e:
            _log.warning(f"设备迁址后同步分组/任务引用失败: {e}")
    return old_serial, name


def claim_device(new_serial, fingerprint):
    """认领：同一台物理设备换了地址，把池中旧记录迁到新 serial 并同步所有引用。

    做三件事（一个事务内）：
      1. 把旧记录的 serial 改成新地址（名称/型号/备注/启用状态/添加时间全部保留）
      2. 同步 device_group.serials 里的旧 serial → 新 serial（否则分组吊着死 IP）
      3. 同步 task_job.target.serial（指定设备模式的任务目标）

    返回 (old_serial, name)；没有匹配到旧记录时返回 (None, "")。
    """
    fingerprint = (fingerprint or "").strip()
    new_serial = (new_serial or "").strip()
    if not fingerprint or not new_serial:
        return None, ""
    with _ctx():
        old = Device.query.filter(Device.fingerprint == fingerprint,
                                  Device.serial != new_serial).first()
        if not old:
            return None, ""
        old_serial, name = _move_device_row(old, new_serial)
        _log.info(f"设备认领: 指纹 {fingerprint} 的『{name or old_serial}』"
                  f"由 {old_serial} 迁到 {new_serial}")
        return old_serial, name


def relocate_device(old_serial, new_serial, fingerprint=""):
    """人工认领：把池中 old_serial 的记录改到 new_serial（换地址的手工兜底）。

    用在"设备已经断联、读不到指纹"的场景：设备换了 IP 后旧地址连不上，
    指纹也没采过，自动认领无从匹配——此时由人工指认"这条就是那台，
    现在在 X"，本函数负责迁移并同步分组/任务引用。

    返回 (ok, name, msg)。
    """
    old_serial = (old_serial or "").strip()
    new_serial = (new_serial or "").strip()
    if not old_serial or not new_serial:
        return False, "", "缺少参数"
    if old_serial == new_serial:
        return False, "", "新旧地址相同，无需迁移"
    with _ctx():
        old = Device.query.get(old_serial)
        if not old:
            return False, "", f"设备池中没有 {old_serial}"
        if Device.query.get(new_serial) is not None:
            return False, "", f"{new_serial} 已在设备池中，请先处理那条记录"
        _, name = _move_device_row(old, new_serial, fingerprint)
        _log.info(f"设备人工认领: 『{name or old_serial}』{old_serial} → {new_serial}")
        return True, name, f"『{name or old_serial}』已迁到 {new_serial}，分组/任务引用已同步"


# ================== 管理（CRUD） ==================
def add_device(serial, name="", note="", enabled=True, fingerprint=""):
    """添加/更新设备（upsert）。返回 True 新增 / False 已存在并更新。"""
    serial = (serial or "").strip()
    if not serial:
        return False
    with _ctx():
        d = Device.query.get(serial)
        if d:
            d.name, d.note, d.enabled = name or "", note or "", enabled
            if fingerprint:
                d.fingerprint = fingerprint
            db.session.commit()
            refresh_names()          # 名字可能改了：通知里要立刻用新名字
            return False
        db.session.add(Device(serial=serial, name=name or "", note=note or "",
                              enabled=enabled, fingerprint=fingerprint or "",
                              created_at=time.strftime("%Y-%m-%d %H:%M")))
        db.session.commit()
        refresh_names()
        _log.info(f"设备池新增设备: {serial}")
        return True


def rename_device(serial, name):
    """重命名设备（唯一性由调用方先校验）。返回是否成功。"""
    name = (name or "").strip()
    with _ctx():
        d = Device.query.get(serial)
        if not d or not name or (d.name or "") == name:
            return False
        old = d.name or ""
        d.name = name
        db.session.commit()
        refresh_names()
        _log.info(f"设备池重命名: {old or serial} → {name}（{serial}）")
        return True


def set_fingerprint(serial, fingerprint):
    """补写设备指纹（认领需要；采集是后来才做的）。"""
    fingerprint = (fingerprint or "").strip()
    if not fingerprint:
        return False
    with _ctx():
        d = Device.query.get(serial)
        if not d or d.fingerprint == fingerprint:
            return False
        d.fingerprint = fingerprint
        db.session.commit()
        _log.info(f"设备池采集指纹: {serial} -> {fingerprint}")
        return True


def remove_device(serial):
    """删除设备。返回是否删除成功。"""
    with _ctx():
        d = Device.query.get(serial)
        if not d:
            return False
        db.session.delete(d)
        db.session.commit()
        refresh_names()
        _log.info(f"设备池删除设备: {serial}")
        return True


def set_enabled(serial, enabled):
    """启用/停用设备（停用后不参与调度）。返回是否成功。"""
    with _ctx():
        d = Device.query.get(serial)
        if not d:
            return False
        d.enabled = enabled
        db.session.commit()
        return True


# ================== 型号采集 ==================
def refresh_model(serial, timeout=8):
    """采集单台设备型号（getprop ro.product.model）并写回 DB。

    IP:5555 设备走本机 adb；USB 设备先查本机 adb，查不到再查 220 远程 adb server。
    采集失败返回空串（静默，不影响主流程）。返回采集到的型号。
    """
    if not serial:
        return ""
    model = ""
    if ":" in serial:
        try:
            r = subprocess.run([ADB_PATH, "-s", serial, "shell",
                                "getprop", "ro.product.model"],
                               capture_output=True, timeout=timeout)
            model = (r.stdout or b"").decode("utf-8", errors="replace").strip()
        except Exception:
            model = ""
    else:
        # USB：先本机，再远程
        try:
            r = subprocess.run([ADB_PATH, "-s", serial, "shell",
                                "getprop", "ro.product.model"],
                               capture_output=True, timeout=timeout)
            model = (r.stdout or b"").decode("utf-8", errors="replace").strip()
        except Exception:
            model = ""
        if not model:
            model = _adb_remote(USB_ADB_HOST, USB_ADB_PORT, "-s", serial,
                                "shell", "getprop", "ro.product.model")
    model = model.strip() or ""
    if not model:
        return ""
    with _ctx():
        d = Device.query.get(serial)
        if d and d.model != model:
            d.model = model
            db.session.commit()
            _log.info(f"设备池采集型号: {serial} -> {model}")
    return model


def refresh_fingerprint(serial, timeout=6):
    """读取并写回设备指纹，返回指纹（失败返回空串）。"""
    fp = read_fingerprint(serial, timeout=timeout)
    if fp:
        set_fingerprint(serial, fp)
    return fp


def refresh_info(serial):
    """采集一台设备的型号 + 指纹（指纹缺失时才读，避免每次启动都白跑一次 adb）。

    返回 (model, fingerprint)。任一失败都静默——设备池展示用，不影响主流程。
    """
    model = refresh_model(serial)
    fp = ""
    with _ctx():
        row = Device.query.get(serial)
        fp = (row.fingerprint or "") if row else ""
    if not fp:
        fp = refresh_fingerprint(serial)
    return model, fp


def refresh_all_models():
    """批量采集池内在线设备的型号（顺带补齐缺失的设备指纹），并发 10，后台线程调用。

    返回成功数。指纹补齐很重要：老库里的设备没有指纹，补上之后换 IP 才能被认领。
    """
    serials = list_online()
    if not serials:
        return 0
    from concurrent.futures import ThreadPoolExecutor, as_completed
    ok = 0
    with ThreadPoolExecutor(max_workers=min(10, len(serials))) as pool:
        futures = {pool.submit(refresh_info, s): s for s in serials}
        for fut in as_completed(futures):
            try:
                model, _fp = fut.result()
                if model:
                    ok += 1
            except Exception:
                pass
    _log.info(f"设备池批量采集型号完成: {ok}/{len(serials)} 台")
    return ok