chuanqi/cq2_sdk/scripts/business/c2_wait_for_ready.py

"""C2: 等游戏就绪 — 纯 TCP poll, 不调 HTTP。

数据流:
    MySQL 拿 latest run (status="running")  →  拿 zone_id  →  拿 zone.host:port
        ↓
    wait_for_game_ready(host, port)  →  asyncio.open_connection 轮询
        ↓
    成功 → RunStore.update_status(run_id, "ready")
    超时 → RunStore.update_status(run_id, "failed", error="timeout")
        ↓
    DebugLogger 留痕

跑法:
    uv run python scripts/business/c2_wait_for_ready.py
    uv run python scripts/business/c2_wait_for_ready.py --user longwu988
    uv run python scripts/business/c2_wait_for_ready.py --run-id 5
    uv run python scripts/business/c2_wait_for_ready.py --timeout 90
    uv run python scripts/business/c2_wait_for_ready.py --check

前置依赖:
    1. A1 跑过 (user.token)
    2. C1 跑过 (runs 记录 + zone.host/port)
"""
from __future__ import annotations
import argparse
import asyncio
import os
import sys
import time
from pathlib import Path
from typing import Any

_HERE = Path(__file__).resolve().parent
_PROJECT = _HERE.parent.parent
sys.path.insert(0, str(_PROJECT / "src"))

from cq2_sdk.env_loader import load_env  # noqa: E402
load_env(_PROJECT / ".env")

from cq2_sdk.debug_log import DebugLogger  # noqa: E402
from cq2_sdk.launcher.navigate import (  # noqa: E402
    GameLaunchTimeout,
    wait_for_game_ready,
)
from cq2_sdk.store import RunStore, UserStore, ZoneStore, connection  # noqa: E402

DEFAULT_USERNAME = "longwu988"
DEFAULT_TIMEOUT = 60.0  # 跟 SDK 内置默认一致
DEFAULT_INTERVAL = 2.0


async def wait_and_persist(
    run_id: int | None,  # None = 自动选 latest
    user_id: int | None,  # run_id=None 时用 user 找 latest
    timeout: float,
    interval: float,
    log: DebugLogger,
) -> dict[str, Any]:
    """Run C2: 等 TCP 通 + 跟新 run 状态。

    Returns: {run_id, host, port, status, elapsed_s, attempts}
    """
    log.step("c2_wait_for_ready_start",
             run_id=run_id, user_id=user_id, timeout=timeout)

    # 1. 找 run
    with connection() as conn:
        runs = RunStore(conn)
        if run_id is None:
            # 自动: 取该 user 的 latest run
            if user_id is None:
                raise RuntimeError(
                    "需要 --run-id 或 --user 来定位 run"
                )
            recent = runs.list_recent(user_id, limit=1)
            if not recent:
                raise RuntimeError(
                    f"user_id={user_id} 在 MySQL 没 run 记录。先跑 C1 创一个。"
                )
            run = recent[0]
        else:
            run = runs.get(run_id)
            if run is None:
                raise RuntimeError(f"run_id={run_id} 不在 MySQL")
        run_db_id = run.id
        run_status = run.status
        zone_db_id = run.zone_id

        # 2. 拿 zone.host:port
        zone = ZoneStore(conn).get_by_id(zone_db_id)
        if zone is None:
            raise RuntimeError(
                f"run_id={run_db_id} 指向 zone_db_id={zone_db_id}, 但 zone 不存在"
            )
        host = zone.host
        port = zone.port
    log.info("loaded_run_and_zone",
             run_id=run_db_id, run_status=run_status,
             zone_db_id=zone_db_id, host=host, port=port)

    if not host or not port:
        log.error("zone_missing_host_or_port", run_id=run_db_id,
                  zone_db_id=zone_db_id, host=host, port=port,
                  note="C1 跑成功的话 zone.host/port 应该有值; 重跑 C1")
        raise RuntimeError(
            f"zone (id={zone_db_id}) 缺 host/port, 没法轮询。"
            f" 重新跑 C1 把它填上。"
        )

    if run_status not in ("running", "ready"):
        log.warn("run_status_unusual_for_c2",
                 run_id=run_db_id, run_status=run_status,
                 note="C2 通常跟 running → ready 转换; 当前 status 看起来已经走过了")

    # 3. TCP poll (用 SDK 自带 wait_for_game_ready)
    log.step("c2_polling_tcp", host=host, port=int(port),
             timeout_s=timeout, interval_s=interval)
    t0 = time.time()
    attempts = 0
    last_err: str | None = None
    try:
        await wait_for_game_ready(
            host=host, port=int(port),
            timeout=timeout, interval=interval,
        )
    except GameLaunchTimeout as e:
        elapsed = time.time() - t0
        log.error("tcp_poll_timeout",
                  run_id=run_db_id, host=host, port=int(port),
                  elapsed_s=round(elapsed, 1), last_err=last_err)
        # 跟新 run 状态为 failed
        with connection() as conn:
            RunStore(conn).update_status(
                run_db_id, "failed", current_task=None,
            )
        raise RuntimeError(
            f"游戏没在 {timeout}s 内起来: {e}"
        )
    elapsed = time.time() - t0
    log.info("tcp_port_reachable",
             run_id=run_db_id, host=host, port=int(port),
             elapsed_s=round(elapsed, 1), attempts=attempts)

    # 4. 跟新 run 状态为 ready
    log.step("c2_persist_run_status", run_id=run_db_id, new_status="ready")
    with connection() as conn:
        updated = RunStore(conn).update_status(run_db_id, "ready")
    log.info("run_status_updated",
             run_id=updated.id, new_status=updated.status)

    return {
        "run_id": run_db_id,
        "host": host,
        "port": int(port),
        "status": "ready",
        "elapsed_s": round(elapsed, 1),
    }


