# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License").
"""Run 10 short, simulation-only LIBERO smoke tests against a GR00T policy server.

What this does (and does NOT do):
  * It does NOT use any physical robot hardware -- LIBERO simulation only.
  * It does NOT rewrite the official evaluation path. Each scenario is run
    through ``gr00t.eval.rollout_policy.run_rollout_gymnasium_policy`` via the
    thin ``_libero_rollout_worker.py`` (which only adds action recording and a
    couple of tiny, optional simulation-only perturbations).
  * It does NOT fake successful rollouts. If the model checkpoint is missing,
    or the server is unreachable, or the LIBERO sim env is not installed, it
    fails with a clear, actionable error message.

Typical usage (two terminals):

  Terminal 1 - start the GR00T inference server::

      uv run python gr00t/eval/run_gr00t_server.py \
          --model-path checkpoints/GR00T-N1.7-LIBERO/libero_10 \
          --embodiment-tag LIBERO_PANDA \
          --use-sim-policy-wrapper

  Terminal 2 - run the smoke tests::

      uv run python examples/LIBERO/smoke_tests/run_10_smoke_tests.py \
          --model-path checkpoints/GR00T-N1.7-LIBERO/libero_10 \
          --manifest examples/LIBERO/smoke_tests/scenarios_10.yaml \
          --output-dir outputs/libero_smoke_tests \
          --max-episode-steps 50 --save-video --render

Pass ``--start-server`` to have this script launch (and later kill) the server
itself, or ``--dry-run`` to validate the manifest / paths without running sims.
"""

from __future__ import annotations

import argparse
import datetime as _dt
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import time
import traceback

REPO_ROOT = Path(__file__).resolve().parents[3]
WORKER = Path(__file__).resolve().parent / "_libero_rollout_worker.py"
DEFAULT_MANIFEST = Path(__file__).resolve().parent / "scenarios_10.yaml"
RESULT_PREFIX = "SMOKE_RESULT_JSON:"

# Path to the dedicated LIBERO uv venv created by setup_libero.sh.
LIBERO_VENV_PYTHON = (
    REPO_ROOT / "gr00t" / "eval" / "sim" / "LIBERO" / "libero_uv" / ".venv" / "bin" / "python"
)

SERVER_CMD_TEMPLATE = [
    "uv", "run", "python", "gr00t/eval/run_gr00t_server.py",
    "--model-path", "{model_path}",
    "--embodiment-tag", "LIBERO_PANDA",
    "--use-sim-policy-wrapper",
]


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _load_manifest(path: Path) -> list[dict]:
    try:
        import yaml
    except ImportError as e:  # pragma: no cover
        raise SystemExit(
            "PyYAML is required to read the scenario manifest. Install it with "
            "`uv pip install pyyaml` (it is a transitive dependency of the gr00t "
            "package, so this normally just works inside the project venv)."
        ) from e
    if not path.exists():
        raise SystemExit(f"Manifest not found: {path}")
    with open(path) as f:
        data = yaml.safe_load(f)
    scenarios = data.get("scenarios") if isinstance(data, dict) else data
    if not isinstance(scenarios, list) or not scenarios:
        raise SystemExit(f"Manifest {path} does not contain a non-empty 'scenarios' list.")
    return scenarios


def _check_model_path(model_path: Path) -> None:
    if model_path.exists():
        # Must look like a real checkpoint dir.
        has_cfg = (model_path / "config.json").exists()
        has_weights = any(model_path.glob("*.safetensors")) or (model_path / "pytorch_model.bin").exists()
        if has_cfg and has_weights:
            return
        missing = []
        if not has_cfg:
            missing.append("config.json")
        if not has_weights:
            missing.append("model-*.safetensors / model.safetensors.index.json")
        raise SystemExit(
            f"Model path {model_path} exists but is missing: {', '.join(missing)}.\n"
            "Re-download the checkpoint:\n\n" + _download_hint()
        )
    raise SystemExit(
        f"Model checkpoint not found at: {model_path}\n\n"
        "Download it first (HuggingFace does not support nested repo paths directly):\n\n"
        + _download_hint()
    )


