"""Rust 进程插件端到端协议验证（多语言演进规划阶段 1 验收）

前置条件:
  1. 本地 Zenoh Router 已运行（scripts/zenoh_router.py）
  2. Rust 参考插件已编译（cargo build --examples 或 workers/rsdemo）

验证内容（docs/plugin-wire-protocol.md 全协议面）:
  1. 路由描述符发布（keyexpr 编码、字段、schema_version/type）
  2. liveliness 服务发现（worker/rsdemo/*）
  3. 普通 RPC（GET 信封 / POST user_info 透传）
  4. 非信封载荷兼容（method 默认 GET）
  5. 流式多回复（逐块）+ 流中异常 {"error": ...} 收尾块
  6. 非流式异常回复格式（result/operation/error/worker）
  7. SIGTERM 优雅退出 → 描述符 DELETE 摘除
  8. 重启 → 10s 内重发布自愈

用法（项目根目录）:
  .venv/bin/python utils/worker-rs/scripts/verify_e2e.py [--type rsdemo] [--count 1] [--bin <路径>]
"""
import argparse
import orjson
import os
import signal
import subprocess
import sys
import time

import zenoh

PASS = 0
FAIL = 0

DEFAULT_BIN = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "target", "debug", "examples", "demo")
)
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
ROUTES_PREFIX = "chongming/routes"


def report(name, cond, detail=""):
    global PASS, FAIL
    if cond:
        PASS += 1
        print(f"✅ {name}" + (f" — {detail}" if detail else ""))
    else:
        FAIL += 1
        print(f"❌ {name}" + (f" — {detail}" if detail else ""))