def check_config(user: str | None, run_id: int | None) -> bool:
    print("[check] 配置检查:")
    if run_id is not None:
        print(f"  run_id  = {run_id}  [OK, 显式指定]")
    elif user:
        print(f"  user    = {user!r}  [OK, 会取该 user 的 latest run]")
    else:
        print("  user    = (未设置)  [FAIL: --user 或 --run-id 至少一个]")
    print("  (实际跑还需要 C1 跑过, MySQL 有 runs 记录 + zone.host/port)")
    return run_id is not None or user is not None


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="C2: 轮询 game server host:port 等游戏就绪, 跟新 run status",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    p.add_argument(
        "--user",
        default=os.environ.get("CQ2_ACCOUNT", DEFAULT_USERNAME),
        help="自动选 latest run 时, 找哪个 user 的 run (默认 .env CQ2_ACCOUNT)",
    )
    p.add_argument(
        "--run-id",
        type=int,
        default=None,
        help="显式指定 run_id, 跳过自动选",
    )
    p.add_argument(
        "--timeout",
        type=float,
        default=DEFAULT_TIMEOUT,
        help=f"TCP poll 超时秒数 (默认 {DEFAULT_TIMEOUT})",
    )
    p.add_argument(
        "--interval",
        type=float,
        default=DEFAULT_INTERVAL,
        help=f"TCP poll 间隔秒数 (默认 {DEFAULT_INTERVAL})",
    )
    p.add_argument(
        "--check",
        action="store_true",
        help="只检查 .env / CLI 配置, 不查 DB 不连 TCP",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()

    if args.check:
        ok = check_config(args.user, args.run_id)
        print(f"\n  result: {'READY' if ok else 'NOT READY'}")
        return 0 if ok else 1

    if args.run_id is None and not args.user:
        check_config(args.user, args.run_id)
        return 1

    with DebugLogger(
        extra={"task": "c2", "user": args.user, "run_id": args.run_id,
               "timeout": args.timeout}
    ) as log:
        log.info("c2_start", user=args.user, run_id=args.run_id,
                 timeout=args.timeout)
        try:
            # user_id 反查 (明确报错, 比传 None 进去好)
            user_id = None
            if args.user and args.run_id is None:
                with connection() as conn:
                    u = UserStore(conn).get_by_username(args.user)
                    if u is None:
                        raise RuntimeError(
                            f"user {args.user!r} 不在 MySQL。先跑 A1 登录。"
                        )
                    user_id = u.id
            result = asyncio.run(wait_and_persist(
                args.run_id, user_id, args.timeout, args.interval, log,
            ))
        except Exception as e:
            log.error(e, step="c2_unexpected")
            print(f"\n[FAIL] {type(e).__name__}: {e}")
            print(f"  log file: {log.path}")
            return 1

    print("\n[OK] C2 游戏就绪:")
    for k, v in result.items():
        print(f"  {k:<10} = {v}")
    print(f"  log file = {log.path}")
    print("  下一步: 进入 D1-D13 (13 个新号流程任务)")
    return 0


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