#!/usr/bin/env python3
"""
工作手机SDK - 逐步操作验证器
每个功能：导航到对应UI页面 -> 执行Frida RPC -> 等待UI响应 -> 截图 -> 返回数据
确保每张截图与功能一一对应
"""
import frida
import json
import time
import os
import subprocess
from datetime import datetime

DEVICE_IP = "192.168.0.12"
FRIDA_PORT = 27042
WECHAT_PID = 7239
DEVICE_SERIAL = "xgfe65eimrrofyws"
WECHAT_PKG = "com.tencent.mm"

SCRIPT_PATH = os.path.expanduser("~/Documents/GitHub/workphone-sdk/sdk/agent/hook/wechat_hook_v2.js")
OUTPUT_DIR = os.path.expanduser("~/Documents/GitHub/workphone-sdk/verification_steps")
os.makedirs(OUTPUT_DIR, exist_ok=True)

LOG_FILE = os.path.join(OUTPUT_DIR, "verification_log.md")


class PhoneController:
    """通过ADB控制手机UI导航"""

    def __init__(self, serial):
        self.serial = serial

    def adb(self, cmd):
        full = f"adb -s {self.serial} {cmd}"
        r = subprocess.run(full, shell=True, capture_output=True, timeout=10)
        return r.stdout.decode("utf-8", errors="ignore").strip()

    def shell(self, cmd):
        return self.adb(f'shell "{cmd}"')

    def screenshot(self, name):
        """截图并保存到本地，返回文件路径"""
        ts = int(time.time())
        fname = f"{name}_{ts}.png"
        fpath = os.path.join(OUTPUT_DIR, fname)
        self.shell("screencap -p /sdcard/verify_sc.png")
        time.sleep(0.5)
        self.adb(f"pull /sdcard/verify_sc.png {fpath}")
        if os.path.exists(fpath) and os.path.getsize(fpath) > 1000:
            return fpath
        return ""

    def tap(self, x, y):
        """点击屏幕坐标"""
        self.shell(f"input tap {x} {y}")
        time.sleep(0.8)

    def swipe(self, x1, y1, x2, y2, duration=300):
        """滑动"""
        self.shell(f"input swipe {x1} {y1} {x2} {y2} {duration}")
        time.sleep(0.5)

    def back(self):
        """返回键"""
        self.shell("input keyevent 4")
        time.sleep(0.5)

    def home(self):
        """Home键"""
        self.shell("input keyevent 3")
        time.sleep(0.5)

    def open_wechat(self):
        """打开微信主页"""
        self.shell(f"am start -n {WECHAT_PKG}/.ui.LauncherUI")
        time.sleep(2)

    def open_wechat_contacts(self):
        """打开微信通讯录Tab"""
        self.open_wechat()
        time.sleep(1)
        # 点击底部通讯录Tab (通常在底部第2个位置)
        # 小米手机分辨率通常是1080x2400
        self.tap(360, 2340)  # 通讯录tab
        time.sleep(1)

    def open_wechat_discover(self):
        """打开微信发现Tab"""
        self.open_wechat()
        time.sleep(1)
        self.tap(720, 2340)  # 发现tab
        time.sleep(1)

    def open_wechat_me(self):
        """打开微信我Tab"""
        self.open_wechat()
        time.sleep(1)
        self.tap(980, 2340)  # 我tab
        time.sleep(1)

    def open_chat(self, name="filehelper"):
        """打开指定聊天窗口"""
        self.open_wechat()
        time.sleep(1)
        # 点击搜索
        self.tap(950, 140)
        time.sleep(1)
        # 输入搜索内容
        self.shell(f"input text {name}")
        time.sleep(1)
        # 点击第一个结果
        self.tap(540, 400)
        time.sleep(1)

    def get_screen_size(self):
        """获取屏幕分辨率"""
        r = self.shell("wm size")
        # Physical size: 1080x2400
        if "x" in r:
            parts = r.split(":")[-1].strip().split("x")
            return int(parts[0]), int(parts[1])
        return 1080, 2400