def _download_hint() -> str:
    return (
        "  uv run hf download nvidia/GR00T-N1.7-LIBERO \\\n"
        "    --include \"libero_10/config.json\" \\\n"
        "              \"libero_10/embodiment_id.json\" \\\n"
        "              \"libero_10/model-*.safetensors\" \\\n"
        "              \"libero_10/model.safetensors.index.json\" \\\n"
        "              \"libero_10/processor_config.json\" \\\n"
        "              \"libero_10/statistics.json\" \\\n"
        "    --local-dir checkpoints/GR00T-N1.7-LIBERO\n"
        "  # (also fetch libero_10/config.json explicitly if the glob above skips it)\n"
    )


def _server_command(model_path: Path) -> list[str]:
    return [tok.format(model_path=str(model_path)) for tok in SERVER_CMD_TEMPLATE]


# The GR00T-N1.7 backbone (nvidia/Cosmos-Reason2-2B, a Qwen3-VL model) is a
# *gated* HuggingFace repo. Loading any GR00T-N1.7 checkpoint pulls that base
# repo's config/processor, so the server will fail to start without HF auth +
# granted access. We surface a clear hint when we detect this.
GATED_BACKBONE_HINT = (
    "The GR00T-N1.7 backbone 'nvidia/Cosmos-Reason2-2B' is a GATED HuggingFace repo.\n"
    "To start the server you must:\n"
    "  1. Request access at https://huggingface.co/nvidia/Cosmos-Reason2-2B (one click, usually instant).\n"
    "  2. Authenticate, e.g.  export HF_TOKEN=hf_xxx   (or: uv run hf auth login)\n"
    "  3. Re-run the server / smoke tests.\n"
)


def _scan_log_for_gated_repo(log_path: Path) -> bool:
    try:
        text = log_path.read_text(errors="replace")
    except OSError:
        return False
    return ("gated repo" in text) or ("Cosmos-Reason2-2B is restricted" in text) or (
        "Access to model nvidia/Cosmos-Reason2-2B" in text
    )


def _ping_server(host: str, port: int, timeout_ms: int = 3000) -> bool:
    """Return True if a GR00T policy server answers on host:port.

    First does a cheap TCP connect (so we don't drag in the heavy torch/gr00t
    import stack when nothing is listening); only if *something* is listening do
    we import ``PolicyClient`` and validate the msgpack ``ping`` endpoint.
    """
    import socket

    try:
        with socket.create_connection((host, port), timeout=timeout_ms / 1000.0):
            pass
    except OSError:
        return False
    # Something is listening -- confirm it actually speaks the GR00T protocol.
    try:
        from gr00t.policy.server_client import PolicyClient

        client = PolicyClient(host=host, port=port, timeout_ms=timeout_ms)
        return bool(client.ping())
    except Exception:
        # Reachable on TCP but the client/import failed; treat as "up enough".
        return True


def _resolve_libero_python(explicit: str | None) -> str:
    if explicit:
        p = Path(explicit)
        if not p.exists():
            raise SystemExit(f"--libero-python {explicit} does not exist.")
        return str(p)
    if LIBERO_VENV_PYTHON.exists():
        return str(LIBERO_VENV_PYTHON)
    # Fall back to the current interpreter, but warn -- LIBERO needs its own venv.
    print(
        "WARNING: the dedicated LIBERO uv venv was not found at\n"
        f"  {LIBERO_VENV_PYTHON}\n"
        "Falling back to the current Python interpreter. If LIBERO / robosuite\n"
        "are not importable there, set up the sim env first:\n"
        "  sudo apt update && sudo apt install libegl1-mesa-dev libglu1-mesa\n"
        "  bash gr00t/eval/sim/LIBERO/setup_libero.sh\n",
        file=sys.stderr,
    )
    return sys.executable


