#!/usr/bin/env python3
"""Run code inside a Colab runtime through its Jupyter kernel, the way `colab exec` does.

Colab's ssh proxy serves one WebSocket client per runtime and wedges server-side when a client dies
without closing cleanly, which leaves no way in short of destroying the VM. The kernel path has no
such limit: `jupyter_kernel_client` speaks the framing (it is the library google-colab-cli uses), so
this only has to hand it the runtime URL and the proxy token.

Usage: colab_kernel_exec.py --host HOST --token TOKEN --code "python"
"""
import argparse
import json
import urllib.request

import jupyter_kernel_client as jkc


def make_client(host, token, kernel_id=None):
    kwargs = {
        "subprotocol": jkc.JupyterSubprotocol.DEFAULT,
        "extra_params": {"colab-runtime-proxy-token": token},
    }
    server_url = f"https://{host}"
    headers = {"X-Colab-Client-Agent": "albedo", "X-Colab-Runtime-Proxy-Token": token}
    if hasattr(jkc, "ColabKernelClient"):
        cls, tok = jkc.ColabKernelClient, "proxy_token"
    elif hasattr(jkc, "KernelClient"):
        cls, tok = jkc.KernelClient, "token"
    else:
        cls, tok = jkc.JupyterKernelClient, "token"
    client = cls(server_url=server_url, **{tok: token}, kernel_id=kernel_id,
                 client_kwargs=kwargs, headers=headers)
    # Never let the library reap a kernel it did not create.
    client._own_kernel = False
    return client


def run(host, token, code, timeout, kernel_id=None):
    if kernel_id:
        client = make_client(host, token, kernel_id)
        client.start()
        return client, kernel_id
    req = urllib.request.Request(
        f"https://{host}/api/kernels?colab-runtime-proxy-token={token}",
        data=json.dumps({"name": "python3"}).encode(), method="POST",
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        kernel_id = json.loads(resp.read())["id"]
    client = make_client(host, token, kernel_id)
    client.start()
    return client, kernel_id


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", required=True)
    ap.add_argument("--token", required=True)
    ap.add_argument("--code", default="")
    ap.add_argument("--file", default="")
    ap.add_argument("--timeout", type=float, default=120.0)
    ap.add_argument("--kernel-id", default="")
    a = ap.parse_args()

    code = Path(a.file).read_text() if a.file else a.code
    client, kernel_id = run(a.host, a.token, code, a.timeout, a.kernel_id or None)
    print(f"KERNEL_ID {kernel_id}", flush=True)
    result = client.execute(code, timeout=a.timeout)
    for out in result.get("outputs", []):
        kind = out.get("output_type")
        if kind == "stream":
            print(out.get("text", ""), end="", flush=True)
        elif kind in ("execute_result", "display_data"):
            print(out.get("data", {}).get("text/plain", ""), flush=True)
        elif kind == "error":
            print("\n".join(out.get("traceback", [])), flush=True)
    print(f"\nEXEC_STATUS {result.get('status')}", flush=True)


if __name__ == "__main__":
    main()