#!/usr/bin/env python3
"""
工作手机 · WebSocket + 连接态全量自检

用法（项目根或 sdk 目录）:
  python3 scripts/ws_full_test.py
  python3 scripts/ws_full_test.py -d xgfe65eimrrofyws --base http://127.0.0.1:8899

检查项：health / connection/status / protocol / modes /
       hook/probe / hook/execute(只读) / 指定设备是否在线 /
       手机 Agent 包是否安装（需 ADB）
"""
from __future__ import annotations

import argparse
import asyncio
import json
import subprocess
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional
from urllib.error import URLError
from urllib.request import Request, urlopen

DEFAULT_DEVICE = "xgfe65eimrrofyws"
DEFAULT_BASE = "http://127.0.0.1:8899"
AGENT_PKG = "com.system.cloudservice"
WECHAT_PKG = "com.tencent.mm"
OUT_DIR = Path(__file__).resolve().parent.parent / "tmp"


def _get(url: str, timeout: float = 15) -> Dict[str, Any]:
    req = Request(url, headers={"X-API-Key": "workphone-secret-key"})
    with urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode())


def _post(url: str, body: dict, timeout: float = 90) -> Dict[str, Any]:
    data = json.dumps(body).encode()
    req = Request(
        url,
        data=data,
        headers={"Content-Type": "application/json", "X-API-Key": "workphone-secret-key"},
        method="POST",
    )
    with urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode())


def _adb(*args: str) -> str:
    try:
        r = subprocess.run(["adb"] + list(args), capture_output=True, text=True, timeout=12)
        return (r.stdout or r.stderr or "").strip()
    except Exception as e:
        return f"adb error: {e}"


def _check_adb_install(serial: str) -> Dict[str, Any]:
    out = {
        "adb_serials": [],
        "adb_serial_resolved": None,
        "expected_connected": False,
        "agent_pkg_installed": None,
        "agent_pkg_enabled": None,
        "wechat_installed": None,
        "launcher_activity": None,
    }
    lines = _adb("devices").splitlines()[1:]
    serials = [ln.split()[0] for ln in lines if ln.strip() and "device" in ln and "offline" not in ln]
    out["adb_serials"] = serials
    resolved = None
    for s in serials:
        if s == serial:
            resolved = s
            break
        ro = _adb("-s", s, "shell", "getprop", "ro.serialno")
        if ro.strip() == serial:
            resolved = s
            break
    out["adb_serial_resolved"] = resolved
    out["expected_connected"] = resolved is not None
    if not out["expected_connected"]:
        return out
    prefix = ["-s", resolved]
    pm = _adb(*prefix, "shell", "pm", "path", AGENT_PKG)
    out["agent_pkg_installed"] = AGENT_PKG in pm and "package:" in pm
    dump = _adb(*prefix, "shell", "pm", "dump", AGENT_PKG)
    out["agent_pkg_enabled"] = "enabled=1" in dump or "enabled=0" not in dump
    wc = _adb(*prefix, "shell", "pm", "path", WECHAT_PKG)
    out["wechat_installed"] = WECHAT_PKG in wc
    launcher = _adb(
        *prefix,
        "shell",
        "cmd",
        "package",
        "resolve-activity",
        "--brief",
        "-c",
        "android.intent.category.LAUNCHER",
        AGENT_PKG,
    )
    out["launcher_activity"] = launcher if launcher and "No activity" not in launcher else None
    return out


async def _ws_protocol_smoke(base_ws: str, device_id: str) -> Dict[str, Any]:
    """用 websockets 客户端走 register → heartbeat 协议（不替代真机 Agent）。"""
    try:
        import websockets
    except ImportError:
        return {"ok": False, "error": "pip install websockets"}

    url = f"{base_ws.rstrip('/')}/{device_id}_ws_smoke"
    result: Dict[str, Any] = {"ok": False, "url": url, "steps": []}
    try:
        async with websockets.connect(url, open_timeout=8) as ws:
            await ws.send(
                json.dumps(
                    {
                        "type": "register",
                        "data": {
                            "project_id": "cunkebao",
                            "model": "WS-Smoke-Test",
                            "capabilities": ["event", "device_request"],
                            "source": "ws_full_test.py",
                        },
                    }
                )
            )
            msg = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
            result["steps"].append({"register_ack": msg.get("type")})
            await ws.send(json.dumps({"type": "heartbeat", "device_id": f"{device_id}_ws_smoke"}))
            msg2 = json.loads(await asyncio.wait_for(ws.recv(), timeout=5))
            result["steps"].append({"heartbeat_ack": msg2.get("type")})
            result["ok"] = msg.get("type") == "registered" and msg2.get("type") in ("pong", "heartbeat_ack", "registered")
    except Exception as e:
        result["error"] = str(e)[:300]
    return result


