#!/usr/bin/env python3
"""
uiautomator2 自动重连 wrapper
用法: from u2_connect import get_device
"""

import uiautomator2 as u2
import subprocess
import time

SERIAL = 'localhost:15556'

def get_device(serial=SERIAL, retry=2):
    """获取设备连接，自动重启 ATX agent"""
    for attempt in range(retry + 1):
        try:
            d = u2.connect(serial)
            # 测试连接是否正常
            d.info
            return d
        except Exception as e:
            if attempt < retry:
                print(f"[u2] 连接失败，重启 ATX agent... ({attempt + 1}/{retry})")
                # 重启 agent
                subprocess.run(
                    ['python', '-m', 'uiautomator2', 'init', '--serial', serial],
                    capture_output=True, timeout=60
                )
                time.sleep(2)
            else:
                raise Exception(f"连接失败: {e}")

def dump_ui(d=None):
    """dump UI 树"""
    if d is None:
        d = get_device()
    return d.dump_hierarchy()

def find_and_click(text, d=None):
    """按文字查找并点击"""
    if d is None:
        d = get_device()
    el = d(textContains=text)
    if el.exists:
        el.click()
        return True
    return False

if __name__ == '__main__':
    d = get_device()
    print(f"连接成功: {d.info['currentPackageName']}")