#!/usr/bin/env python3
"""
工作手机SDK - 真机验证脚本
通过WiFi连接frida-server，加载wechat_hook_v2.js，逐个验证所有RPC方法
"""
import frida
import json
import time
import os
import sys
import subprocess
from datetime import datetime
from typing import Dict, Any, List, Tuple

# 配置
DEVICE_IP = "192.168.0.12"
FRIDA_PORT = 27042
WECHAT_PACKAGE = "com.tencent.mm"
DEVICE_SERIAL = "xgfe65eimrrofyws"

# Hook脚本路径
SCRIPT_PATH = os.path.join(os.path.dirname(__file__), '..', 'agent', 'hook', 'wechat_hook_v2.js')
if not os.path.exists(SCRIPT_PATH):
    # 尝试从GitHub仓库路径
    SCRIPT_PATH = os.path.expanduser("~/Documents/GitHub/workphone-sdk/sdk/agent/hook/wechat_hook_v2.js")

# 截图保存目录
SCREENSHOT_DIR = os.path.join(
    os.path.dirname(__file__), '..', '..', '..',
    '开发文档', '8、部署', '05-测试验收', 'screenshots', 'verification_110',
)
os.makedirs(SCREENSHOT_DIR, exist_ok=True)

# 结果文件
RESULT_FILE = os.path.join(os.path.dirname(__file__), '..', '..', 'verification_results.json')


def take_screenshot(name: str) -> str:
    """截取手机屏幕"""
    filename = f"{name}_{int(time.time())}.png"
    filepath = os.path.join(SCREENSHOT_DIR, filename)
    try:
        subprocess.run(
            ["adb", "-s", DEVICE_SERIAL, "shell", "screencap", "-p", f"/sdcard/screenshot.png"],
            capture_output=True, timeout=10
        )
        subprocess.run(
            ["adb", "-s", DEVICE_SERIAL, "pull", "/sdcard/screenshot.png", filepath],
            capture_output=True, timeout=10
        )
        if os.path.exists(filepath):
            return filepath
    except Exception as e:
        print(f"  截图失败: {e}")
    return ""


def connect_frida():
    """连接Frida并attach微信"""
    print(f"[1/4] 连接 Frida Server ({DEVICE_IP}:{FRIDA_PORT})...")
    mgr = frida.get_device_manager()
    device = mgr.add_remote_device(f"{DEVICE_IP}:{FRIDA_PORT}")
    print(f"  设备: {device.name}")

    # 找微信主进程
    print("[2/4] 查找微信主进程...")
    procs = [p for p in device.enumerate_processes() if p.name == WECHAT_PACKAGE]
    if not procs:
        # 尝试通过应用列表
        apps = [a for a in device.enumerate_applications() if a.identifier == WECHAT_PACKAGE]
        if apps and apps[0].pid > 0:
            pid = apps[0].pid
        else:
            print("  微信主进程未找到！尝试启动微信...")
            subprocess.run(
                ["adb", "-s", DEVICE_SERIAL, "shell", "am", "start", "-n", 
                 "com.tencent.mm/.ui.LauncherUI"],
                capture_output=True, timeout=10
            )
            time.sleep(3)
            procs = [p for p in device.enumerate_processes() if p.name == WECHAT_PACKAGE]
            if not procs:
                raise Exception("微信主进程无法启动")
            pid = procs[0].pid
    else:
        pid = procs[0].pid

    print(f"  微信 PID: {pid}")

    # Attach
    print("[3/4] Attach 微信进程...")
    session = device.attach(pid)
    print("  Attach 成功!")

    # 加载脚本
    print(f"[4/4] 加载 Hook 脚本...")
    if not os.path.exists(SCRIPT_PATH):
        raise FileNotFoundError(f"Hook脚本不存在: {SCRIPT_PATH}")

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

    script = session.create_script(source)
    
    messages = []
    def on_message(message, data):
        if message.get('type') == 'send':
            messages.append(message.get('payload', {}))
        elif message.get('type') == 'error':
            print(f"  [脚本错误] {message.get('description', '')}")

    script.on('message', on_message)
    script.load()
    
    rpc = script.exports_sync
    
    # 验证连接
    try:
        pong = rpc.ping()
        print(f"  RPC ping: {pong}")
    except Exception as e:
        print(f"  RPC ping 失败: {e}")

    return device, session, script, rpc


