#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UIAutomator2 + Frida 联合验证
真正控制微信界面发送消息，截图全过程
"""
import uiautomator2 as u2
import frida, time, json, subprocess, os
from datetime import datetime

PHONE_IP = "192.168.110.80"
ADB_SERIAL = "192.168.110.80:5555"
FRIDA_PORT = 27042
WECHAT_PID = 9462
BASE = "/Users/karuo/Documents/开发/2、私域银行/工作手机"
HOOK_SCRIPT = BASE + "/sdk/agent/hook/wechat_hook_v2.js"
SC_DIR = BASE + "/开发文档/6、测试/live_verify_20260518/screenshots"
RESULT_FILE = BASE + "/verification_u2_20260518.json"

os.makedirs(SC_DIR, exist_ok=True)

def adb_sc(name, wait=0.8):
    """ADB截图"""
    time.sleep(wait)
    remote = f"/sdcard/{name}.png"
    local = f"{SC_DIR}/{name}.png"
    subprocess.run(["adb", "-s", ADB_SERIAL, "shell", "screencap", "-p", remote], capture_output=True)
    subprocess.run(["adb", "-s", ADB_SERIAL, "pull", remote, local], capture_output=True)
    size = os.path.getsize(local) if os.path.exists(local) else 0
    print(f"  [截图] {name}.png ({size//1024}KB)")
    return local if size > 10000 else None

ts = datetime.now().strftime("%H%M%S")
results = {}

# ============================================================
# 连接UIAutomator2
# ============================================================
print(f"[{datetime.now().strftime('%H:%M:%S')}] 连接UIAutomator2...")
d = u2.connect(PHONE_IP)
info = d.info
print(f"[OK] 设备: {info.get('productName')} | 屏幕: {'亮' if info.get('screenOn') else '暗'} | 当前包: {info.get('currentPackageName')}")

# ============================================================
# 连接Frida（用于数据读取）
# ============================================================
print(f"[{datetime.now().strftime('%H:%M:%S')}] 连接Frida...")
dm = frida.get_device_manager()
device = dm.add_remote_device(f"{PHONE_IP}:{FRIDA_PORT}")
session = device.attach(WECHAT_PID)
with open(HOOK_SCRIPT, encoding="utf-8") as f:
    src = f.read()
script = session.create_script(src)
script.on("message", lambda m, d: None)
script.load()
time.sleep(3)
rpc = script.exports_sync
print("[OK] Frida已连接")

# ============================================================
# 阶段1：真实数据统计（Frida读取）
# ============================================================
print("\n" + "="*50)
print("阶段1：真实数据统计")
print("="*50)

r_type3 = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM rcontact WHERE type=3"})
type3_count = r_type3.get("rows", [{}])[0].get("cnt", "?")

r_groups = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM chatroom"})
groups_count = r_groups.get("rows", [{}])[0].get("cnt", "?")

r_total = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM rcontact"})
total_count = r_total.get("rows", [{}])[0].get("cnt", "?")

print(f"  真实好友(type=3): {type3_count}")
print(f"  群组数量: {groups_count}")
print(f"  rcontact总记录: {total_count}")

results["data_stats"] = {
    "type3_friends": type3_count,
    "groups": groups_count,
    "total_rcontact": total_count
}

# ============================================================
# 阶段2：UIAutomator2发送消息（真实控制）
# ============================================================
print("\n" + "="*50)
print("阶段2：UIAutomator2真实发送消息")
print("="*50)

# 2a. 确保微信在前台
d.app_start("com.tencent.mm", activity=".ui.LauncherUI")
time.sleep(2)
adb_sc(f"u2_step1_main_{ts}")

# 2b. 搜索文件传输助手
print("  2b. 点击搜索...")
# 尝试点击搜索按钮
try:
    d(resourceId="com.tencent.mm:id/f8t").click()  # 搜索按钮
except:
    pass
try:
    d(description="搜索").click()
except:
    pass
time.sleep(1)
adb_sc(f"u2_step2_search_{ts}")

# 2c. 输入"文件传输助手"
print("  2c. 搜索文件传输助手...")
try:
    d(focused=True).set_text("文件传输助手")
except:
    d.send_keys("文件传输助手")
time.sleep(1.5)
adb_sc(f"u2_step3_search_result_{ts}")

# 2d. 点击文件传输助手
print("  2d. 点击文件传输助手...")
try:
    d(text="文件传输助手").click()
except:
    # 点击第一个搜索结果
    d(resourceId="com.tencent.mm:id/b5e").click()
time.sleep(1.5)
adb_sc(f"u2_step4_chat_{ts}")

# 2e. 点击输入框
print("  2e. 点击输入框...")
try:
    d(resourceId="com.tencent.mm:id/bkk").click()  # 消息输入框
except:
    try:
        d(className="android.widget.EditText").click()
    except:
        pass
time.sleep(0.5)

# 2f. 输入消息
msg_content = f"[UIAutomator2真机验证 {datetime.now().strftime('%H:%M:%S')}] 好友:{type3_count}人 | 群组:{groups_count}个 | 总联系人:{total_count}条 | Frida+U2控制成功"
print(f"  2f. 输入消息: {msg_content[:50]}...")
try:
    d(focused=True).set_text(msg_content)
except:
    d.send_keys(msg_content)
time.sleep(0.8)
adb_sc(f"u2_step5_typed_{ts}")

# 2g. 点击发送按钮
print("  2g. 点击发送...")
try:
    d(text="发送").click()
except:
    try:
        d(resourceId="com.tencent.mm:id/bkl").click()  # 发送按钮
    except:
        # 按回车
        d.press("enter")
time.sleep(1.5)

# 2h. 截图（发送后）
sc_after = adb_sc(f"u2_step6_after_send_{ts}")
print(f"  发送后截图: {sc_after}")

# 2i. 验证DB中是否有消息
print("\n  验证DB消息...")
r_db = rpc.raw_sql({"sql": "SELECT msgId, talker, content, isSend, createTime FROM message WHERE talker='filehelper' ORDER BY createTime DESC LIMIT 5"})
print(f"  DB查询结果: {r_db.get('count', 0)}条")
for row in r_db.get("rows", []):
    print(f"    isSend={row.get('isSend')} | {str(row.get('content',''))[:60]}")
results["message_send_u2"] = {
    "content": msg_content,
    "db_messages": r_db.get("rows", [])
}

# ============================================================
# 阶段3：UIAutomator2发朋友圈
# ============================================================
print("\n" + "="*50)
print("阶段3：UIAutomator2发朋友圈")
print("="*50)

# 3a. 返回主界面
d.press("back")
d.press("back")
time.sleep(1)

# 3b. 点击"发现"tab
print("  3b. 点击发现...")
try:
    d(text="发现").click()
except:
    pass
time.sleep(1)
adb_sc(f"u2_step7_discover_{ts}")

# 3c. 点击"朋友圈"
print("  3c. 点击朋友圈...")
try:
    d(text="朋友圈").click()
except:
    pass
time.sleep(2)
adb_sc(f"u2_step8_moments_{ts}")

# 3d. 点击发布按钮（相机图标）
print("  3d. 点击发布朋友圈...")
try:
    d(resourceId="com.tencent.mm:id/ej6").click()  # 相机按钮
except:
    try:
        d(description="拍摄，按钮").click()
    except:
        # 长按相机按钮发纯文字
        try:
            d(className="android.widget.ImageView", instance=0).long_click()
        except:
            pass
time.sleep(1.5)
adb_sc(f"u2_step9_moment_compose_{ts}")

# 3e. 输入朋友圈内容
moment_content = f"[SDK验证{datetime.now().strftime('%m/%d %H:%M')}] 工作手机Frida+UIAutomator2真机验证成功 好友:{type3_count}人"
print(f"  3e. 输入朋友圈内容: {moment_content[:40]}...")
try:
    d(resourceId="com.tencent.mm:id/bkk").set_text(moment_content)
except:
    try:
        d(className="android.widget.EditText").set_text(moment_content)
    except:
        d.send_keys(moment_content)
time.sleep(0.8)
adb_sc(f"u2_step10_moment_typed_{ts}")

# 3f. 点击发送
print("  3f. 发送朋友圈...")
try:
    d(text="发表").click()
except:
    try:
        d(text="发送").click()
    except:
        pass
time.sleep(2)
adb_sc(f"u2_step11_moment_sent_{ts}")

results["moment_send_u2"] = {"content": moment_content}

# ============================================================
# 保存结果
# ============================================================
report = {
    "title": "工作手机SDK - UIAutomator2真机验证报告",
    "time": datetime.now().isoformat(),
    "device": PHONE_IP,
    "real_data": {
        "type3_friends": type3_count,
        "groups": groups_count,
        "total_rcontact": total_count
    },
    "screenshots": [f"{SC_DIR}/{f}" for f in [
        f"u2_step1_main_{ts}.png",
        f"u2_step4_chat_{ts}.png",
        f"u2_step5_typed_{ts}.png",
        f"u2_step6_after_send_{ts}.png",
        f"u2_step8_moments_{ts}.png",
        f"u2_step11_moment_sent_{ts}.png",
    ]],
    "results": results
}
with open(RESULT_FILE, "w", encoding="utf-8") as f:
    json.dump(report, f, ensure_ascii=False, indent=2, default=str)
print(f"\n结果已保存: {RESULT_FILE}")

script.unload()
session.detach()
print("完成！")