"""Install a model trained on Colab into this agent.
  python tools/install_model.py path/to/union-agent-trained.zip
- backs up the current model/ to model_backup_<timestamp>/
- copies model.npz + config.json from the zip's export/ into model/ (runtime code in model/ is kept)
- copies the zip's eval summary to eval/model_eval_summary.json (used by the report) and train log to training/
- runs a quick load + 3-prompt sanity check
Roll back: delete model/ and rename the backup folder to model/.
"""
import json, os, shutil, sys, time, zipfile
HERE = os.path.dirname(os.path.abspath(__file__)); ROOT = os.path.dirname(HERE)

def main(zpath):
    if not os.path.exists(zpath): sys.exit(f"not found: {zpath}")
    z = zipfile.ZipFile(zpath)
    names = z.namelist()
    def find(suffix):
        c = [n for n in names if n.endswith(suffix) and "__MACOSX" not in n]
        return min(c, key=len) if c else None
    npz, cfg, summ, log = find("model.npz"), find("config.json"), find("eval/summary.json"), find("train_log.jsonl")
    if not npz or not cfg: sys.exit("zip has no model.npz/config.json (expected the export/ folder from the Colab notebook)")
    cfg_new = json.loads(z.read(cfg))
    need = {"n_in", "n_out", "d", "h", "f", "n_enc", "n_dec", "max_in", "max_out"}
    if not need <= set(cfg_new): sys.exit(f"{cfg} is not a model config (missing {need - set(cfg_new)})")
    cur_cfg = json.load(open(os.path.join(ROOT, "model", "config.json")))
    if (cfg_new["n_in"], cfg_new["n_out"]) != (cur_cfg["n_in"], cur_cfg["n_out"]):
        sys.exit(f"vocab mismatch: new model n_in/n_out {cfg_new['n_in']}/{cfg_new['n_out']} vs agent {cur_cfg['n_in']}/{cur_cfg['n_out']} "
                 "(model trained with a different tokenizer/catalog; copy the whole export/ folder over model/ instead)")
    bak = os.path.join(ROOT, "model_backup_" + time.strftime("%Y%m%d_%H%M%S"))
    shutil.copytree(os.path.join(ROOT, "model"), bak, ignore=shutil.ignore_patterns("__pycache__"))
    open(os.path.join(ROOT, "model", "model.npz"), "wb").write(z.read(npz))
    json.dump({k: cfg_new[k] for k in need}, open(os.path.join(ROOT, "model", "config.json"), "w"), indent=1)
    if summ: open(os.path.join(ROOT, "eval", "model_eval_summary.json"), "wb").write(z.read(summ))
    if log: open(os.path.join(ROOT, "training", "train_log_colab.jsonl"), "wb").write(z.read(log))
    print(f"installed {npz} (d={cfg_new['d']}, enc={cfg_new['n_enc']}, dec={cfg_new['n_dec']}); previous model backed up to {os.path.basename(bak)}")
    sys.path.insert(0, os.path.join(ROOT, "model"))
    from runtime import CommandModel
    m = CommandModel(os.path.join(ROOT, "model"))
    print("params:", sum(v.size for v in m.W.values()))
    for p in ("awaaz pachaas pe rakh do", "wifi Airtel_5G se connect karo password Wifi@2024", "board pin 11 pe led jalao"):
        o = m.predict(p); print(" ", p, "->", {k: v for k, v in o.items() if k != "_meta"})
    if summ:
        s = json.loads(z.read(summ))
        for k in ("val", "test_strict", "curated", "wifi_holdout"):
            if k in s: print(f"  {k:14s} exact {s[k]['exact']*100:.2f}%  action {s[k]['action']*100:.2f}%")
    print("next: ./run.sh test   and   ./run.sh report   (Windows: run.bat test / run.bat report)")

if __name__ == "__main__":
    if len(sys.argv) != 2: sys.exit(__doc__)
    main(sys.argv[1])