"""Pull files off a Colab runtime over the HTTP runtime proxy, with no ssh.

The ssh tunnel is a websocket bridged through `albedo_plugins.ws`, and Colab allows exactly **one**
ssh session per runtime. If the process holding that bridge dies without releasing it (a kernel or
daemon restart is enough), the runtime answers `429 {"error":"already-active-session"}` to every new
tunnel and there is no local remedy: `pgrep albedo_plugins.ws` is empty, the token is still valid for
an hour, and the API's only escape hatch is `/tun/m/unassign`, which destroys the VM.

The runtime proxy itself is a different path and is not locked. The same `colab-runtime-proxy-token`
authenticates plain HTTP against it, so the Jupyter REST API stays reachable while ssh is stuck:

    GET  /api/status          -> {"connections":..,"kernels":..,"last_activity":..}   (is it alive)
    GET  /api/contents/<p>    -> directory listing, or one file as text
    GET  /files/<p>           -> the file; honours Range, so `curl -C -` style resume works

Usage:

    python3 scripts/pull_colab_http.py --host <runtime.netloc> --token <tok> \
        --remote content/swift15/checkpoints/x/latest.pt --local out/latest.pt

`--host` is the netloc of `instance.url` and `--token` is `instance.token`; both expire after about an
hour, so re-list the instance rather than caching them.
"""
import argparse
import os
import sys
import time
import urllib.error
import urllib.request

CHUNK = 1 << 20  # 1 MiB: stream to disk, never hold the file in memory


def size_and_status(url, headers):
    req = urllib.request.Request(url, headers={**headers, "Range": "bytes=0-0"})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            cr = r.headers.get("Content-Range") or ""
            total = int(cr.split("/")[-1]) if "/" in cr else int(r.headers.get("Content-Length") or 0)
            return total, r.status
    except urllib.error.HTTPError as e:
        return -1, e.code


def pull(url, local, headers, resume=True, quiet=False):
    os.makedirs(os.path.dirname(os.path.abspath(local)) or ".", exist_ok=True)
    # A checkpoint written every optimizer step keeps identical byte sizes (same tensor shapes), so
    # "same size" is NOT "same content": resume-by-size would silently never refresh latest.pt.
    # Pass resume=False for anything that is rewritten in place.
    have = os.path.getsize(local) if (resume and os.path.exists(local)) else 0
    total, status = size_and_status(url, headers)
    if status == 404:
        raise SystemExit(f"{local}: remote 404")
    if total > 0 and have == total:
        print(f"{local}: already complete ({total:,} bytes)")
        return total
    h = dict(headers)
    if have:
        h["Range"] = f"bytes={have}-"
    t0 = time.time()
    try:
        r = urllib.request.urlopen(urllib.request.Request(url, headers=h), timeout=120)
    except urllib.error.HTTPError as e:
        if e.code == 416:
            print(f"{local}: already complete")
            return have
        raise
    with r as resp, open(local, "ab" if have else "wb") as fh:
        mode = "resumed" if have else "new"
        while True:
            chunk = resp.read(CHUNK)
            if not chunk:
                break
            fh.write(chunk)
            have += len(chunk)
            if not quiet:
                pct = (100.0 * have / total) if total > 0 else 0.0
                sys.stderr.write(f"\r  {local}: {have:,}/{total:,} ({pct:5.1f}%)")
                sys.stderr.flush()
    dt = time.time() - t0
    rate = (have / dt / 1e6) if dt else 0.0
    print(f"{local}: {mode}, {have:,} bytes in {dt:.1f}s ({rate:.1f} MB/s)")
    return have


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--host", required=True, help="netloc of colab instance.url")
    ap.add_argument("--token", required=True, help="colab instance.token (query param)")
    ap.add_argument("--remote", required=True, help="path under the runtime's /content root")
    ap.add_argument("--local", required=True, help="destination path on this host")
    ap.add_argument("--no-resume", action="store_true")
    ap.add_argument("--quiet", action="store_true")
    a = ap.parse_args(argv)
    remote = a.remote.lstrip("/")
    url = f"https://{a.host}/files/{remote}?colab-runtime-proxy-token={a.token}"
    headers = {"X-Colab-Client-Agent": "colab-cli"}
    pull(url, a.local, headers, resume=not a.no_resume, quiet=a.quiet)
    return 0


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