#!/usr/bin/env python3
"""Launch this release's bundled vllm-fast-text runtime, evaluate the pinned test set, stop server."""
import argparse
from datetime import datetime, timezone
import hashlib
import json
import os
from pathlib import Path
import socket
import subprocess
import sys
import time
import urllib.request


def verify_release(model):
    for line in (model / "SHA256SUMS").read_text().splitlines():
        digest, name = line.split("  ", 1)
        rel = Path(name)
        if rel.is_absolute() or ".." in rel.parts:
            raise ValueError("非法 SHA256SUMS 路径: " + name)
        h = hashlib.sha256()
        with (model / rel).open("rb") as stream:
            for block in iter(lambda: stream.read(8 * 1024 * 1024), b""):
                h.update(block)
        if h.hexdigest() != digest:
            raise ValueError("发布文件已改变，请重新打包和评测: " + name)


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model-dir", type=Path, required=True)
    p.add_argument("--output-dir", type=Path, required=True,
                   help="必须是新目录，避免把不同模型的预测混用")
    p.add_argument("--port", type=int, default=8080)
    p.add_argument("--device", default="cuda")
    p.add_argument("--startup-timeout", type=float, default=900)
    a = p.parse_args()
    if sys.version_info < (3, 11):
        raise SystemExit("NeoHorse 推理需要 Python >= 3.11；请使用推理环境中的 python 执行。")
    import importlib.metadata
    model, out = a.model_dir.resolve(), a.output_dir.resolve()
    if out.exists():
        raise SystemExit("评测输出目录已存在，请选择新目录。")
    verify_release(model)
    # Refuse an occupied port so a different running model cannot be scored by mistake.
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", a.port))
    out.mkdir(parents=True)
    try:
        gpu = subprocess.check_output(["nvidia-smi", "--query-gpu=name,memory.total,driver_version",
                                       "--format=csv,noheader"], text=True).strip()
    except (OSError, subprocess.CalledProcessError):
        gpu = "Unavailable (CPU or nvidia-smi unavailable)"
    versions = {}
    for name in ("torch", "transformers", "safetensors", "flash-linear-attention"):
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = "not installed"
    meta = json.loads((model / "model_manifest.json").read_text())
    metadata = {"started_at": datetime.now(timezone.utc).isoformat(), "device": a.device,
                "gpu_inventory": gpu, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES", "not set"),
                "python": sys.version, "versions": versions, "temperature": meta["temperature"],
                "backbone_dtype": "bfloat16", "head_dtype": "float32",
                "request_shape": "all five questions per case, serial HTTP requests, no warmup exclusion",
                "release_sha256sums_sha256": hashlib.sha256((model / "SHA256SUMS").read_bytes()).hexdigest()}
    (out / "run_metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
    env = os.environ.copy()
    # Bind to loopback only; do not inherit API authentication into a private local run.
    env.pop("NEOHORSE_API_KEY", None)
    # Launch the bundled vllm-fast-text adapter directly; it manages its own
    # PYTHONPATH for spawned vLLM workers and has no --device option.
    command = [sys.executable, str(model / "vllm-fast-text" / "server.py"),
               "--bundle", str(model), "--host", "127.0.0.1", "--port", str(a.port)]
    with (out / "server.log").open("w") as log:
        server = subprocess.Popen(command, env=env, stdout=log, stderr=subprocess.STDOUT)
        try:
            deadline = time.monotonic() + a.startup_timeout
            while True:
                if server.poll() is not None:
                    raise RuntimeError("模型启动失败，请查看 " + str(out / "server.log"))
                try:
                    with urllib.request.urlopen(f"http://127.0.0.1:{a.port}/healthz", timeout=2) as response:
                        if response.status == 200:
                            break
                except OSError:
                    pass
                if time.monotonic() > deadline:
                    raise RuntimeError("等待模型启动超时，请查看 server.log")
                time.sleep(1)
            subprocess.run([sys.executable, str(model / "evaluation_tools/evaluate.py"),
                            "--endpoint", f"http://127.0.0.1:{a.port}/v1/systemone",
                            "--output-dir", str(out)], check=True)
        finally:
            server.terminate()
            try:
                server.wait(timeout=15)
            except subprocess.TimeoutExpired:
                server.kill()
                server.wait()


if __name__ == "__main__":
    main()