#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
工作手机SDK - 截图验证朋友圈和消息发送全过程
"""
import frida, time, json, os, subprocess
from datetime import datetime

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

os.makedirs(SCREENSHOT_DIR, exist_ok=True)

def pull_screenshot(remote_path, local_name):
    """从手机拉取截图到本地"""
    local_path = f"{SCREENSHOT_DIR}/{local_name}"
    result = subprocess.run(
        ["adb", "-s", "192.168.110.80:5555", "pull", remote_path, local_path],
        capture_output=True, text=True
    )
    if result.returncode == 0:
        print(f"    截图已保存: {local_name}")
        return local_path
    else:
        print(f"    截图拉取失败: {result.stderr}")
        return None

print(f"[{datetime.now().strftime('%H:%M:%S')}] 连接Frida: {PHONE_IP}:{FRIDA_PORT}")
dm = frida.get_device_manager()
device = dm.add_remote_device(f"{PHONE_IP}:{FRIDA_PORT}")
session = device.attach(WECHAT_PID)
print(f"[OK] 附加微信 PID={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()
print("[OK] Hook已加载，等待3秒...")
time.sleep(3)
rpc = script.exports_sync

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

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

# 联系人统计
r_total = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM rcontact"})
total_contacts = r_total.get("rows", [{}])[0].get("cnt", "?")
print(f"  rcontact总记录: {total_contacts}")

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

r_groups = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM chatroom"})
groups_count = r_groups.get("rows", [{}])[0].get("cnt", "?")
print(f"  群组数量(chatroom表): {groups_count}")

r_labels = rpc.raw_sql({"sql": "SELECT COUNT(*) as cnt FROM contactlabelinfo"})
labels_count = r_labels.get("rows", [{}])[0] if r_labels.get("rows") else {}
labels_count = labels_count.get("cnt", "表不存在")
print(f"  标签数: {labels_count}")

# 群组名称（chatroom表有名称的）
r_named_groups = rpc.raw_sql({"sql": "SELECT chatroomname, chatroomnick FROM chatroom WHERE chatroomnick != '' AND chatroomnick IS NOT NULL LIMIT 20"})
print(f"  有名称的群组: {len(r_named_groups.get('rows', []))}个")
for g in r_named_groups.get("rows", [])[:5]:
    print(f"    {g.get('chatroomname','')} -> {g.get('chatroomnick','')}")

results["data_stats"] = {
    "total_rcontact": total_contacts,
    "type3_friends": type3_count,
    "groups_count": groups_count,
    "labels_count": labels_count,
    "named_groups_sample": r_named_groups.get("rows", [])[:10]
}

# ============================================================
# 阶段2：截图验证朋友圈发送全过程
# ============================================================
print("\n" + "="*50)
print("阶段2：朋友圈发送全过程截图验证")
print("="*50)

# 2a. 导航到主界面
print("  2a. 导航到微信主界面...")
rpc.navigate_to_main({})
time.sleep(1.5)
rpc.take_screenshot({"path": "/sdcard/moment_step0_main.png"})
pull_screenshot("/sdcard/moment_step0_main.png", f"moment_step0_main_{ts}.png")

# 2b. 发送朋友圈前截图
moment_content = f"[SDK验证{datetime.now().strftime('%m/%d %H:%M')}] 工作手机Frida真机验证-朋友圈功能正常 #工作手机SDK"
print(f"  2b. 准备发送朋友圈: {moment_content[:30]}...")
rpc.take_screenshot({"path": "/sdcard/moment_step1_before.png"})
pull_screenshot("/sdcard/moment_step1_before.png", f"moment_step1_before_{ts}.png")

# 2c. 发送朋友圈
print("  2c. 执行发送朋友圈...")
moment_result = rpc.post_moments({"content": moment_content, "image_urls": []})
print(f"  发送结果: {json.dumps(moment_result, ensure_ascii=False)}")
results["moment_send"] = moment_result
time.sleep(2)

# 2d. 发送后截图
rpc.take_screenshot({"path": "/sdcard/moment_step2_after.png"})
pull_screenshot("/sdcard/moment_step2_after.png", f"moment_step2_after_{ts}.png")

# 2e. 读取朋友圈验证是否发送成功
print("  2e. 读取朋友圈验证...")
moments = rpc.get_moments({"limit": 5})
print(f"  朋友圈条数: {moments.get('count', 0)}")
if moments.get("moments"):
    for m in moments["moments"][:3]:
        print(f"    sns_id={m.get('sns_id','')} content={str(m.get('content',''))[:50]}")
results["moments_after"] = moments

# ============================================================
# 阶段3：截图验证消息发送全过程
# ============================================================
print("\n" + "="*50)
print("阶段3：消息发送全过程截图验证")
print("="*50)

# 3a. 导航到文件传输助手
print("  3a. 导航到文件传输助手...")
rpc.navigate_to_chat({"wxid": "filehelper"})
time.sleep(1.5)
rpc.take_screenshot({"path": "/sdcard/msg_step1_before.png"})
pull_screenshot("/sdcard/msg_step1_before.png", f"msg_step1_before_{ts}.png")

# 3b. 发送消息
msg_content = f"[SDK验证{datetime.now().strftime('%H:%M:%S')}] 消息发送功能正常 | 好友:{type3_count}人 | 群组:{groups_count}个"
print(f"  3b. 发送消息: {msg_content[:50]}...")
msg_result = rpc.send_message({"to_id": "filehelper", "content": msg_content, "msg_type": "text"})
print(f"  发送结果: {json.dumps(msg_result, ensure_ascii=False)}")
results["message_send"] = msg_result
time.sleep(1.5)

# 3c. 发送后截图（消息应该出现在聊天界面）
rpc.take_screenshot({"path": "/sdcard/msg_step2_after.png"})
pull_screenshot("/sdcard/msg_step2_after.png", f"msg_step2_after_{ts}.png")

# 3d. 再发一条带时间戳的消息确认
msg2 = f"[SDK验证] 第2条测试消息 {datetime.now().strftime('%H:%M:%S')}"
rpc.send_message({"to_id": "filehelper", "content": msg2, "msg_type": "text"})
time.sleep(1)
rpc.take_screenshot({"path": "/sdcard/msg_step3_confirm.png"})
pull_screenshot("/sdcard/msg_step3_confirm.png", f"msg_step3_confirm_{ts}.png")

# ============================================================
# 阶段4：修复getContacts返回真实数量
# ============================================================
print("\n" + "="*50)
print("阶段4：验证修复后的联系人接口")
print("="*50)

# 用getContactsFull获取真实数量
contacts_full = rpc.get_contacts_full({"limit": 500, "filter_type": 3})
print(f"  getContactsFull(type=3): {contacts_full.get('count', 0)}个（limit=500）")
print(f"  真实总数(SQL): {type3_count}个")

# 用原始getContacts（limit=500）
contacts_orig = rpc.get_contacts({"limit": 500})
print(f"  getContacts(limit=500): {contacts_orig.get('count', 0)}个")

results["contacts_comparison"] = {
    "get_contacts_500": contacts_orig.get("count", 0),
    "get_contacts_full_type3_500": contacts_full.get("count", 0),
    "raw_sql_type3": type3_count
}

# ============================================================
# 保存结果
# ============================================================
report = {
    "title": "工作手机SDK - 截图验证报告",
    "time": datetime.now().isoformat(),
    "device": PHONE_IP,
    "real_data": {
        "total_rcontact": total_contacts,
        "type3_friends": type3_count,
        "groups": groups_count,
        "labels": labels_count
    },
    "results": results,
    "screenshots": [
        f"moment_step0_main_{ts}.png",
        f"moment_step1_before_{ts}.png",
        f"moment_step2_after_{ts}.png",
        f"msg_step1_before_{ts}.png",
        f"msg_step2_after_{ts}.png",
        f"msg_step3_confirm_{ts}.png",
    ]
}
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}")
print(f"截图目录: {SCREENSHOT_DIR}")

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