#!/usr/bin/env python3
"""安装经验证的 CPU LeRobot 环境；不修改当前解释器、上游源码或训练环境。"""

import argparse
import json
import os
import subprocess
import sys
import venv
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--venv", type=Path, default=ROOT / "build/venvs/lerobot")
    args = parser.parse_args()
    target = args.venv.resolve()
    if sys.version_info[:2] != (3, 12):
        parser.error("此兼容环境仅验证 Python 3.12，请使用 Python 3.12 执行安装脚本")
    if not target.is_relative_to(ROOT / "build") or target == Path(sys.prefix).resolve():
        parser.error("目标必须位于仓库 build/ 下，并且不能是当前环境")
    env = {**os.environ, "PYTHONNOUSERSITE": "1", "PYTHONPATH": ""}
    if not (target / "bin/python").exists():
        venv.EnvBuilder(with_pip=True).create(target)
    python = str(target / "bin/python")
    prefix = subprocess.check_output(
        [python, "-I", "-c", "import sys; print(sys.prefix)"], env=env, text=True
    ).strip()
    config = (target / "pyvenv.cfg").read_text().lower()
    if Path(prefix).resolve() != target or "include-system-site-packages = false" not in config:
        parser.error("目标不是独立 venv；拒绝对共享/其他 Python 环境执行 pip")

    def run(*arguments):
        subprocess.run([python, "-m", "pip", *arguments], env=env, cwd=ROOT, check=True)

    run(
        "install",
        "--index-url",
        "https://download.pytorch.org/whl/cpu",
        "torch==2.11.0+cpu",
        "torchvision==0.26.0+cpu",
    )
    constraints = "integrations/lerobot/requirements-cpu.txt"
    run("install", "-r", constraints)
    run(
        "install",
        "--no-build-isolation",
        "-c",
        constraints,
        "-e",
        "control_bridge",
        "-e",
        "integrations/lerobot",
    )
    run("check")
    # Record versions, not pip freeze's private VCS origins or credentials.
    result = subprocess.check_output(
        [python, "-m", "pip", "list", "--format=json"], env=env, text=True
    )
    report = ROOT / "build/lerobot-environment.json"
    report.write_text(json.dumps({"python": sys.version, "packages": json.loads(result)}, indent=2))
    print(f"完成：{python}\n版本记录：{report}\n请在此环境执行插件测试和示例。")


if __name__ == "__main__":
    main()