"""Re-run NIUA's robot-file benchmark yourself.

We broke our own robot arm 8 ways (one break per copy) and kept 5 good
robots. This script downloads those 13 files, then asks three tools about
each one:

  - NIUA's check (urdf.sim-ready@1), through the public API
  - yourdfpy: load the robot with its meshes, then validate()
  - MuJoCo: compile the robot into a simulation model

and prints how many of the 8 broken copies each tool flags, next to the
numbers we published.

    pip install mujoco==3.16.0 yourdfpy==0.0.60 trimesh==4.10.0
    python3 reproduce.py        (Windows: py reproduce.py)

NIUA's column runs live on api.ohao.tech, no account needed. If the API
is busy or you hit the rate limit, the script shows our published NIUA
column, clearly marked, and still runs the other two tools on your
machine. NIUA_API_KEY is optional. Nothing you run here is stored.
Simulation only, no hardware tested.
"""
from __future__ import annotations

import io, json, os, sys, tempfile, urllib.error, urllib.request, zipfile
from importlib.metadata import version
from pathlib import Path

SITE = os.environ.get("NIUA_SITE", "https://ohao.tech")
API = os.environ.get("NIUA_API", "https://api.ohao.tech")
CONTRACT = "urdf.sim-ready@1"
PINNED = {"mujoco": "3.16.0", "yourdfpy": "0.0.60", "trimesh": "4.10.0"}
UA = {"User-Agent": "niua-benchmark/1"}

try:
    import mujoco
    import yourdfpy
except ImportError:
    sys.exit("Install the two open-source tools first:\n  pip install mujoco==3.16.0 yourdfpy==0.0.60 trimesh==4.10.0")


def get(url: str) -> bytes:
    try:
        with urllib.request.urlopen(urllib.request.Request(url, headers=UA), timeout=60) as r:
            return r.read()
    except OSError as e:  # URLError, and read timeouts (socket.timeout on 3.9)
        raise SystemExit(f"Could not download {url}: {e}. Check your connection and try again.")


# Why the NIUA column could not run live, for the closing note.
SKIPPED = {"why": ""}