def _start_server(model_path: Path, host: str, port: int, log_path: Path):
    cmd = [
        sys.executable, str(REPO_ROOT / "gr00t" / "eval" / "run_gr00t_server.py"),
        "--model-path", str(model_path),
        "--embodiment-tag", "LIBERO_PANDA",
        "--use-sim-policy-wrapper",
        "--host", host,
        "--port", str(port),
    ]
    log_f = open(log_path, "w")
    print(f"Starting GR00T server (logging to {log_path}):\n  {' '.join(cmd)}")
    proc = subprocess.Popen(cmd, stdout=log_f, stderr=subprocess.STDOUT, cwd=str(REPO_ROOT))
    return proc, log_f


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #
def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model-path", default="checkpoints/GR00T-N1.7-LIBERO/libero_10")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=5555)
    ap.add_argument("--manifest", default=str(DEFAULT_MANIFEST))
    ap.add_argument("--output-dir", default="outputs/libero_smoke_tests")
    ap.add_argument("--max-episode-steps", type=int, default=50,
                    help="Upper cap (and default) on max_episode_steps for every scenario.")
    ap.add_argument("--n-action-steps", type=int, default=8)
    ap.add_argument("--save-video", action="store_true",
                    help="Force-enable video recording for every scenario.")
    ap.add_argument("--render", action="store_true",
                    help="Also dump decoded frames into <scenario>/frames/.")
    ap.add_argument("--dry-run", action="store_true",
                    help="Validate manifest/model path/server, write metadata, run no sims.")
    ap.add_argument("--resume", action="store_true",
                    help="Skip scenarios that already have a rollout_summary.json in the run dir.")
    ap.add_argument("--run-dir", default=None,
                    help="Existing <output-dir>/<timestamp> dir to resume into "
                         "(default: create a new timestamped dir).")
    ap.add_argument("--libero-python", default=None,
                    help="Path to the LIBERO uv venv python (default: auto-detect).")
    ap.add_argument("--start-server", action="store_true",
                    help="Launch the GR00T server in a subprocess and kill it at the end.")
    ap.add_argument("--per-scenario-timeout", type=int, default=900,
                    help="Hard timeout (seconds) for each scenario subprocess.")
    args = ap.parse_args()

    os.chdir(REPO_ROOT)

    model_path = Path(args.model_path)
    manifest_path = Path(args.manifest)
    scenarios = _load_manifest(manifest_path)

    print(f"Loaded {len(scenarios)} scenarios from {manifest_path}")
    n_normal = sum(1 for s in scenarios if s.get("label") == "normal")
    n_abnormal = sum(1 for s in scenarios if s.get("label") == "abnormal_probe")
    print(f"  labels: {n_normal} normal, {n_abnormal} abnormal_probe")

    server_cmd_str = " ".join(_server_command(model_path))
    print("\nExpected GR00T server command (run this in a separate terminal):\n  " + server_cmd_str + "\n")
    print("NOTE: " + GATED_BACKBONE_HINT)

    # ---- validate model path -------------------------------------------------
    _check_model_path(model_path)
    print(f"Model checkpoint OK: {model_path}")

    # ---- output dir ----------------------------------------------------------
    out_root = Path(args.output_dir)
    if args.run_dir:
        run_dir = Path(args.run_dir)
        run_dir.mkdir(parents=True, exist_ok=True)
    else:
        ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir = out_root / ts
        run_dir.mkdir(parents=True, exist_ok=True)
    print(f"Run directory: {run_dir}")

    # ---- optionally start server --------------------------------------------
    server_proc = None
    server_log_f = None
    if args.start_server and not args.dry_run:
        server_proc, server_log_f = _start_server(model_path, args.host, args.port,
                                                  run_dir / "server.log")

    # ---- check server reachability ------------------------------------------
    reachable = _ping_server(args.host, args.port)
    if args.start_server and not args.dry_run and not reachable:
        # give the server time to load the (multi-GB) model
        print("Waiting for the GR00T server to come up (loading model can take a few minutes)...")
        for _ in range(120):
            time.sleep(5)
            if server_proc is not None and server_proc.poll() is not None:
                hint = ""
                if _scan_log_for_gated_repo(run_dir / "server.log"):
                    hint = "\n\n" + GATED_BACKBONE_HINT
                raise SystemExit(
                    f"The GR00T server process exited (rc={server_proc.returncode}); "
                    f"see {run_dir / 'server.log'}.{hint}"
                )
            if _ping_server(args.host, args.port):
                reachable = True
                break

    if not reachable and not args.dry_run:
        gated_hint = ""
        if args.start_server and _scan_log_for_gated_repo(run_dir / "server.log"):
            gated_hint = "\n" + GATED_BACKBONE_HINT
        msg = (
            f"\nERROR: no GR00T policy server is reachable at {args.host}:{args.port}.\n\n"
            "Start it first in a separate terminal:\n\n  " + server_cmd_str + "\n\n"
            "...then re-run this script. (Or pass --start-server to have this script\n"
            "launch it for you, or --dry-run to validate everything without running sims.)\n"
            + gated_hint
        )
        print(msg, file=sys.stderr)
        return 2
    if reachable:
        print(f"GR00T server reachable at {args.host}:{args.port}.")
    else:
        print("(dry-run) skipping server reachability requirement.")

    libero_python = _resolve_libero_python(args.libero_python)
    print(f"LIBERO rollout interpreter: {libero_python}")

    # ---- run scenarios -------------------------------------------------------
    rows: list[dict] = []
    try:
        for i, sc in enumerate(scenarios, 1):
            sid = sc["id"]
            label = sc.get("label", "normal")
            seed = int(sc.get("seed", 0))
            # Keep a scenario's own (smaller) budget, but cap larger ones at --max-episode-steps.
            mes = min(int(sc.get("max_episode_steps", args.max_episode_steps)), args.max_episode_steps)
            n_action_steps = int(sc.get("n_action_steps", args.n_action_steps))
            env_name = sc["env_name"]
            save_video = bool(sc.get("save_video", True)) or args.save_video
            obs_noise_std = float(sc.get("obs_noise_std", 0.0) or 0.0)
            action_repeat = int(sc.get("action_repeat", 1) or 1)
            instruction_override = sc.get("instruction_override")

            sc_dir = run_dir / sid
            sc_dir.mkdir(parents=True, exist_ok=True)

            print(f"\n[{i}/{len(scenarios)}] {sid}  ({label})  env={env_name}  "
                  f"seed={seed}  max_episode_steps={mes}")

            # metadata.json (always written)
            metadata = {
                "scenario_id": sid,
                "label": label,
                "manifest_entry": sc,
                "resolved": {
                    "env_name": env_name,
                    "seed": seed,
                    "max_episode_steps": mes,
                    "n_action_steps": n_action_steps,
                    "save_video": save_video,
                    "save_frames": bool(args.render),
                    "obs_noise_std": obs_noise_std,
                    "action_repeat": action_repeat,
                    "instruction_override": instruction_override,
                },
                "model_path": str(model_path),
                "policy_server": {"host": args.host, "port": args.port},
                "libero_python": libero_python,
                "timestamp": _dt.datetime.now().isoformat(timespec="seconds"),
            }
            with open(sc_dir / "metadata.json", "w") as f:
                json.dump(metadata, f, indent=2)

            if instruction_override:
                print(f"  NOTE: instruction_override is set ({instruction_override!r}) but is not "
                      "yet wired into the LIBERO env path; it will be ignored by this rollout.")

            # resume?
            if args.resume and (sc_dir / "rollout_summary.json").exists():
                try:
                    prev = json.loads((sc_dir / "rollout_summary.json").read_text())
                except Exception:
                    prev = {}
                print(f"  --resume: existing rollout_summary.json found, skipping.")
                rows.append({
                    "scenario_id": sid, "label": label, "seed": seed,
                    "rollout_started": prev.get("error") is None,
                    "actions_produced": bool(prev.get("actions_path")),
                    "video_saved": bool(prev.get("video_path")),
                    "success": prev.get("success", "unknown"),
                    "output_dir": str(sc_dir),
                    "error_if_any": prev.get("error"),
                    "resumed": True,
                })
                continue

            if args.dry_run:
                with open(sc_dir / "rollout_summary.json", "w") as f:
                    json.dump({"status": "dry_run", "env_name": env_name, "seed": seed}, f, indent=2)
                rows.append({
                    "scenario_id": sid, "label": label, "seed": seed,
                    "rollout_started": False, "actions_produced": False, "video_saved": False,
                    "success": "unknown", "output_dir": str(sc_dir),
                    "error_if_any": "dry-run (no rollout executed)",
                })
                continue

            # build worker command
            cmd = [
                libero_python, str(WORKER),
                "--env-name", env_name,
                "--host", args.host, "--port", str(args.port),
                "--max-episode-steps", str(mes),
                "--n-action-steps", str(n_action_steps),
                "--seed", str(seed),
                "--out-dir", str(sc_dir),
                "--obs-noise-std", str(obs_noise_std),
                "--action-repeat", str(action_repeat),
            ]
            cmd += ["--save-video"] if save_video else ["--no-save-video"]
            if args.render and save_video:
                cmd += ["--save-frames"]

            stdout_path = sc_dir / "stdout.log"
            stderr_path = sc_dir / "stderr.log"
            t0 = time.time()
            try:
                with open(stdout_path, "w") as so, open(stderr_path, "w") as se:
                    proc = subprocess.run(
                        cmd, stdout=so, stderr=se, cwd=str(REPO_ROOT),
                        timeout=args.per_scenario_timeout,
                    )
                rc = proc.returncode
            except subprocess.TimeoutExpired:
                rc = -9
                with open(stderr_path, "a") as se:
                    se.write(f"\n[runner] scenario timed out after {args.per_scenario_timeout}s\n")
            elapsed = time.time() - t0

            # parse worker result
            result = None
            try:
                for line in reversed(stdout_path.read_text().splitlines()):
                    line = line.strip()
                    if line.startswith(RESULT_PREFIX):
                        result = json.loads(line[len(RESULT_PREFIX):].strip())
                        break
            except Exception:
                pass
            if result is None:
                err_tail = ""
                try:
                    err_tail = "\n".join(stderr_path.read_text().splitlines()[-5:])
                except Exception:
                    pass
                result = {
                    "ok": False, "rollout_started": False, "actions_produced": False,
                    "video_saved": False, "success": None,
                    "error": (f"worker exited rc={rc} with no result line; "
                              f"see {stderr_path}. tail:\n{err_tail}"),
                }

            status = "ok" if result.get("ok") else "FAILED"
            print(f"  -> {status}  rc={rc}  {elapsed:.1f}s  "
                  f"actions={'yes' if result.get('actions_produced') else 'no'}  "
                  f"video={'yes' if result.get('video_saved') else 'no'}  "
                  f"success={result.get('success')}")
            if result.get("error"):
                print(f"     error: {result['error']}")

            rows.append({
                "scenario_id": sid, "label": label, "seed": seed,
                "rollout_started": bool(result.get("rollout_started")),
                "actions_produced": bool(result.get("actions_produced")),
                "video_saved": bool(result.get("video_saved")),
                "success": result.get("success") if result.get("success") is not None else "unknown",
                "output_dir": str(sc_dir),
                "error_if_any": result.get("error"),
                "rc": rc,
                "elapsed_sec": round(elapsed, 1),
                "n_action_calls": result.get("n_action_calls"),
                "episode_length": result.get("episode_length"),
                "video_path": result.get("video_path"),
                "actions_path": result.get("actions_path"),
            })
    finally:
        if server_proc is not None:
            print("\nStopping GR00T server subprocess...")
            server_proc.terminate()
            try:
                server_proc.wait(timeout=15)
            except subprocess.TimeoutExpired:
                server_proc.kill()
            if server_log_f is not None:
                server_log_f.close()

    # ---- summary -------------------------------------------------------------
    summary = {
        "run_dir": str(run_dir),
        "manifest": str(manifest_path),
        "model_path": str(model_path),
        "policy_server": {"host": args.host, "port": args.port},
        "server_command": server_cmd_str,
        "dry_run": args.dry_run,
        "n_scenarios": len(scenarios),
        "n_normal": n_normal,
        "n_abnormal_probe": n_abnormal,
        "timestamp": _dt.datetime.now().isoformat(timespec="seconds"),
        "scenarios": rows,
        "totals": {
            "rollout_started": sum(1 for r in rows if r.get("rollout_started")),
            "actions_produced": sum(1 for r in rows if r.get("actions_produced")),
            "video_saved": sum(1 for r in rows if r.get("video_saved")),
            "errors": sum(1 for r in rows if r.get("error_if_any") and "dry-run" not in str(r.get("error_if_any"))),
        },
    }
    write_summary(run_dir, summary)
    print_table(rows)
    print(f"\nWrote: {run_dir / 'summary.json'}")
    print(f"Wrote: {run_dir / 'summary.md'}")
    print(f"\nReview the videos with:\n  python {Path(__file__).resolve().parent / 'review_smoke_tests.py'} "
          f"--run-dir {run_dir}\n")

    # exit code: 0 if dry-run or every scenario at least started; else 1
    if args.dry_run:
        return 0
    failed = summary["totals"]["errors"]
    if failed == 0:
        return 0
    print(f"{failed}/{len(scenarios)} scenario(s) reported an error -- exiting non-zero.", file=sys.stderr)
    return 1


