iOSAI/Utils/AiUtils.py

import os
import re
import cv2
import numpy as np
import wda
from Utils.LogManager import LogManager
import xml.etree.ElementTree as ET
from lxml import etree
from wda import Client

# 工具类
class AiUtils(object):

    # 在屏幕中找到对应的图片
    @classmethod
    def findImageInScreen(cls, target, udid):
        try:
            # 加载原始图像和模板图像
            image_path = AiUtils.imagePathWithName(udid,"bgv")  # 替换为你的图像路径
            template_path = AiUtils.imagePathWithName("", target)  # 替换为你的模板路径

            # 读取图像和模板，确保它们都是单通道灰度图
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)

            if image is None:
                LogManager.error("加载背景图失败")
                return -1, -1

            if template is None:
                LogManager.error("加载模板图失败")
                return -1, -1

            # 获取模板的宽度和高度
            w, h = template.shape[::-1]

            # 使用模板匹配方法
            res = cv2.matchTemplate(image, template, cv2.TM_CCOEFF_NORMED)
            threshold = 0.7  # 匹配度阈值，可以根据需要调整
            loc = np.where(res >= threshold)

            # 检查是否有匹配结果
            if loc[0].size > 0:
                # 取第一个匹配位置
                pt = zip(*loc[::-1]).__next__()  # 获取第一个匹配点的坐标
                center_x = int(pt[0] + w // 2)
                center_y = int(pt[1] + h // 2)
                # print(f"第一个匹配到的小心心中心坐标: ({center_x}, {center_y})")
                return center_x, center_y
            else:
                return -1, -1
        except Exception as e:
            LogManager.error(f"加载素材失败：{e}", udid)
            print(e)
            return -1, -1

    # 使用正则查找字符串中的数字
    @classmethod
    def findNumber(cls, str):
        # 使用正则表达式匹配数字
        match = re.search(r'\d+', str)
        if match:
            return int(match.group())  # 将匹配到的数字转换为整数
        return None  # 如果没有找到数字，返回 None


    # 选择截图
    @classmethod
    def screenshot(cls):
        client = wda.USBClient("eca000fcb6f55d7ed9b4c524055214c26a7de7aa")
        session = client.session()
        image = session.screenshot()
        image_path = "screenshot.png"
        image.save(image_path)
        image = cv2.imread(image_path)

        # 如果图像过大，缩小显示
        scale_percent = 50  # 缩小比例
        width = int(image.shape[1] * scale_percent / 100)
        height = int(image.shape[0] * scale_percent / 100)
        dim = (width, height)
        resized_image = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)

        # 创建一个窗口并显示缩小后的图像
        cv2.namedWindow("Image")
        cv2.imshow("Image", resized_image)
        print("请在图像上选择爱心图标区域，然后按Enter键确认。")

        # 使用selectROI函数手动选择区域
        roi = cv2.selectROI("Image", resized_image, showCrosshair=True, fromCenter=False)

        # 将ROI坐标按原始图像尺寸放大
        x, y, w, h = roi
        x = int(x * image.shape[1] / resized_image.shape[1])
        y = int(y * image.shape[0] / resized_image.shape[0])
        w = int(w * image.shape[1] / resized_image.shape[1])
        h = int(h * image.shape[0] / resized_image.shape[0])

        # 根据选择的ROI提取爱心图标
        if w > 0 and h > 0:  # 确保选择的区域有宽度和高度
            heart_icon = image[y:y + h, x:x + w]

            # 转换为HSV颜色空间
            hsv = cv2.cvtColor(heart_icon, cv2.COLOR_BGR2HSV)

            # 定义红色的HSV范围
            lower_red1 = np.array([0, 120, 70])
            upper_red1 = np.array([10, 255, 255])
            lower_red2 = np.array([170, 120, 70])
            upper_red2 = np.array([180, 255, 255])

            # 创建掩模
            mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
            mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
            mask = mask1 + mask2

            # 反转掩模，因为我们想要的是爱心图标，而不是背景
            mask_inv = cv2.bitwise_not(mask)

            # 应用掩模
            heart_icon = cv2.bitwise_and(heart_icon, heart_icon, mask=mask_inv)

            # 创建一个全透明的背景
            height, width, channels = heart_icon.shape
            roi = np.zeros((height, width, channels), dtype=np.uint8)

            # 将爱心图标粘贴到透明背景上
            for c in range(channels):
                roi[:, :, c] = np.where(mask_inv == 255, heart_icon[:, :, c], roi[:, :, c])

            # 图片名称
            imgName = "temp.png"

            # 保存结果
            cv2.imwrite(imgName, roi)

            # 显示结果
            cv2.imshow("Heart Icon with Transparent Background", roi)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        else:
            print("未选择有效区域。")


    # 根据名称获取图片完整地址
    @classmethod
    def imagePathWithName(cls, udid, name):
        current_file_path = os.path.abspath(__file__)
        # 获取当前文件所在的目录（即script目录）
        current_dir = os.path.dirname(current_file_path)
        # 由于script目录位于项目根目录下一级，因此需要向上一级目录移动两次
        project_root = os.path.abspath(os.path.join(current_dir, '..'))
        # 构建资源文件的完整路径，向上两级目录，然后进入 resources 目录
        resource_path = os.path.abspath(os.path.join(project_root, "resources", udid, name + ".png")).replace('/', '\\')
        return resource_path

    # 获取根目录
    @classmethod
    def getRootDir(cls):
        current_file = os.path.abspath(__file__)
        # 获取当前文件所在的目录
        current_dir = os.path.dirname(current_file)
        # 获取项目根目录（假设根目录是当前文件的父目录的父目录）
        project_root = os.path.dirname(current_dir)
        # 返回根目录
        return project_root


    # 创建一个以udid命名的目录
    @classmethod
    def makeUdidDir(cls, udid):
        # 获取项目根目录
        home = cls.getRootDir()
        # 拼接 resources 目录的路径
        resources_dir = os.path.join(home, "resources")
        # 拼接 udid 目录的路径
        udid_dir = os.path.join(resources_dir, udid)
        # 检查 udid 目录是否存在，如果不存在则创建
        if not os.path.exists(udid_dir):
            try:
                os.makedirs(udid_dir)
                LogManager.info(f"目录 {udid_dir} 创建成功", udid)
                print(f"目录 {udid_dir} 创建成功")
            except Exception as e:
                print(f"创建目录时出错: {e}")
                LogManager.error(f"创建目录时出错: {e}", udid)
        else:
            LogManager.info(f"目录 {udid_dir} 已存在，跳过创建", udid)
            print(f"目录 {udid_dir} 已存在，跳过创建")

    # 查找首页按钮
    # uuid 设备id
    # click 是否点击该按钮
    @classmethod
    def findHomeButton(cls, udid="eca000fcb6f55d7ed9b4c524055214c26a7de7aa"):
        client = wda.USBClient(udid)
        session = client.session()
        session.appium_settings({"snapshotMaxDepth": 10})
        homeButton = session.xpath("//*[@label='首页']")
        try:
            if homeButton.label == "首页":
                print("找到了")
                return homeButton
            else:
                print("没找到")
                return None
        except Exception as e:
            print(e)
            return None


    # 查找关闭按钮
    @classmethod
    def findLiveCloseButton(cls,udid="eca000fcb6f55d7ed9b4c524055214c26a7de7aa"):
        client = wda.USBClient(udid)
        session = client.session()
        session.appium_settings({"snapshotMaxDepth": 10})
        r = session.xpath("//XCUIElementTypeButton[@name='关闭屏幕']")
        try:
            if r.label == "关闭屏幕":
                return r
            else:
                return None
        except Exception as e:
            print(e)
            return None

    # 获取直播间窗口数量
    @classmethod
    def count_add_by_xml(cls, session):
        xml = session.source()
        root = ET.fromstring(xml)
        return sum(
            1 for e in root.iter()
            if e.get('type') in ('XCUIElementTypeButton', 'XCUIElementTypeImage')
            and (e.get('name') == '添加' or e.get('label') == '添加')
            and (e.get('visible') in (None, 'true'))
        )

    # 获取关注按钮
    @classmethod
    def getFollowButton(cls, session: Client):
        followButton = session.xpath("//Window[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[2]/Other[2]/Other[1]/Other[1]/Other[3]/Other[1]/Other[1]/Button[1]")
        if followButton.exists:
            print("找到了")
            return followButton
        else:
            print("没找到")
            return None

    # 查找发消息按钮
    @classmethod
    def getSendMesageButton(cls, session: Client):
        msgButton = session.xpath("//Window[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[1]/Other[2]/Other[2]/Other[1]/Other[1]/Other[3]/Other[1]/Other[1]")
        if msgButton.exists:
            print("找到了")
            return msgButton
        else:
            print("没找到")
            return None

    # 获取当前屏幕上的节点
    @classmethod
    def getCurrentScreenSource(cls):
        client = wda.USBClient("eca000fcb6f55d7ed9b4c524055214c26a7de7aa")
        print(client.source())


    # 查找app主页上的收件箱按钮
    @classmethod
    def getMsgBoxButton(cls, session: Client):
        box = session.xpath("//XCUIElementTypeButton[name='a11y_vo_inbox']")
        if box.exists:
            return box
        else:
            return None

    # 获取收件箱中的未读消息数
    @classmethod
    def getUnReadMsgCount(cls, session: Client):
        btn = cls.getMsgBoxButton(session)
        return cls.findNumber(btn.label)

    # 获取聊天页面的聊天信息
    @classmethod
    def extract_messages_from_xml(cls, xml: str):
        """
        输入 WDA 的页面 XML，输出按时间顺序的消息列表：
        每项形如：
          {'type': 'time', 'text': '昨天 下午8:48'}
          {'type': 'msg', 'dir': 'in'|'out', 'text': 'hello'}
        """
        root = etree.fromstring(xml.encode("utf-8"))
        items = []

        # 屏幕宽度（用于右对齐判断）
        app = root.xpath('/XCUIElementTypeApplication')
        screen_w = cls.parse_float(app[0], 'width', 414.0) if app else 414.0

        # 1) 时间分隔
        for t in root.xpath('//XCUIElementTypeStaticText[contains(@traits, "Header")]'):
            txt = t.get('label') or t.get('name') or t.get('value') or ''
            y = cls.parse_float(t, 'y')
            if txt.strip():
                items.append({'type': 'time', 'text': txt.strip(), 'y': y})

        # 2) 消息气泡
        msg_nodes = root.xpath(
            '//XCUIElementTypeTable//XCUIElementTypeCell'
            '//XCUIElementTypeOther[@name or @label]'
        )

        EXCLUDES = {'Heart', 'Lol', 'ThumbsUp', '分享发布内容', '视频贴纸标签页', '双击发送表情'}

        for o in msg_nodes:
            text = (o.get('label') or o.get('name') or '').strip()
            if not text or text in EXCLUDES:
                continue

            # 拿到所在的 Cell（用来找头像按钮）
            cell = o.getparent()
            while cell is not None and cell.get('type') != 'XCUIElementTypeCell':
                cell = cell.getparent()

            x = cls.parse_float(o, 'x')
            y = cls.parse_float(o, 'y')
            w = cls.parse_float(o, 'width')
            right_edge = x + w

            direction = None

            # 2.1 依据同 Cell 内“图片头像”的位置判定（优先，最稳）
            if cell is not None:
                avatar_btns = cell.xpath('.//XCUIElementTypeButton[@name="图片头像" or @label="图片头像"]')
                if avatar_btns:
                    ax = cls.parse_float(avatar_btns[0], 'x')
                    # 头像在左侧 → 对方；头像在右侧 → 自己
                    if ax < screen_w / 2:
                        direction = 'in'
                    else:
                        direction = 'out'

            # 2.2 退化规则：看是否右对齐
            if direction is None:
                # 离右边 <= 20px 视为右对齐（自己发的）
                if right_edge > screen_w - 20:
                    direction = 'out'
                else:
                    direction = 'in'

            items.append({'type': 'msg', 'dir': direction, 'text': text, 'y': y})

        # 3) 按 y 排序并清理
        items.sort(key=lambda i: i['y'])
        for it in items:
            it.pop('y', None)
        return items

    @classmethod
    def parse_float(cls, el, attr, default=0.0):
        try:
            return float(el.get(attr, default))
        except Exception:
            return default

    # 从导航栏读取主播名称。找不到时返回空字符串
    @classmethod
    def get_navbar_anchor_name(cls, session, timeout: float = 2.0) -> str:
        """只从导航栏读取主播名称。找不到时返回空字符串。"""
        # 可选：限制快照深度，提升解析速度/稳定性
        try:
            session.appium_settings({"snapshotMaxDepth": 22})
        except Exception:
            pass

        def _text_of(el) -> str:
            info = getattr(el, "info", {}) or {}
            return (info.get("label") or info.get("name") or info.get("value") or "").strip()

        def _clean_tail(s: str) -> str:
            return re.sub(r"[，、,。.\s]+$", "", s).strip()

        # 导航栏容器：从“返回”按钮向上找最近祖先，且该祖先内包含“更多/举报”（多语言兜底）
        NAV_CONTAINER = (
            "//XCUIElementTypeButton"
            "[@name='返回' or @label='返回' or @name='Back' or @label='Back' or @name='戻る' or @label='戻る']"
            "/ancestor::XCUIElementTypeOther"
            "[ .//XCUIElementTypeButton"
            "   [@name='更多' or @label='更多' or @name='More' or @label='More' or "
            "    @name='その他' or @label='その他' or @name='詳細' or @label='詳細' or "
            "    @name='举报' or @label='举报' or @name='Report' or @label='Report' or "
            "    @name='報告' or @label='報告']"
            "][1]"
        )

        # ① 优先：可访问的 Other（自身有文本，且不含子 Button），更贴近 TikTok 的实现
        XPATH_TITLE_OTHER = (
                NAV_CONTAINER +
                "//XCUIElementTypeOther[@accessible='true' and count(.//XCUIElementTypeButton)=0 "
                "  and (string-length(@name)>0 or string-length(@label)>0 or string-length(@value)>0)"
                "][1]"
        )
        # ② 退路：第一个 StaticText
        XPATH_TITLE_STATIC = (
                NAV_CONTAINER +
                "//XCUIElementTypeStaticText[string-length(@value)>0 or string-length(@label)>0 or string-length(@name)>0][1]"
        )

        # 尝试 ①
        q = session.xpath(XPATH_TITLE_OTHER)
        if q.wait(timeout):
            t = _clean_tail(_text_of(q.get()))
            if t:
                return t

        # 尝试 ②
        q2 = session.xpath(XPATH_TITLE_STATIC)
        if q2.wait(1.0):
            t = _clean_tail(_text_of(q2.get()))
            if t:
                return t

        return ""

# AiUtils.getCurrentScreenSource()