class StepVerifier:
    """逐步验证器"""

    def __init__(self):
        self.phone = PhoneController(DEVICE_SERIAL)
        self.exports = None
        self.script = None
        self.session = None
        self.log_lines = []
        self.results = []

    def connect(self):
        """连接Frida"""
        print("[CONNECT] Connecting to Frida...")
        device = frida.get_device_manager().add_remote_device(f"{DEVICE_IP}:{FRIDA_PORT}")
        self.session = device.attach(WECHAT_PID)

        with open(SCRIPT_PATH, "r", encoding="utf-8") as f:
            source = f.read()

        self.script = self.session.create_script(source)
        self.script.on("message", lambda m, d: None)
        self.script.load()
        time.sleep(2)

        self.exports = self.script.exports_sync
        avail = [x for x in dir(self.exports) if not x.startswith("_")]
        print(f"  Loaded {len(avail)} methods")
        pong = self.exports.ping()
        print(f"  ping: {pong}")

        # 获取屏幕尺寸
        w, h = self.phone.get_screen_size()
        print(f"  Screen: {w}x{h}")
        return True

    def log(self, text):
        print(text)
        self.log_lines.append(text)

    def verify_step(self, module, method, description, nav_action, params=None, wait_after=1):
        """
        执行一个验证步骤:
        1. 导航到对应页面
        2. 截图(操作前)
        3. 执行RPC
        4. 等待UI响应
        5. 截图(操作后)
        6. 记录结果
        """
        self.log(f"\n{'='*60}")
        self.log(f"  [{module}] {method}")
        self.log(f"  描述: {description}")
        self.log(f"{'='*60}")

        # Step 1: 导航
        self.log(f"  [NAV] {nav_action.__doc__ or 'navigating...'}")
        nav_action()
        time.sleep(1)

        # Step 2: 操作前截图
        before_img = self.phone.screenshot(f"{module}__{method}__before")
        self.log(f"  [BEFORE] 截图: {os.path.basename(before_img) if before_img else 'FAILED'}")

        # Step 3: 执行RPC
        fn = getattr(self.exports, method, None)
        if fn is None:
            self.log(f"  [ERROR] 方法不存在: {method}")
            self.results.append({
                "module": module, "method": method, "status": "MISSING",
                "description": description, "data": None,
                "before_img": before_img, "after_img": "",
            })
            return

        start = time.time()
        try:
            if params:
                result = fn(params)
            else:
                result = fn()
            latency = int((time.time() - start) * 1000)
        except Exception as e:
            latency = int((time.time() - start) * 1000)
            result = {"success": False, "error": str(e)[:200]}

        self.log(f"  [EXEC] {method}({json.dumps(params, ensure_ascii=False)[:50] if params else ''})")
        self.log(f"  [TIME] {latency}ms")

        # 格式化结果
        if isinstance(result, str):
            self.log(f"  [DATA] {result[:100]}")
            data_preview = result
        elif isinstance(result, dict):
            data_preview = json.dumps(result, ensure_ascii=False, indent=2)[:300]
            self.log(f"  [DATA] {data_preview[:150]}")
        elif isinstance(result, list):
            data_preview = f"[{len(result)} items] " + json.dumps(result[:2], ensure_ascii=False)[:200]
            self.log(f"  [DATA] {data_preview[:150]}")
        else:
            data_preview = str(result)[:100]
            self.log(f"  [DATA] {data_preview}")

        # Step 4: 等待UI响应
        time.sleep(wait_after)

        # Step 5: 操作后截图
        after_img = self.phone.screenshot(f"{module}__{method}__after")
        self.log(f"  [AFTER] 截图: {os.path.basename(after_img) if after_img else 'FAILED'}")

        # 判断状态
        if isinstance(result, dict) and result.get("success") == False:
            status = "EXEC_ERR"
        else:
            status = "PASS"

        self.log(f"  [STATUS] {status}")

        self.results.append({
            "module": module, "method": method, "status": status,
            "description": description, "latency_ms": latency,
            "data": result if isinstance(result, (dict, list, str, bool, int, float)) else str(result),
            "before_img": before_img, "after_img": after_img,
        })

    def run_all(self):
        """运行所有验证步骤"""
        self.log("# 工作手机SDK - 逐步真机验证")
        self.log(f"> 时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        self.log(f"> 设备: {DEVICE_IP}:{FRIDA_PORT}")
        self.log("")

        # ===== 模块1: 系统状态 =====
        self.log("\n## 模块1: 系统状态\n")

        self.verify_step(
            "SYS", "getProcessInfo",
            "获取微信进程信息(PID/架构/模块数)",
            lambda: self.phone.open_wechat(),
        )

        self.verify_step(
            "SYS", "getWechatVersion",
            "获取微信版本号",
            lambda: None,  # 不需要导航，直接在当前页面
        )

        self.verify_step(
            "SYS", "getHookStatus",
            "检查Hook是否已激活",
            lambda: None,
        )

        # ===== 模块2: 联系人 =====
        self.log("\n## 模块2: 联系人\n")

        self.verify_step(
            "H16_联系人", "getContacts",
            "打开通讯录 -> 获取联系人列表",
            lambda: self.phone.open_wechat_contacts(),
            params={"limit": 10},
        )

        self.verify_step(
            "H16_联系人", "getContactInfo",
            "获取文件传输助手的详细信息",
            lambda: None,  # 保持在通讯录页面
            params={"wxid": "filehelper"},
        )

        self.verify_step(
            "H16_联系人", "searchContacts",
            "搜索包含'文件'的联系人",
            lambda: None,
            params={"keyword": "文件", "limit": 5},
        )

        # ===== 模块3: 消息发送 =====
        self.log("\n## 模块3: 消息发送\n")

        msg_content = f"[SDK验证] {datetime.now().strftime('%H:%M:%S')} 测试消息"
        self.verify_step(
            "H17_消息发送", "sendMessage",
            f"打开文件传输助手对话 -> 发送文本消息: '{msg_content}'",
            lambda: self.phone.open_chat("filehelper"),
            params={"to_id": "filehelper", "content": msg_content, "msg_type": "text"},
            wait_after=2,
        )

        # ===== 模块4: 消息接收 =====
        self.log("\n## 模块4: 消息接收\n")

        self.verify_step(
            "H15_消息接收", "getRecentMessages",
            "获取最近消息列表",
            lambda: self.phone.open_wechat(),
            params={"limit": 5},
        )

        self.verify_step(
            "H15_消息接收", "getMessages",
            "获取filehelper的消息记录",
            lambda: None,
            params={"conversation_id": "filehelper", "limit": 5},
        )

        # ===== 模块5: 好友管理 =====
        self.log("\n## 模块5: 好友管理\n")

        self.verify_step(
            "H19_好友管理", "setFriendRemark",
            "修改文件传输助手的备注名为'SDK文件助手'",
            lambda: self.phone.open_wechat_contacts(),
            params={"wxid": "filehelper", "remark": "SDK文件助手"},
            wait_after=2,
        )

        self.verify_step(
            "H18_好友请求", "getFriendRequests",
            "获取好友请求列表",
            lambda: None,
            params={"limit": 5},
        )

        # ===== 模块6: 群管理 =====
        self.log("\n## 模块6: 群管理\n")

        self.verify_step(
            "H22_群管理", "getGroups",
            "获取群聊列表",
            lambda: self.phone.open_wechat(),
            params={"limit": 10},
        )

        # ===== 模块7: 个人信息 =====
        self.log("\n## 模块7: 个人信息\n")

        self.verify_step(
            "H23_账号管理", "getProfile",
            "打开'我'页面 -> 获取个人资料",
            lambda: self.phone.open_wechat_me(),
            params={},
        )

        self.verify_step(
            "H23_账号管理", "checkAccountStatus",
            "检查账号状态",
            lambda: None,
            params={},
        )

        # ===== 模块8: 标签 =====
        self.log("\n## 模块8: 标签\n")

        self.verify_step(
            "H28_标签", "getLabels",
            "获取标签列表",
            lambda: None,
            params={},
        )

        # ===== 模块9: 朋友圈 =====
        self.log("\n## 模块9: 朋友圈\n")

        self.verify_step(
            "H21_朋友圈", "getMoments",
            "打开发现页 -> 获取朋友圈动态",
            lambda: self.phone.open_wechat_discover(),
            params={"wxid": "", "limit": 3},
        )

        # ===== 模块10: 搜索 =====
        self.log("\n## 模块10: 搜索\n")

        self.verify_step(
            "H31_搜索", "globalSearch",
            "全局搜索'微信'",
            lambda: self.phone.open_wechat(),
            params={"keyword": "微信", "limit": 5},
        )

        # ===== 模块11: 设备信息 =====
        self.log("\n## 模块11: 设备信息\n")

        self.verify_step(
            "H39_设备信息", "getDeviceInfo",
            "获取设备信息",
            lambda: None,
            params={},
        )

        self.verify_step(
            "H39_设备信息", "getNetworkInfo",
            "获取网络信息",
            lambda: None,
            params={},
        )

        self.verify_step(
            "H39_设备信息", "getStorageInfo",
            "获取存储信息",
            lambda: None,
            params={},
        )

        # ===== 模块12: 账号安全 =====
        self.log("\n## 模块12: 账号安全\n")

        self.verify_step(
            "H24_账号安全", "getLoginDevices",
            "获取登录设备列表",
            lambda: self.phone.open_wechat_me(),
            params={},
        )

        self.verify_step(
            "H35_注册登录", "checkLoginState",
            "检查登录状态",
            lambda: None,
            params={},
        )

        # ===== 汇总 =====
        self.summarize()

    def summarize(self):
        """汇总结果"""
        total = len(self.results)
        passed = sum(1 for r in self.results if r["status"] in ("PASS", "EXEC_ERR"))
        missing = sum(1 for r in self.results if r["status"] == "MISSING")

        self.log(f"\n{'='*60}")
        self.log(f"  验证完成!")
        self.log(f"  总计: {total}")
        self.log(f"  通过: {passed}")
        self.log(f"  缺失: {missing}")
        self.log(f"  通过率: {passed/total*100:.1f}%")
        self.log(f"{'='*60}")

        # 保存日志
        with open(LOG_FILE, "w", encoding="utf-8") as f:
            f.write("\n".join(self.log_lines))
        print(f"\n日志: {LOG_FILE}")

        # 保存JSON
        result_file = os.path.join(OUTPUT_DIR, "step_results.json")
        with open(result_file, "w", encoding="utf-8") as f:
            json.dump({
                "time": datetime.now().isoformat(),
                "total": total,
                "passed": passed,
                "missing": missing,
                "results": self.results,
            }, f, ensure_ascii=False, indent=2, default=str)
        print(f"结果: {result_file}")

    def cleanup(self):
        try:
            self.script.unload()
            self.session.detach()
        except:
            pass


def main():
    v = StepVerifier()
    try:
        if v.connect():
            v.run_all()
    finally:
        v.cleanup()


if __name__ == "__main__":
    main()