iOSAI/Module/DeviceInfo.py

import json
import os
import socket
import threading
import time
import subprocess
from typing import Dict, Optional

import tidevice
# import wda  # 目前先不用 wda 获取屏幕，避免触发 tidevice 的那套 WDA 启动
from tidevice import Usbmux, ConnectionType

from Entity.DeviceModel import DeviceModel
from Entity.Variables import WdaAppBundleId, wdaScreenPort, wdaFunctionPort
from Module.FlaskSubprocessManager import FlaskSubprocessManager
from Module.IOSActivator import IOSActivator
from Utils.LogManager import LogManager


class DeviceInfo:
    _instance = None
    _instance_lock = threading.Lock()

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            with cls._instance_lock:
                if not cls._instance:
                    cls._instance = super().__new__(cls)
        return cls._instance

    REMOVE_GRACE_SEC = float(os.getenv("REMOVE_GRACE_SEC", "6.0"))

    def __init__(self) -> None:
        if getattr(self, "_initialized", False):
            return

        self._lock = threading.RLock()
        self._models: Dict[str, DeviceModel] = {}
        self._manager = FlaskSubprocessManager.get_instance()
        self.screenPort = 9110

        # 设备心跳时间
        self._last_seen: Dict[str, float] = {}

        # iproxy 子进程：udid -> Popen
        self._iproxy_process: Dict[str, subprocess.Popen] = {}

        # Windows 下隐藏子进程窗口（给 iproxy 用）
        self._creationflags = 0
        self._startupinfo = None
        if os.name == "nt":
            try:
                self._creationflags = subprocess.CREATE_NO_WINDOW  # type: ignore[attr-defined]
            except Exception:
                self._creationflags = 0

            si = subprocess.STARTUPINFO()
            si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
            si.wShowWindow = 0  # SW_HIDE
            self._startupinfo = si

        LogManager.info("DeviceInfo 初始化完成", udid="system")
        print("[Init] DeviceInfo 初始化完成")
        self._initialized = True

    # ---------------------------
    # 主循环
    # ---------------------------
    def listen(self):
        LogManager.method_info("进入主循环", "listen", udid="system")
        print("[Listen] 开始监听设备上下线...")

        while True:
            try:
                usb = Usbmux().device_list()
                online = {d.udid for d in usb if d.conn_type == ConnectionType.USB}
            except Exception as e:
                LogManager.warning(f"[device_list] 异常：{e}", udid="system")
                time.sleep(1)
                continue

            with self._lock:
                known = set(self._models.keys())

            # ---------- 1. 新设备 ----------
            for udid in online:
                self._last_seen[udid] = time.time()
                if udid not in known:
                    try:
                        self._add_device(udid)
                    except Exception as e:
                        # 关键：单设备异常不能干掉整个循环
                        LogManager.warning(f"[Add] 处理设备 {udid} 异常: {e}", udid=udid)
                        print(f"[Add] 处理设备 {udid} 异常: {e}")

            # ---------- 2. 可能离线设备 ----------
            now = time.time()
            for udid in list(known):
                if udid not in online:
                    last = self._last_seen.get(udid, 0)
                    if now - last > self.REMOVE_GRACE_SEC:
                        try:
                            self._remove_device(udid)
                        except Exception as e:
                            LogManager.method_error(f"移除失败：{e}", "listen", udid=udid)
                            print(f"[Remove] 移除失败 {udid}: {e}")

            time.sleep(1)

    # ---------------------------
    # 添加设备
    # ---------------------------
    def _add_device(self, udid: str):
        with self._lock:
            if udid in self._models:
                print(f"[Add] 已存在，跳过: {udid}")
                return
            print(f"[Add] 新增设备 {udid}")

        # 判断 iOS 版本
        try:
            t = tidevice.Device(udid)
            version_major = float(t.product_version.split(".")[0])
        except Exception as e:
            print(f"[Add] 获取系统版本失败 {udid}: {e}")
            version_major = 0

        # 启动 WDA
        if version_major >= 17.0:
            threading.Thread(
                target=IOSActivator().activate_ios17,
                args=(udid,),
                daemon=True,
            ).start()
        else:
            try:
                tidevice.Device(udid).app_start(WdaAppBundleId)
            except Exception as e:
                print(f"[Add] 使用 tidevice 启动 WDA 失败 {udid}: {e}")

        # 分配投屏端口 & 写入模型
        with self._lock:
            self.screenPort += 1
            screen_port = self.screenPort

            model = DeviceModel(
                deviceId=udid,
                screenPort=screen_port,
                width=0,
                height=0,
                scale=0,
                type=1
            )
            self._models[udid] = model

        print(f"[Add] 新设备完成 {udid}, screenPort={screen_port}")
        self._manager_send()

        # 启动 iproxy（投屏转发）
        try:
            self._start_iproxy(udid, screen_port)
        except Exception as e:
            print(f"[iproxy] 启动失败 {udid}: {e}")

        # 异步等待 WDA + 更新屏幕尺寸（后面你要真用尺寸，再把 _screen_info 换成 wda 版本即可）
        threading.Thread(
            target=self._async_wait_wda_and_update_screen,
            args=(udid,),
            daemon=True
        ).start()

    # ---------------------------
    # 启动 iproxy（投屏）
    # ---------------------------
    def _start_iproxy(self, udid: str, local_port: int):
        iproxy_path = self._find_iproxy()

        # 已有进程并且还在跑，就不重复起
        p = self._iproxy_process.get(udid)
        if p is not None and p.poll() is None:
            print(f"[iproxy] 已存在运行中的进程，跳过 {udid}")
            return

        args = [
            iproxy_path,
            "-u", udid,
            str(local_port),    # 本地端口（投屏）
            "9567"              # 手机端口（go-ios screencast）
        ]

        print(f"[iproxy] 启动进程: {args}")

        # 不用 PIPE，防止没人读导致缓冲爆掉；窗口用前面配置隐藏
        proc = subprocess.Popen(
            args,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            creationflags=self._creationflags,
            startupinfo=self._startupinfo,
        )

        self._iproxy_process[udid] = proc

    # ---------------------------
    # 停止 iproxy
    # ---------------------------
    def _stop_iproxy(self, udid: str):
        p = self._iproxy_process.get(udid)
        if not p:
            return
        try:
            p.terminate()
            try:
                p.wait(timeout=2)
            except Exception:
                p.kill()
        except Exception:
            pass
        self._iproxy_process.pop(udid, None)
        print(f"[iproxy] 已停止 {udid}")

    # ---------------------------
    # 异步等待 WDA 然后更新屏幕大小
    # ---------------------------
    def _async_wait_wda_and_update_screen(self, udid: str):
        print(f"[WDA] 等待 WDA 就绪 {udid}")

        try:
            # 最长等待 20 秒（你后面要真用屏幕尺寸，再把 _screen_info 换回 wda 的实现）
            for _ in range(20):
                w, h, s = self._screen_info(udid)
                if w > 0:
                    print(f"[WDA] 屏幕信息成功 {udid} {w}x{h} scale={s}")
                    with self._lock:
                        m = self._models.get(udid)
                        if m:
                            m.width = w
                            m.height = h
                            m.scale = s
                    self._manager_send()
                    return
                time.sleep(1)
        except Exception as e:
            print(f"[WDA] 获取屏幕信息异常 {udid}: {e}")

        print(f"[WDA] 屏幕信息获取失败（超时） {udid}")

    # ---------------------------
    # 移除设备
    # ---------------------------
    def _remove_device(self, udid: str):
        print(f"[Remove] 移除设备 {udid}")

        # 先停 iproxy（只管自己的，不影响其他设备）
        self._stop_iproxy(udid)

        with self._lock:
            self._models.pop(udid, None)
            self._last_seen.pop(udid, None)

        self._manager_send()

    # ---------------------------
    # WDA 屏幕查询（当前保持空实现）
    # ---------------------------
    def _screen_info(self, udid: str):
        # 现在先不通过 wda 取屏幕，避免触发 tidevice 那套 WDA 启动逻辑
        # 你如果确认 go-ios 跑起来后用 facebook-wda 取尺寸是安全的，可以改回下面这种：
        #
        # try:
        #     c = wda.USBClient(udid, wdaFunctionPort)
        #     size = c.window_size()
        #     return int(size.width), int(size.height), float(c.scale)
        # except Exception as e:
        #     print(f"[Screen] 获取屏幕信息异常: {e} {udid}")
        #     return 0, 0, 0.0
        return 0, 0, 0.0

    # ---------------------------
    # 找 iproxy
    # ---------------------------
    def _find_iproxy(self) -> str:
        base = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        name = "iproxy.exe" if os.name == "nt" else "iproxy"
        path = os.path.join(base, "resources", "iproxy", name)
        return path

    # ---------------------------
    # 数据同步到 Flask
    # ---------------------------
    def _manager_send(self):
        try:
            self._send_snapshot_to_flask()
        except Exception:
            try:
                self._manager.start()
            except Exception:
                pass
            try:
                self._send_snapshot_to_flask()
            except Exception:
                pass

    def _send_snapshot_to_flask(self):
        with self._lock:
            devices = [m.toDict() for m in self._models.values()]

        payload = json.dumps({"devices": devices}, ensure_ascii=False)

        port = int(os.getenv("FLASK_COMM_PORT", "34566"))
        with socket.create_connection(("127.0.0.1", port), timeout=1.5) as s:
            s.sendall(payload.encode() + b"\n")

        print(f"[SNAPSHOT] 已发送 {len(devices)} 台设备")