def niua_check(data: bytes, key: str | None) -> dict | None:
    """NIUA's verdict and failing rules, or None if it cannot run live right now."""
    headers = {**UA, "Content-Type": "application/octet-stream"}
    if key:
        headers["Authorization"] = f"Bearer {key}"
    req = urllib.request.Request(f"{API}/api/v1/playground/check?contract={CONTRACT}", data=data, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            body = json.load(r)
    except urllib.error.HTTPError as e:
        if e.code in (401, 403):
            SKIPPED["why"] = ("Your NIUA_API_KEY was refused." if key else
                              "The API asked for an account; set NIUA_API_KEY (https://ohao.tech/niua/console/keys) to run it live.")
            return None
        if e.code in (429, 503):
            SKIPPED["why"] = "The API is busy or you hit the rate limit; try again in a few minutes."
            return None
        raise SystemExit(f"NIUA API answered {e.code}: {e.read()[:300].decode('utf-8', 'replace')}")
    except OSError as e:  # URLError, and read timeouts (socket.timeout on 3.9)
        SKIPPED["why"] = f"Could not reach the API ({e})."
        return None
    return {"verdict": body["verdict"], "rules": sorted({f["rule"] for f in body.get("failures", [])})}


def urdf_in(path: Path, tmp: Path) -> Path:
    if path.suffix == ".urdf":
        return path
    with zipfile.ZipFile(path) as z:
        z.extractall(tmp)
        return tmp / next(n for n in z.namelist() if n.endswith(".urdf"))


def yourdf(urdf: Path) -> str:
    try:
        root = urdf.parent.parent
        def handler(fname):
            return str(root / fname.split("/", 3)[-1]) if fname.startswith("package://") else str(urdf.parent / fname)
        r = yourdfpy.URDF.load(str(urdf), filename_handler=handler, load_meshes=True, build_collision_scene_graph=True, load_collision_meshes=True)
        return "ok" if r.validate() else "invalid"
    except Exception as e:
        return f"error: {type(e).__name__}"


def mj(urdf: Path, scratch: Path) -> str:
    # MuJoCo's URDF importer finds meshes by file name in meshdir.
    xml = urdf.read_text(encoding="utf-8")
    meshdir = (urdf.parent.parent / "meshes").as_posix()
    head, rest = xml.split("<robot", 1)
    end = rest.index(">") + 1
    tmp = scratch / "_mj.urdf"
    tmp.write_text(head + "<robot" + rest[:end] + f'<mujoco><compiler meshdir="{meshdir}" strippath="true"/></mujoco>' + rest[end:], encoding="utf-8")
    try:
        mujoco.MjModel.from_xml_path(str(tmp))
        return "compiles"
    except Exception as e:
        return "error: " + str(e).splitlines()[0][:60]


def main() -> int:
    print(f"Downloading the 13 test robots from {SITE} ...")
    published = json.loads(get(f"{SITE}/evidence/urdf-oracle/result.json"))
    pub_rows = {r["fixture"]: r for r in published["rows"]}
    key = os.environ.get("NIUA_API_KEY") or None
    with tempfile.TemporaryDirectory() as t:
        fx = Path(t) / "fixtures"
        zipfile.ZipFile(io.BytesIO(get(f"{SITE}/evidence/urdf-oracle/fixtures.zip"))).extractall(fx)
        expected = json.loads((fx / "expected.json").read_text(encoding="utf-8"))
        rows, niua_live = [], True
        for name in sorted(expected):
            g = niua_check((fx / name).read_bytes(), key) if niua_live else None
            if g is None:
                niua_live = False
                p = pub_rows[name]
                g = {"verdict": p["gate"].split()[0], "rules": [] if p["gate_rules"] == "-" else p["gate_rules"].split(", ")}
            with tempfile.TemporaryDirectory() as s:
                urdf = urdf_in(fx / name, Path(s))
                rows.append({"fixture": name, "broken": expected[name]["verdict"] == "fail", "niua": g, "yourdfpy": yourdf(urdf), "mujoco": mj(urdf, Path(s))})

    w = max(len(r["fixture"]) for r in rows)
    print(f"\n{'robot file':<{w}}  {'should':<6}  {'NIUA':<5}  {'yourdfpy':<16}  MuJoCo")
    for r in rows:
        print(f"{r['fixture']:<{w}}  {'fail' if r['broken'] else 'pass':<6}  {r['niua']['verdict']:<5}  {r['yourdfpy']:<16}  {r['mujoco']}")

    broken = [r for r in rows if r["broken"]]
    flagged = {
        "NIUA": sum(r["niua"]["verdict"] == "fail" for r in broken),
        "yourdfpy": sum(r["yourdfpy"] != "ok" for r in broken),
        "MuJoCo": sum(r["mujoco"] != "compiles" for r in broken),
    }
    want = {"NIUA": published["flagged"]["gate"], "yourdfpy": published["flagged"]["yourdfpy"], "MuJoCo": published["flagged"]["mujoco"]}
    false_alarms = sum(r["niua"]["verdict"] == "fail" for r in rows if not r["broken"])

    print(f"\nBroken copies flagged (of {len(broken)}):")
    for tool in flagged:
        note = "" if flagged[tool] == want[tool] else f"   <- we published {want[tool]}/{len(broken)}"
        print(f"  {tool:<9} {flagged[tool]}/{len(broken)}{note}")
    print(f"Good robots NIUA wrongly failed: {false_alarms} of {len(rows) - len(broken)}")
    if not niua_live:
        print("\nNIUA column: our published result, not re-run. " + SKIPPED["why"])
    else:
        print("\nNIUA column: run live on api.ohao.tech just now.")
    tools = {"mujoco": mujoco.__version__, "yourdfpy": version("yourdfpy"), "trimesh": version("trimesh")}
    drift = [f"{k} {v} (we used {PINNED[k]})" for k, v in tools.items() if v != PINNED[k]]
    if drift:
        print("Different tool versions than ours, results may differ: " + ", ".join(drift))
    same = flagged == want and false_alarms == 0
    print("\nResult: " + ("matches what we published." if same else "differs from what we published. Please tell us on Discord: https://discord.gg/NJZQcTzD5C"))
    return 0 if same else 1


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