# -*- coding: utf-8 -*-
"""端到端验证中转服务器。

从 Windows 走公网隧道，模拟「被控端」和「手表」两条连接，
验证：注册 / 接入 / 双向收发 / 单设备单配对码 / 延迟。
"""
from __future__ import annotations

import json
import sys
import threading
import time
import urllib.error
import urllib.request


# Cloudflare 的 Browser Integrity Check 会拦 Python 默认 UA
# （返回 403 + error code 1010），所以请求都要带上这个。
_UA = {"User-Agent": "LanRemote/1.4.0"}

URL = sys.argv[1] if len(sys.argv) > 1 else ""
ADMIN = sys.argv[2] if len(sys.argv) > 2 else ""
# 每次跑用不同的配对码：中转服务器会持久化「设备->配对码」绑定，
# 固定码第二次跑就会撞上一轮的绑定（这正是单设备单码制在起作用）
CODE = str(int(time.time()) % 1000000).zfill(6)
CODE_NEW = str((int(time.time()) + 7) % 1000000).zfill(6)
BAD_CODE = "000000" if CODE != "000000" else "111111"
if not URL:
    print("用法: python _build/test_relay_e2e.py <隧道地址>")
    sys.exit(1)

passed = 0
failed = 0


def check(name, ok, detail=""):
    global passed, failed
    if ok:
        passed += 1
        print("  [PASS] " + name)
    else:
        failed += 1
        print("  [FAIL] " + name + "  " + str(detail))


def post(path, body: bytes = b"", headers: dict | None = None,
         timeout: float = 40.0) -> tuple:
    req = urllib.request.Request(URL + path, data=body, method="POST", headers=_UA)
    req.add_header("Content-Type", "application/octet-stream")
    for key, value in (headers or {}).items():
        req.add_header(key, value)
    started = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            data = resp.read()
            return resp.status, data, time.time() - started
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read(), time.time() - started
    except Exception as exc:                              # noqa: BLE001
        return 0, str(exc).encode(), time.time() - started


def post_json(path, payload: dict, headers: dict | None = None) -> tuple:
    return post(path, json.dumps(payload, ensure_ascii=False).encode("utf-8"),
                headers)


def get(path, timeout: float = 30.0) -> tuple:
    try:
        _req = urllib.request.Request(URL + path, headers=_UA)
        with urllib.request.urlopen(_req, timeout=timeout) as resp:
            return resp.status, resp.read(), 0.0
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read(), 0.0
    except Exception as exc:                              # noqa: BLE001
        return 0, str(exc).encode(), 0.0