def start_worker(worker_type, count=1, bin_path=None):
    """启动 Rust 插件进程（CWD=项目根，读取全局 config.toml）"""
    binary = bin_path or DEFAULT_BIN
    proc = subprocess.Popen(
        [binary, "--type", worker_type, "--count", str(count)],
        cwd=ROOT,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    # 等待 Queryable 就绪（日志出现"进程就绪"）
    deadline = time.monotonic() + 20
    while time.monotonic() < deadline:
        if proc.poll() is not None:
            out = proc.stdout.read().decode() if proc.stdout else ""
            raise RuntimeError(f"Worker 启动失败，退出码 {proc.returncode}:\n{out}")
        time.sleep(0.2)
        line = proc.stdout.readline() if proc.stdout else b""
        if "等待退出信号" in line.decode("utf-8", "ignore"):
            return proc
    raise RuntimeError("Worker 20s 内未就绪")


def wait_worker_exit(proc, timeout=10):
    try:
        proc.wait(timeout=timeout)
        return True
    except subprocess.TimeoutExpired:
        return False


def make_subscriber(session, worker_type):
    """建立常驻订阅，持续记录 PUT 描述符与 DELETE 事件（不重放历史值，
    调用方须在事件发生前建立订阅）"""
    samples = {}

    def cb(sample):
        key = str(sample.key_expr)
        if sample.kind == zenoh.SampleKind.DELETE:
            samples[key] = None  # 摘除标记
        else:
            samples[key] = orjson.loads(bytes(sample.payload))

    sub = session.declare_subscriber(f"{ROUTES_PREFIX}/{worker_type}/**", cb)
    return sub, samples


def rpc(session, worker_expr, subject, payload, timeout=10.0):
    """按信封协议发起 Zenoh Query，返回所有回复块 JSON 列表"""
    full_key = f"{worker_expr}/{subject}"
    replies = session.get(
        full_key,
        payload=orjson.dumps(payload),
        consolidation=zenoh.ConsolidationMode.NONE,
        timeout=timeout,
    )
    result = []
    for reply in replies:
        if reply.ok:
            result.append(orjson.loads(bytes(reply.ok.payload)))
        else:
            result.append({"__zenoh_err__": str(reply.err.payload)})
    return result


def get_worker_exprs(session, worker_type):
    """liveliness 发现 worker/worker_type/* 副本列表"""
    exprs = []
    for reply in session.liveliness().get(f"worker/{worker_type}/*"):
        if reply.ok:
            exprs.append(str(reply.ok.key_expr))
    return exprs


def main():
    global PASS, FAIL
    parser = argparse.ArgumentParser()
    parser.add_argument("--type", default="rsdemo")
    parser.add_argument("--count", type=int, default=1)
    parser.add_argument("--bin", default=None, help="插件二进制路径（默认 SDK 示例 demo）")
    args = parser.parse_args()
    wt = args.type

    session = zenoh.open(zenoh.Config.from_file("config/zenoh_config.json5"))
    print(f"=== Rust 进程插件协议验证 (type={wt}, bin={args.bin or DEFAULT_BIN}) ===")

    # ========== 1. 路由描述符发布（先订阅再启动，zenoh 不重放历史值） ==========
    sub1, samples1 = make_subscriber(session, wt)
    proc = start_worker(wt, args.count, args.bin)
    time.sleep(1.5)

    expect_keys = {
        "rsdemo": {
            "chongming/routes/rsdemo/rsdemo_hello__GET",
            "chongming/routes/rsdemo/rsdemo_echo__POST",
            "chongming/routes/rsdemo/rsdemo_stream__POST",
            "chongming/routes/rsdemo/rsdemo_error__GET",
        },
        "demo": {
            "chongming/routes/demo/demo_hello__GET",
            "chongming/routes/demo/demo_storage_upload_bytes__POST",
            "chongming/routes/demo/demo_storage_download_bytes__GET",
            "chongming/routes/demo/demo_storage_presign_upload__POST",
            "chongming/routes/demo/demo_storage_presign_download__GET",
        },
    }[wt]
    missing = expect_keys - set(samples1.keys())
    report(
        "路由描述符 keyexpr 编码正确",
        not missing,
        f"缺失: {missing or '无'}",
    )

    if "chongming/routes/rsdemo/rsdemo_hello__GET" in samples1:
        d = samples1["chongming/routes/rsdemo/rsdemo_hello__GET"]
        report("描述符 schema_version=1", d.get("schema_version") == 1, f"实际: {d.get('schema_version')}")
        report("描述符 type=插件名", d.get("type") == wt, f"实际: {d.get('type')}")
        report(
            "描述符字段完整（subject/method/path/params/response_model）",
            d.get("subject") == "rsdemo.hello"
            and d.get("method") == "GET"
            and d.get("path") == "/rsdemo/hello"
            and d.get("params") == ["name: str"]
            and isinstance(d.get("response_model"), dict),
        )
        report(
            "defaults 合并（timeout=30.0/auth_required=false）",
            d.get("timeout") == 30.0 and d.get("auth_required") is False,
        )
    sub1.undeclare()

    # ========== 2. liveliness 服务发现 ==========
    exprs = get_worker_exprs(session, wt)
    report(
        "liveliness 服务发现副本",
        len(exprs) == args.count,
        f"实际: {exprs}",
    )
    worker_expr = exprs[0] if exprs else f"worker/{wt}/missing"

    # ========== 3. 普通 RPC：GET 信封 ==========
    replies = rpc(session, worker_expr, "rsdemo.hello", {"meta": {"method": "GET"}, "data": {"name": "rust"}})
    ok = (
        len(replies) == 1
        and replies[0].get("greeting") == "Hello, rust!"
        and replies[0].get("worker_id") == worker_expr.split("/")[-1]
        and isinstance(replies[0].get("timestamp"), (int, float))
    )
    report("RPC GET（信封分派 + worker_id 注入）", ok, f"实际: {replies}")

    # ========== 4. 非信封载荷兼容（method 默认 GET） ==========
    replies = rpc(session, worker_expr, "rsdemo.hello", {"name": "legacy"})
    report(
        "非信封载荷兼容（按 GET 处理）",
        len(replies) == 1 and replies[0].get("greeting") == "Hello, legacy!",
        f"实际: {replies}",
    )

    # ========== 5. POST + user_info 透传 ==========
    replies = rpc(
        session, worker_expr, "rsdemo.echo",
        {"meta": {"method": "POST", "user_info": {"username": "alice"}}, "data": {"msg": "hi"}},
    )
    ok = (
        len(replies) == 1
        and replies[0].get("echo") == "hi"
        and replies[0].get("user_info") == {"username": "alice"}
    )
    report("RPC POST（user_info 信封透传）", ok, f"实际: {replies}")

    # ========== 6. 流式多回复 ==========
    replies = rpc(
        session, worker_expr, "rsdemo.stream",
        {"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 5}},
        timeout=15.0,
    )
    report(
        "流式多回复（5 块逐块推送）",
        len(replies) == 5
        and all(r.get("index") == i for i, r in enumerate(replies))
        and all("error" not in r for r in replies),
        f"块数: {len(replies)}",
    )

    # 流中异常 → {"error": ...} 收尾块（第 26 块抛错，共 26 块）
    replies = rpc(
        session, worker_expr, "rsdemo.stream",
        {"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 60}},
        timeout=20.0,
    )
    ok = (
        len(replies) == 26
        and all("error" not in r for r in replies[:25])
        and isinstance(replies[25].get("error"), str)
        and "模拟流中异常" in replies[25]["error"]
    )
    report("流中异常 → error 收尾块终止", ok, f"块数: {len(replies)}, 末块: {replies[-1] if replies else None}")

    # ========== 7. 非流式异常回复格式 ==========
    replies = rpc(session, worker_expr, "rsdemo.error", {"meta": {"method": "GET"}, "data": {}})
    ok = (
        len(replies) == 1
        and replies[0].get("result") == 0.0
        and replies[0].get("operation") == "error"
        and "模拟业务异常" in replies[0].get("error", "")
        and replies[0].get("worker") == worker_expr.split("/")[-1]
    )
    report("非流式异常回复格式（result/operation/error/worker）", ok, f"实际: {replies}")

    # ========== 8. SIGTERM 优雅退出 → 描述符摘除（先订阅再 kill） ==========
    sub2, samples2 = make_subscriber(session, wt)
    os.kill(proc.pid, signal.SIGTERM)
    exited = wait_worker_exit(proc, timeout=10)
    report("SIGTERM 优雅退出", exited, f"退出码: {proc.returncode if exited else '超时'}")
    time.sleep(1.5)
    removed = [k for k, v in samples2.items() if v is None]
    report(
        "退出时 DELETE 摘除全部描述符",
        len(removed) == len(expect_keys),
        f"摘除: {sorted(removed)}",
    )
    sub2.undeclare()

    # ========== 9. 重启 → 10s 内重发布自愈 ==========
    sub3, samples3 = make_subscriber(session, wt)
    proc = start_worker(wt, args.count, args.bin)
    time.sleep(11.5)  # 覆盖一个完整重发布周期
    ok = all(k in samples3 and samples3[k] is not None for k in expect_keys)
    report(
        "重启后周期重发布自愈（≤10s）",
        ok,
        f"恢复: {sorted(k for k in samples3 if samples3[k] is not None)}",
    )
    sub3.undeclare()

    # 清理
    os.kill(proc.pid, signal.SIGTERM)
    wait_worker_exit(proc, timeout=10)
    session.close()

    print(f"\n=== 结果: {PASS} 通过, {FAIL} 失败 ===")
    sys.exit(1 if FAIL else 0)


if __name__ == "__main__":
    main()