fetch_models.py

"""Get the basic models from OUR vault - nobody else's site.

The console ships without weights (they are 4-8 GB each and every one carries
its author's licence). This fetcher pulls them from the steward's own store,
verifies each SHA-256 from models.lock.json, and drops them where the console
scans. Standard library only, so the bundled Python can run it anywhere.

  python fetch_models.py --list
  python fetch_models.py --profile basic
  python fetch_models.py --profile full --dest "I:/LYGO_MODELS"
  python fetch_models.py --check                 # verify what is already here
  python fetch_models.py --profile basic --mirror

Nothing is downloaded from a URL that is not in models.lock.json, and no file
is accepted unless its SHA-256 matches the lock. Re-running is safe: files
already correct are skipped (and a half-finished download resumes).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

HERE = Path(__file__).resolve().parent
DEFAULT_LOCK = HERE / "models.lock.json"
CHUNK = 1024 * 1024


def human(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024 or unit == "TB":
            return "%.1f %s" % (n, unit)
        n /= 1024.0
    return "%.1f TB" % n


def load_lock(path: Path) -> dict:
    if not path.exists():
        sys.exit("no lock file at %s - the console install is incomplete" % path)
    lock = json.loads(path.read_text(encoding="utf-8"))
    if not lock.get("models"):
        sys.exit("lock file %s lists no models" % path)
    return lock


def pick_dest(explicit: str | None) -> Path:
    if explicit:
        return Path(explicit).expanduser().resolve()
    env = os.environ.get("LYGO_MODELS")
    if env:
        return Path(env).expanduser().resolve()
    for cand in (HERE.parent / "models", HERE / "models"):
        if cand.is_dir():
            return cand.resolve()
    return (HERE.parent / "models").resolve()


def sha256_of(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for block in iter(lambda: fh.read(CHUNK), b""):
            h.update(block)
    return h.hexdigest()


def url_for(lock: dict, name: str, use_mirror: bool) -> str:
    hosts = lock.get("hosts") or {}
    key = "mirror" if use_mirror else "primary"
    base = hosts.get(key) or hosts.get("primary")
    if not base:
        sys.exit("lock file has no %s host" % key)
    if not base.endswith("/"):
        base += "/"
    return base + name


def host_allowed(lock: dict, url: str) -> bool:
    allowed = [h for h in (lock.get("hosts") or {}).values() if isinstance(h, str) and h.startswith("http")]
    return any(url.startswith(a if a.endswith("/") else a + "/") or url.startswith(a) for a in allowed)


def download(url: str, dest: Path, entry: dict, lock: dict) -> tuple[bool, str]:
    if not host_allowed(lock, url):
        return False, "refused: %s is not one of our hosts" % url
    want_bytes = int(entry["bytes"])
    want_sha = entry["sha256"].lower()
    part = dest.with_suffix(dest.suffix + ".part")

    have = part.stat().st_size if part.exists() else 0
    if have > want_bytes:
        part.unlink()
        have = 0

    req = urllib.request.Request(url, headers={"User-Agent": "lygo-console-fetch/1"})
    if have:
        req.add_header("Range", "bytes=%d-" % have)
        print("   resuming at %s" % human(have))
    mode = "ab" if have else "wb"
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=120) as r, part.open(mode) as out:
            if have and r.status != 206:
                out.close()
                part.unlink(missing_ok=True)
                return download(url, dest, entry, lock)
            done = have
            while True:
                block = r.read(CHUNK)
                if not block:
                    break
                out.write(block)
                done += len(block)
                pct = 100.0 * done / want_bytes if want_bytes else 0.0
                speed = done / max(1e-6, time.time() - t0)
                sys.stdout.write("\r   %5.1f%%  %s / %s  (%s/s)   " % (pct, human(done), human(want_bytes), human(speed)))
                sys.stdout.flush()
    except urllib.error.HTTPError as e:
        return False, "http %s from %s" % (e.code, url)
    except Exception as e:  # noqa: BLE001 - report anything the network throws
        print()
        return False, "%s: %s" % (e.__class__.__name__, e)
    print()

    got = part.stat().st_size
    if got != want_bytes:
        return False, "wrong size: got %s, lock says %s (kept as %s)" % (got, want_bytes, part.name)
    digest = sha256_of(part)
    if digest != want_sha:
        return False, "SHA-256 mismatch - refusing it (got %s..., lock says %s...)" % (digest[:16], want_sha[:16])
    shutil.move(str(part), str(dest))
    return True, "sha256 ok"


def write_vault_manifest(dest: Path, lock: dict) -> None:
    """Keep the vault manifest the console reads in sync (same shape as before)."""
    path = dest / "manifest.json"
    data: dict = {}
    if path.exists():
        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except Exception:  # noqa: BLE001 - a broken manifest must not block a fetch
            data = {}
    for entry in lock["models"]:
        local = dest / entry["file"]
        if entry.get("kind") == "mmproj" or not local.exists():
            continue
        vid = entry["vault_id"]
        rec = data.get(vid) or {"id": vid}
        rec["path"] = str(local)
        rec["bytes"] = local.stat().st_size
        rec["sha256"] = sha256_of(local)
        rec["ok"] = rec["sha256"] == entry["sha256"].lower()
        rec.setdefault("mmproj", None)
        for other in lock["models"]:
            if other.get("kind") == "mmproj" and other.get("vault_id") == vid:
                p = dest / other["file"]
                if p.exists():
                    rec["mmproj"] = str(p)
                    rec["mmproj_sha256"] = sha256_of(p)
        data[vid] = rec
    path.write_text(json.dumps(data, indent=1), encoding="utf-8")


def check_only(dest: Path, lock: dict) -> int:
    bad = 0
    for entry in lock["models"]:
        local = dest / entry["file"]
        if not local.exists():
            print("  MISSING  %s" % entry["file"])
            bad += 1
            continue
        size_ok = local.stat().st_size == int(entry["bytes"])
        digest = sha256_of(local) if size_ok else ""
        if size_ok and digest == entry["sha256"].lower():
            print("  ok       %s  (%s, %s)" % (entry["file"], human(local.stat().st_size), entry["licence"]))
        else:
            print("  BAD      %s  (%s, sha256 %s)" % (entry["file"], human(local.stat().st_size), digest[:16] or "size"))
            bad += 1
    print()
    print("  %s" % ("all present and verified" if not bad else "%d file(s) not right" % bad))
    return 1 if bad else 0


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Fetch the basic LYGO models from our own vault.")
    ap.add_argument("--lock", default=str(DEFAULT_LOCK))
    ap.add_argument("--dest", default=None, help="where to put them (default: LYGO_MODELS, else the console's models\\)")
    ap.add_argument("--profile", default=None, help="core / basic / full / coder / embed (default: basic)")
    ap.add_argument("--mirror", action="store_true", help="use the mirror host instead of the primary")
    ap.add_argument("--list", action="store_true", help="show what the vault offers and exit")
    ap.add_argument("--check", action="store_true", help="verify the models already on this machine, download nothing")
    ap.add_argument("--yes", action="store_true", help="do not stop for confirmation")
    args = ap.parse_args(argv)

    lock = load_lock(Path(args.lock))
    by_id = {m["id"]: m for m in lock["models"]}
    hosts = lock.get("hosts") or {}

    print("LYGO model vault  ·  %s" % lock.get("product", "console"))
    print("  store   : %s" % hosts.get("repo_page", hosts.get("primary", "?")))
    print("  mirror  : %s" % hosts.get("mirror_page", "?"))
    print("  revision: %s" % hosts.get("revision", "main"))
    print()

    if args.list:
        for pid, ids in sorted((lock.get("profiles") or {}).items()):
            total = sum(int(by_id[i]["bytes"]) for i in ids if i in by_id)
            print("  profile %-6s  %s" % (pid, human(total)))
            for i in ids:
                e = by_id.get(i)
                if e:
                    print("      %-22s %9s  %-11s %s" % (e["file"], human(e["bytes"]), e["licence"], ",".join(e.get("modalities") or ["text"])))
        print()
        return 0

    dest = pick_dest(args.dest)
    dest.mkdir(parents=True, exist_ok=True)

    if args.check:
        print("  checking %s" % dest)
        return check_only(dest, lock)

    profile = (args.profile or "basic").lower()
    if profile not in (lock.get("profiles") or {}):
        sys.exit("unknown profile %r - try --list (known: %s)" % (profile, ", ".join(lock.get("profiles") or {})))
    wanted = [by_id[i] for i in lock["profiles"][profile] if i in by_id]
    total = sum(int(m["bytes"]) for m in wanted)

    print("  profile : %s  (%d file(s), %s)" % (profile, len(wanted), human(total)))
    print("  into    : %s" % dest)
    for m in wanted:
        print("      %-24s %9s  %s  %s" % (m["file"], human(m["bytes"]), m["licence"], m["author"]))
    print()
    if not args.yes:
        try:
            if input("  download now? [y/N] ").strip().lower() not in ("y", "yes"):
                print("  stopped - nothing downloaded")
                return 0
        except EOFError:
            print("  no answer available - stopping (use --yes for unattended runs)")
            return 0

    failures = []
    for m in wanted:
        target = dest / m["file"]
        if target.exists() and target.stat().st_size == int(m["bytes"]) and sha256_of(target) == m["sha256"].lower():
            print("  have     %s (already verified)" % m["file"])
            continue
        print("  get      %s  %s  %s" % (m["file"], human(m["bytes"]), m["licence"]))
        print("           from %s" % url_for(lock, m["file"], args.mirror))
        ok, why = download(url_for(lock, m["file"], args.mirror), target, m, lock)
        print("           %s: %s" % ("done" if ok else "FAILED", why))
        if not ok:
            failures.append((m["file"], why))

    write_vault_manifest(dest, lock)
    print()
    print("  licence/attribution for what was fetched:")
    for m in wanted:
        print("      %-24s %-11s %s - %s" % (m["file"], m["licence"], m["author"], m["upstream"]))
    print("      full Apache-2.0 text: %s" % hosts.get("license_text", "https://www.apache.org/licenses/LICENSE-2.0.txt"))
    print()
    if failures:
        print("  %d of %d failed - nothing else was touched:" % (len(failures), len(wanted)))
        for name, why in failures:
            print("      %s: %s" % (name, why))
        return 2
    print("  all set - start the console and it will pick the model up")
    return 0


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