def main() -> int:
    print(f"隧道地址: {URL}")
    print()

    # ---------------------------------------------------------------
    print("一、健康检查")
    status, body, _ = get("/v1/health")
    health = {}
    try:
        health = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("GET /v1/health 返回 200", status == 200, status)
    check("健康检查 ok", health.get("ok") is True, str(health)[:100])
    print(f"       版本 {health.get('version')}  内存 {health.get('memoryKB', 0) // 1024} MB")

    # ---------------------------------------------------------------
    print()
    print("二、被控端注册")
    device_id = "win-test-" + str(int(time.time()))
    status, body, cost = post_json("/v1/agent/register", {
        "deviceId": device_id, "code": CODE,
        "name": "测试电脑", "version": "1.4.0"})
    reg = {}
    try:
        reg = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("注册成功", reg.get("ok") is True, str(reg)[:140])
    agent_token = reg.get("token", "")
    check("拿到被控端令牌", len(agent_token) == 48, len(agent_token))
    print(f"       注册耗时 {cost * 1000:.0f} ms")

    # ---------------------------------------------------------------
    print()
    print("三、单设备单配对码制")
    status, body, _ = post_json("/v1/agent/register", {
        "deviceId": device_id, "code": CODE_NEW,
        "name": "测试电脑", "version": "1.4.0"})
    again = {}
    try:
        again = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("同一台设备报新码时被改回旧码",
          again.get("code") == CODE,
          f"返回 {again.get('code')}  note={again.get('note')}")
    check("并且给出了说明", "单设备单配对码" in str(again.get("note", "")),
          str(again.get("note"))[:80])
    agent_token = again.get("token", agent_token)         # 令牌会刷新

    # 另一台设备用同一个码 -> 应该被拒
    status, body, _ = post_json("/v1/agent/register", {
        "deviceId": "another-device-xyz", "code": CODE,
        "name": "另一台电脑", "version": "1.4.0"})
    clash = {}
    try:
        clash = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("同一个码不能被两台设备用", clash.get("ok") is False,
          str(clash)[:120])

    # ---------------------------------------------------------------
    print()
    print("四、手表接入")
    status, body, _ = post_json("/v1/client/attach", {
        "code": CODE, "name": "测试手表"})
    att = {}
    try:
        att = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("接入成功", att.get("ok") is True, str(att)[:140])
    session = att.get("session", "")
    check("拿到会话令牌", len(session) == 48, len(session))
    check("看到了设备名", att.get("deviceName") == "测试电脑",
          att.get("deviceName"))

    # 错的码
    status, body, _ = post_json("/v1/client/attach", {
        "code": BAD_CODE, "name": "坏手表"})
    bad = {}
    try:
        bad = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    check("不存在的配对码被拒", bad.get("ok") is False, str(bad)[:100])

    # ---------------------------------------------------------------
    print()
    print("五、被控端 -> 手表 传字节")
    payload = "来自电脑的消息-" + str(time.time()) + "\n"
    status, body, cost = post("/v1/agent/send", payload.encode("utf-8"),
                              {"X-Token": agent_token})
    check("被控端发送成功", status == 200, status)

    started = time.time()
    status, got, _ = post("/v1/client/poll", b"", {"X-Session": session},
                          timeout=40)
    elapsed = time.time() - started
    check("手表收到了", got.decode("utf-8", "replace") == payload,
          repr(got[:80]))
    print(f"       端到端延迟 {elapsed * 1000:.0f} ms")

    # ---------------------------------------------------------------
    print()
    print("六、手表 -> 被控端 传字节")
    reply = "来自手表的消息-" + str(time.time()) + "\n"
    status, body, _ = post("/v1/client/send", reply.encode("utf-8"),
                           {"X-Session": session})
    check("手表发送成功", status == 200, status)

    started = time.time()
    status, got, _ = post("/v1/agent/poll", b"", {"X-Token": agent_token},
                          timeout=40)
    elapsed = time.time() - started
    check("被控端收到了", got.decode("utf-8", "replace") == reply,
          repr(got[:80]))
    print(f"       端到端延迟 {elapsed * 1000:.0f} ms")

    # ---------------------------------------------------------------
    print()
    print("七、长轮询：没数据时应该挂住而不是立刻返回")
    started = time.time()
    status, got, _ = post("/v1/client/poll", b"", {"X-Session": session},
                          timeout=40)
    elapsed = time.time() - started
    check("没数据时挂够时间再返回", 15 < elapsed < 32,
          f"挂了 {elapsed:.1f} 秒")
    check("挂完返回空", got == b"", repr(got[:40]))

    # ---------------------------------------------------------------
    print()
    print("八、连续多段数据不丢不乱序")
    chunks = [f"片段{i}-{'x' * 50}\n" for i in range(8)]
    for piece in chunks:
        post("/v1/agent/send", piece.encode("utf-8"), {"X-Token": agent_token})
    status, got, _ = post("/v1/client/poll", b"", {"X-Session": session},
                          timeout=40)
    check("8 段全部收到且顺序正确",
          got.decode("utf-8", "replace") == "".join(chunks),
          f"收到 {len(got)} 字节，期望 {sum(len(c) for c in chunks)}")

    # ---------------------------------------------------------------
    print()
    print("九、大块数据（模拟一个 256KB 分块）")
    big = ("B" * (256 * 1024)) + "\n"
    status, body, cost = post("/v1/agent/send", big.encode("utf-8"),
                              {"X-Token": agent_token}, timeout=60)
    check("256KB 发送成功", status == 200, status)
    status, got, cost2 = post("/v1/client/poll", b"", {"X-Session": session},
                              timeout=60)
    check("256KB 完整收到", len(got) == len(big.encode("utf-8")),
          f"收到 {len(got)}，期望 {len(big.encode('utf-8'))}")
    print(f"       上行 {cost * 1000:.0f} ms / 下行 {cost2 * 1000:.0f} ms")

    # ---------------------------------------------------------------
    print()
    print("十、令牌校验")
    status, body, _ = post("/v1/agent/poll", b"", {"X-Token": "fake-token-xyz"},
                           timeout=20)
    check("假被控端令牌被拒", status == 401,
          f"状态 {status}  响应 {body[:80]}")
    status, body, _ = post("/v1/client/poll", b"", {"X-Session": "fake-session-xyz"},
                           timeout=20)
    check("假会话被拒", status == 401,
          f"状态 {status}  响应 {body[:80]}")

    # ---------------------------------------------------------------
    print()
    print("十一、管理页与统计")
    status, body, _ = get("/v1/health")
    snap = {}
    try:
        snap = json.loads(body)
    except Exception:                                     # noqa: BLE001
        pass
    names = [d.get("name") for d in snap.get("devices", [])]
    check("统计里能看到刚注册的设备", "测试电脑" in names, str(names))
    check("统计里有会话数", snap.get("sessions", 0) >= 1, snap.get("sessions"))
    codes = [d.get("code") for d in snap.get("devices", [])]
    check("统计里能看到配对码（便于管理）", CODE in codes, str(codes))

    # 没登录 -> 给登录页（200），但不给数据
    status, body, _ = get("/")
    check("管理页没登录时给登录页", status == 200, f"状态 {status}")
    text = body.decode("utf-8", "replace")
    check("登录页确实是登录页", "管理密码" in text or "登录" in text,
          text[:80])
    check("登录前不泄露配对码", CODE.encode() not in body,
          "居然把配对码给出去了")

    if ADMIN:
        import urllib.parse
        status, body, _ = get("/?key=" + urllib.parse.quote(ADMIN))
        check("带钥匙能打开管理页",
              status == 200 and b"<html" in body.lower(), f"状态 {status}")
        check("管理页能打开", b"<html" in body.lower(), body[:60])
        status, body, _ = get("/v1/admin/devices")
        check("管理接口没登录也被拒", status == 401, f"状态 {status}")
        req = urllib.request.Request(URL + "/v1/admin/devices", headers=_UA)
        req.add_header("X-Admin", ADMIN)
        try:
            with urllib.request.urlopen(req, timeout=20) as resp:
                admin_snap = json.loads(resp.read())
            check("管理接口带钥匙能拿到设备列表",
                  admin_snap.get("ok") is True, str(admin_snap)[:80])
        except Exception as exc:                          # noqa: BLE001
            check("管理接口带钥匙能拿到设备列表", False, str(exc))

    # 清理：把这次用的配对码从服务器抹掉，免得越积越多
    if ADMIN:
        try:
            req = urllib.request.Request(
                URL + "/v1/admin/forget",
                data=json.dumps({"code": CODE}).encode("utf-8"),
                method="POST", headers=_UA)
            req.add_header("X-Admin", ADMIN)
            req.add_header("Content-Type", "application/json")
            with urllib.request.urlopen(req, timeout=20) as resp:
                print()
                print("  已清理测试配对码: " + resp.read().decode("utf-8", "replace"))
        except Exception as exc:
            print(f"  清理失败（不影响结果）: {exc}")

    print()
    print("=" * 66)
    print(f"通过 {passed} 项，失败 {failed} 项")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())