# 所有需要验证的操作（按模块分组）
VERIFICATION_PLAN = {
    "H15_消息接收": {
        "get_messages": {"conversation_id": "", "limit": 3},
        "get_recent_messages": {"limit": 5},
        "search_messages": {"keyword": "你好", "limit": 3},
    },
    "H16_联系人": {
        "get_contacts": {"limit": 10},
        "get_contact_info": {"wxid": "filehelper"},
        "search_contacts": {"keyword": "文件", "limit": 5},
    },
    "H17_消息发送": {
        "send_message": {"to_id": "filehelper", "content": "[SDK验证] 消息发送测试 " + datetime.now().strftime("%H:%M:%S"), "msg_type": "text"},
    },
    "H18_好友请求": {
        "get_friend_requests": {"limit": 5},
    },
    "H19_好友管理": {
        "set_friend_remark": {"wxid": "filehelper", "remark": "文件传输助手"},
    },
    "H20_朋友圈发布": {
        # 跳过实际发布，只验证接口存在
    },
    "H21_朋友圈浏览": {
        "get_moments": {"wxid": "", "limit": 3},
    },
    "H22_群管理": {
        "get_groups": {"limit": 10},
    },
    "H23_账号管理": {
        "get_profile": {},
        "check_account_status": {},
    },
    "H24_账号安全": {
        "get_login_devices": {},
    },
    "H25_支付": {
        "get_wallet_balance": {},
    },
    "H26_二维码": {
        "generate_my_qr_code": {},
    },
    "H27_视频号": {
        "browse_channels": {"limit": 3},
    },
    "H28_标签": {
        "get_labels": {},
    },
    "H29_收藏": {
        "get_favorites": {"limit": 5},
    },
    "H30_设置": {
        # 只读操作
    },
    "H31_搜索": {
        "global_search": {"keyword": "微信", "limit": 5},
    },
    "H32_小程序": {
        "get_recent_mini_programs": {},
    },
    "H33_文件传输": {
        "send_image": {"to_id": "filehelper", "image_path": "/sdcard/DCIM/test.jpg"},
    },
    "H34_消息转发": {
        # 需要msg_id，跳过
    },
    "H35_注册登录": {
        "check_login_state": {},
        "get_sim_phone": {},
    },
    "H36_公众号": {
        "get_official_accounts": {"limit": 5},
    },
    "H37_表情": {
        # 需要emoji_md5，跳过
    },
    "H38_浮窗": {
        # 需要特定条件，跳过
    },
    "H39_设备信息": {
        "get_device_info": {},
        "get_storage_info": {},
        "get_network_info": {},
    },
    "系统_Hook状态": {
        "get_hook_status": {},
        "get_process_info": {},
        "get_wechat_version": {},
    },
}

# RPC方法名映射（snake_case -> camelCase）
def to_camel(name: str) -> str:
    parts = name.split('_')
    return parts[0] + ''.join(p.capitalize() for p in parts[1:])


def run_verification():
    """运行完整验证"""
    print("=" * 60)
    print("  工作手机SDK - Frida无线真机验证")
    print(f"  设备: {DEVICE_IP}:{FRIDA_PORT}")
    print(f"  时间: {datetime.now().isoformat()}")
    print("=" * 60)
    print()

    # 连接
    device, session, script, rpc = connect_frida()
    print("\n✅ Frida 连接成功！开始验证...\n")

    # 初始截图
    take_screenshot("00_initial_state")

    results = []
    total = 0
    passed = 0
    failed = 0
    
    for module_name, actions in VERIFICATION_PLAN.items():
        print(f"\n{'='*50}")
        print(f"  模块: {module_name}")
        print(f"{'='*50}")

        for action, params in actions.items():
            total += 1
            rpc_method = to_camel(action)
            print(f"\n  [{total}] {action} -> rpc.{rpc_method}()")
            print(f"      参数: {json.dumps(params, ensure_ascii=False)[:80]}")

            start_time = time.time()
            try:
                fn = getattr(rpc, rpc_method, None)
                if fn is None:
                    # 尝试下划线版本
                    fn = getattr(rpc, action, None)
                
                if fn is None:
                    result = {"success": False, "error": f"RPC方法不存在: {rpc_method}"}
                    status = "FAIL"
                    failed += 1
                else:
                    result = fn(params) if params else fn({})
                    if isinstance(result, dict) and result.get("success", False):
                        status = "PASS"
                        passed += 1
                    elif isinstance(result, dict) and "error" in result:
                        status = "FAIL"
                        failed += 1
                    else:
                        # 有返回就算通过
                        status = "PASS"
                        passed += 1

            except Exception as e:
                result = {"success": False, "error": str(e)}
                status = "FAIL"
                failed += 1

            latency = int((time.time() - start_time) * 1000)
            
            icon = "✅" if status == "PASS" else "❌"
            print(f"      {icon} {status} ({latency}ms)")
            if isinstance(result, dict):
                # 打印简要结果
                preview = json.dumps(result, ensure_ascii=False)[:120]
                print(f"      结果: {preview}")

            # 截图
            screenshot = take_screenshot(f"{module_name}_{action}")

            results.append({
                "module": module_name,
                "action": action,
                "rpc_method": rpc_method,
                "params": params,
                "status": status,
                "latency_ms": latency,
                "result": result if isinstance(result, dict) else str(result),
                "screenshot": screenshot,
                "timestamp": datetime.now().isoformat(),
            })

            # 每个操作间隔一下，避免触发风控
            time.sleep(0.5)

    # 最终截图
    take_screenshot("99_final_state")

    # 保存结果
    report = {
        "title": "工作手机SDK - Frida无线真机验证报告",
        "device": {
            "ip": DEVICE_IP,
            "port": FRIDA_PORT,
            "serial": DEVICE_SERIAL,
            "model": "2312DRAABC",
            "android": "13",
        },
        "test_time": datetime.now().isoformat(),
        "total": total,
        "passed": passed,
        "failed": failed,
        "pass_rate": f"{passed/total*100:.1f}%" if total > 0 else "0%",
        "results": results,
    }

    with open(RESULT_FILE, 'w', encoding='utf-8') as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    print(f"\n{'='*60}")
    print(f"  验证完成!")
    print(f"  总计: {total} | 通过: {passed} | 失败: {failed}")
    print(f"  通过率: {passed/total*100:.1f}%" if total > 0 else "  无测试")
    print(f"  结果: {RESULT_FILE}")
    print(f"  截图: {SCREENSHOT_DIR}")
    print(f"{'='*60}")

    # 清理
    try:
        script.unload()
        session.detach()
    except:
        pass

    return report


if __name__ == "__main__":
    try:
        report = run_verification()
    except Exception as e:
        print(f"\n❌ 验证失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)