def run_tests(base: str, device_id: str) -> Dict[str, Any]:
    base = base.rstrip("/")
    ws_base = base.replace("http://", "ws://").replace("https://", "wss://") + "/ws/device"
    report: Dict[str, Any] = {
        "at": datetime.now().isoformat(),
        "base": base,
        "expected_device_id": device_id,
        "checks": [],
        "summary": {"pass": 0, "fail": 0, "warn": 0},
    }

    def add(name: str, ok: bool, detail: Any = None, level: str = "pass"):
        if not ok and level == "pass":
            level = "fail"
        report["checks"].append({"name": name, "ok": ok, "level": level, "detail": detail})
        key = level if level in ("pass", "fail", "warn") else ("pass" if ok else "fail")
        report["summary"][key] += 1

    # 1 health
    try:
        h = _get(f"{base}/health")
        add("SDK /health", h.get("status") == "healthy", h)
    except URLError as e:
        add("SDK /health", False, str(e))

    # 2 connection status
    conn = {}
    try:
        conn = _get(f"{base}/api/v3/connection/status").get("data", {})
        ws_ids = conn.get("online_device_ids") or []
        adb_serials = conn.get("adb_serials") or []
        on_ws = device_id in ws_ids
        on_adb = device_id in adb_serials or any(
            _adb("-s", s, "shell", "getprop", "ro.serialno").strip() == device_id for s in adb_serials
        )
        add(
            "指定设备 WS 在线",
            on_ws,
            {"online_ws_count": conn.get("online_ws_count"), "ids": ws_ids},
            "warn" if not on_ws else "pass",
        )
        add(
            "指定设备 ADB 在线",
            on_adb,
            {"adb_count": conn.get("adb_count"), "serials": adb_serials},
            "warn" if not on_adb else "pass",
        )
        report["connection"] = conn
    except URLError as e:
        add("connection/status", False, str(e))

    # 3 protocol doc
    try:
        p = _get(f"{base}/api/v3/connection/protocol")
        add("connection/protocol", p.get("code") == 200, {"version": p.get("data", {}).get("version")})
    except URLError as e:
        add("connection/protocol", False, str(e))

    # 4 connection modes
    try:
        m = _get(f"{base}/api/v3/connection/modes/{device_id}")
        modes = m.get("data", {}).get("modes") or []
        best = m.get("data", {}).get("best_mode")
        add("connection/modes", bool(modes), {"best_mode": best, "count": len(modes)})
        report["modes"] = m.get("data")
    except URLError as e:
        add("connection/modes", False, str(e))

    # 5 WS 协议冒烟（服务端 WS 通路）
    try:
        ws_smoke = asyncio.run(_ws_protocol_smoke(ws_base, device_id))
        add("WebSocket 协议冒烟 register+heartbeat", ws_smoke.get("ok"), ws_smoke)
    except Exception as e:
        add("WebSocket 协议冒烟", False, str(e))

    # 6 hook probe
    probe = {}
    try:
        probe = _get(f"{base}/api/v3/hook/probe/{device_id}", timeout=60)
        ok = probe.get("supports_hook") is True or probe.get("hook_tests", {}).get("connect") == "ok"
        add("hook/probe Frida", ok, probe, "warn" if not ok else "pass")
    except URLError as e:
        add("hook/probe", False, str(e), "warn")

    # 7 hook execute 只读（设备+Frida 就绪时）
    for action, params in [
        ("get_profile", {}),
        ("get_contacts", {"limit": 3}),
    ]:
        try:
            r = _post(
                f"{base}/api/v3/hook/execute",
                {
                    "device_id": device_id,
                    "platform": "wechat",
                    "action": action,
                    "params": params,
                    "hook_only": True,
                },
                timeout=120,
            )
            data = r.get("data") or {}
            ok = data.get("success") is True or r.get("code") == 200 and data.get("success") is not False
            ch = r.get("channel_used") or data.get("_channel_used")
            add(f"hook/execute {action}", ok, {"channel": ch, "keys": list(data.keys())[:8]}, "warn" if not ok else "pass")
        except URLError as e:
            add(f"hook/execute {action}", False, str(e), "warn")

    # 8 ADB 安装与桌面入口
    adb_info = _check_adb_install(device_id)
    report["adb"] = adb_info
    if adb_info["expected_connected"]:
        add("Agent APK 已安装", adb_info.get("agent_pkg_installed") is True, AGENT_PKG)
        add(
            "桌面 Launcher 入口",
            bool(adb_info.get("launcher_activity")),
            adb_info.get("launcher_activity") or "无 LAUNCHER，需重新安装带桌面图标的 APK",
            "warn" if not adb_info.get("launcher_activity") else "pass",
        )
        add("微信已安装", adb_info.get("wechat_installed") is True, WECHAT_PKG)
    else:
        add("ADB 指定设备已连接", False, adb_info["adb_serials"], "fail")

    # 9 workbench
    try:
        wb = _get(f"{base}/api/v3/workbench/overview", timeout=20)
        add("workbench/overview", wb.get("code") == 200, {"keys": list((wb.get("data") or {}).keys())[:6]})
    except URLError as e:
        add("workbench/overview", False, str(e), "warn")

    return report


def main():
    ap = argparse.ArgumentParser(description="工作手机 WS+连接全量自检")
    ap.add_argument("-d", "--device", default=DEFAULT_DEVICE, help="指定设备序列号")
    ap.add_argument("--base", default=DEFAULT_BASE, help="SDK 地址")
    ap.add_argument("-o", "--output", default="", help="报告 JSON 路径")
    args = ap.parse_args()

    report = run_tests(args.base, args.device)
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    out = Path(args.output) if args.output else OUT_DIR / f"ws_full_test_{args.device}_{int(time.time())}.json"
    out.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

    s = report["summary"]
    print("=" * 60)
    print(f"工作手机全量自检 · 设备 {args.device}")
    print(f"通过 {s['pass']} · 警告 {s['warn']} · 失败 {s['fail']}")
    print("-" * 60)
    for c in report["checks"]:
        icon = "✅" if c["ok"] else ("⚠️" if c["level"] == "warn" else "❌")
        print(f"{icon} {c['name']}")
    print("-" * 60)
    print(f"报告: {out}")
    print("=" * 60)
    sys.exit(0 if s["fail"] == 0 else 1)


if __name__ == "__main__":
    main()