#!/usr/bin/env python3
"""Make a downloaded checkpoint loadable on the robot with NO Hugging Face access.

Two edits, both on a COPY (the downloaded checkpoint is never modified):

1. policy_preprocessor.json -> tokenizer_name points at the bundled PaliGemma
   tokenizer instead of "google/paligemma-3b-pt-224". That repo is gated, so on
   a robot without a logged-in HF account the gated download would fail at
   startup.
2. config.json -> drop `pretrained_revision` (per Gustavo: rosetta's
   policy_bridge_node rejects it).

Usage:
    python prepare_deploy.py <checkpoint_dir> [--tokenizer DIR] [--out DIR]
"""
import argparse, json, shutil
from pathlib import Path

here = Path(__file__).resolve().parent
ap = argparse.ArgumentParser()
ap.add_argument("checkpoint", help="dir containing model.safetensors, config.json, policy_*.json")
ap.add_argument("--tokenizer", default=str(here / "paligemma_tokenizer"))
ap.add_argument("--out", default=None, help="default: <checkpoint>_deploy")
a = ap.parse_args()

src = Path(a.checkpoint).resolve()
tok = Path(a.tokenizer).resolve()
out = Path(a.out).resolve() if a.out else src.with_name(src.name + "_deploy")

for need in ("model.safetensors", "config.json", "policy_preprocessor.json"):
    if not (src / need).exists():
        raise SystemExit(f"missing {need} in {src}")
if not (tok / "tokenizer.json").exists():
    raise SystemExit(f"no tokenizer.json in {tok}")

if out.exists():
    shutil.rmtree(out)
out.mkdir(parents=True)
for f in src.iterdir():
    if f.is_file():
        dst = out / f.name
        if f.name == "model.safetensors":
            try:
                dst.symlink_to(f)        # 8.9 GB: link, don't copy
            except OSError:
                shutil.copy2(f, dst)
        else:
            shutil.copy2(f, dst)

cfg = json.loads((out / "config.json").read_text())
removed = [k for k in ("pretrained_revision",) if k in cfg]
for k in removed:
    cfg.pop(k)
(out / "config.json").write_text(json.dumps(cfg, indent=4))

pre = json.loads((out / "policy_preprocessor.json").read_text())
changed = 0
for st in pre.get("steps", []):
    c = st.get("config") or {}
    if "tokenizer_name" in c:
        c["tokenizer_name"] = str(tok)
        changed += 1
(out / "policy_preprocessor.json").write_text(json.dumps(pre, indent=2))

if changed == 0:
    raise SystemExit("no tokenizer_name found in policy_preprocessor.json -- format changed?")
print(f"deploy-ready checkpoint: {out}")
print(f"  config.json: removed {removed or 'nothing'}")
print(f"  tokenizer_name -> {tok}")