def write_summary(run_dir: Path, summary: dict) -> None:
    """Write summary.json and summary.md. Returns nothing; raises on JSON errors."""
    with open(run_dir / "summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    # markdown
    lines = []
    lines.append(f"# LIBERO smoke-test summary\n")
    lines.append(f"- run dir: `{summary['run_dir']}`")
    lines.append(f"- model: `{summary['model_path']}`")
    lines.append(f"- server: `{summary['policy_server']['host']}:{summary['policy_server']['port']}`")
    lines.append(f"- server command: `{summary['server_command']}`")
    lines.append(f"- scenarios: {summary['n_scenarios']} ({summary['n_normal']} normal, "
                 f"{summary['n_abnormal_probe']} abnormal_probe)")
    if summary.get("dry_run"):
        lines.append(f"- **DRY RUN** (no rollouts executed)")
    t = summary["totals"]
    lines.append(f"- totals: rollout_started={t['rollout_started']}, actions_produced={t['actions_produced']}, "
                 f"video_saved={t['video_saved']}, errors={t['errors']}\n")
    lines.append("| scenario_id | label | seed | rollout_started | actions_produced | video_saved | success | output_dir | error_if_any |")
    lines.append("|---|---|---|---|---|---|---|---|---|")
    for r in summary["scenarios"]:
        err = (str(r.get("error_if_any")) or "").replace("\n", " ").replace("|", "\\|")
        if len(err) > 160:
            err = err[:157] + "..."
        lines.append("| {id} | {label} | {seed} | {rs} | {ap} | {vs} | {succ} | {od} | {err} |".format(
            id=r["scenario_id"], label=r["label"], seed=r["seed"],
            rs="yes" if r.get("rollout_started") else "no",
            ap="yes" if r.get("actions_produced") else "no",
            vs="yes" if r.get("video_saved") else "no",
            succ=r.get("success"), od=r["output_dir"], err=err or "",
        ))
    (run_dir / "summary.md").write_text("\n".join(lines) + "\n")


def print_table(rows: list[dict]) -> None:
    headers = ["scenario_id", "label", "seed", "rollout_started", "actions_produced",
               "video_saved", "success", "output_dir", "error_if_any"]
    def cell(r, h):
        if h == "rollout_started":
            return "yes" if r.get("rollout_started") else "no"
        if h == "actions_produced":
            return "yes" if r.get("actions_produced") else "no"
        if h == "video_saved":
            return "yes" if r.get("video_saved") else "no"
        if h == "error_if_any":
            e = str(r.get("error_if_any") or "")
            e = e.replace("\n", " ")
            return (e[:57] + "...") if len(e) > 60 else e
        return str(r.get(h, ""))
    table = [headers] + [[cell(r, h) for h in headers] for r in rows]
    widths = [max(len(row[i]) for row in table) for i in range(len(headers))]
    print("\n" + "=" * 8 + " SMOKE-TEST SUMMARY " + "=" * 8)
    for ri, row in enumerate(table):
        print("  " + " | ".join(c.ljust(widths[i]) for i, c in enumerate(row)))
        if ri == 0:
            print("  " + "-+-".join("-" * w for w in widths))


if __name__ == "__main__":
    raise SystemExit(main())