#!/usr/bin/env python
"""Relay for Colab runs that does not depend on the laptop staying up.

Runs on the Oracle VM (python:alpine container `goldenai-share`, Caddy host
goldenshare.<vm>.nip.io): GET serves the staged training bundle from /srv,
PUT stores checkpoints and the final adapter into /srv/uploads. The laptop
pushes the bundle with scp before the run and pulls the uploads with scp
after it, so a Claude session restart or a dead quick tunnel no longer loses
a checkpoint (2026-09-04: three trycloudflare tunnels in a row stopped
resolving mid-run).

    python vm_share.py --dir /srv --port 8123
"""

from __future__ import annotations

import argparse
import http.server
from functools import partial
from pathlib import Path


class Handler(http.server.SimpleHTTPRequestHandler):
    uploads: Path

    def do_PUT(self):  # noqa: N802
        name = Path(self.path.split("?")[0]).name
        if not name or name.startswith("."):
            self.send_response(400)
            self.end_headers()
            return
        length = int(self.headers.get("Content-Length", 0))
        dest = self.uploads / name
        tmp = self.uploads / (name + ".part")
        with open(tmp, "wb") as sink:
            remaining = length
            while remaining > 0:
                chunk = self.rfile.read(min(1 << 20, remaining))
                if not chunk:
                    break
                sink.write(chunk)
                remaining -= len(chunk)
        tmp.replace(dest)
        print(f"upload received: {name} ({dest.stat().st_size:,} bytes)", flush=True)
        self.send_response(201)
        self.end_headers()

    def log_message(self, fmt, *args):  # quieter than the default, still visible in docker logs
        print(self.address_string(), fmt % args, flush=True)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--dir", default="/srv")
    ap.add_argument("--port", type=int, default=8123)
    args = ap.parse_args()
    root = Path(args.dir)
    Handler.uploads = root / "uploads"
    Handler.uploads.mkdir(parents=True, exist_ok=True)
    httpd = http.server.ThreadingHTTPServer(("0.0.0.0", args.port), partial(Handler, directory=str(root)))
    print(f"share on 0.0.0.0:{args.port} serving {root}, uploads in {Handler.uploads}", flush=True)
    httpd.serve_forever()
    return 0


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