bootstrap.py

"""Runtime setup for the AntiSite Space (Gradio SDK — no Dockerfile available).

A Gradio Space gives us pip and apt, but no image build steps, so the pieces that
a Dockerfile would normally bake in are assembled here on first boot:

  * the AntiSite and ParaSurf sources (cloned, not vendored, so the demo cannot
    drift from the published release),
  * DMS, compiled from ParaSurf's bundled source into a user-writable prefix —
    the shipped binary needs GLIBC 2.34 and its data paths are compiled in, so
    building with our own LIBDIR is what makes it work without root,
  * the pdb2pqr executable bit, which git does not always preserve,
  * the frozen ParaSurf weights.

Everything lands under $ANTISITE_WORK (default ~/antisite_runtime) and is skipped
if already present, so a warm container restarts quickly.
"""

from __future__ import annotations

import os
import shutil
import stat
import subprocess
import sys
from pathlib import Path

ANTISITE_REPO = "https://github.com/aggelos-michael-papadopoulos/AntiSite.git"
PARASURF_REPO = "https://github.com/aggelos-michael-papadopoulos/ParaSurf.git"

# Pin the demo to a released commit/tag so it always matches the manuscript.
# Override with the ANTISITE_REF variable in the Space settings.
ANTISITE_REF = os.environ.get("ANTISITE_REF", "main")

# Frozen ParaSurf (Paragraph-expanded) weights.
PARASURF_GDRIVE_ID = "1nd3npYK303e8owDBvW8Ygd5m9SD1puhR"

HOME = Path(os.environ.get("HOME", "/home/user"))
WORK = Path(os.environ.get("ANTISITE_WORK", HOME / "antisite_runtime"))
ANTISITE_ROOT = WORK / "AntiSite"
PARASURF_ROOT = ANTISITE_ROOT / "ParaSurf"   # ParaSurf must sit inside the repo root
PREFIX = WORK / "local"                       # user-writable install prefix for DMS
PARASURF_WEIGHTS = PARASURF_ROOT / "Paragraph_expanded_entire_dataset_best.pth"


def _run(cmd: list[str], cwd: Path | None = None) -> None:
    print(f"[bootstrap] $ {' '.join(str(c) for c in cmd)}", flush=True)
    subprocess.run([str(c) for c in cmd], cwd=str(cwd) if cwd else None, check=True)


def _clone(url: str, dest: Path, ref: str | None = None) -> None:
    if (dest / ".git").exists():
        print(f"[bootstrap] {dest.name} already present", flush=True)
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    if ref and ref != "main":
        # A pinned commit may not be the tip, so fetch the history it needs.
        _run(["git", "clone", url, dest])
        _run(["git", "checkout", ref], cwd=dest)
    else:
        _run(["git", "clone", "--depth", "1", url, dest])


def _build_dms() -> None:
    """Compile DMS into PREFIX.

    ParaSurf's GNUmakefile bakes LIBDIR into the binary via -DDESTLIB, so a build
    with our own prefix looks for its radii file somewhere we can actually write.
    Verified to produce byte-identical surfaces to a root install.
    """
    dms_bin = PREFIX / "bin" / "dms"
    if dms_bin.exists():
        print("[bootstrap] dms already built", flush=True)
        return
    for sub in ("bin", "lib", "man/man1"):
        (PREFIX / sub).mkdir(parents=True, exist_ok=True)
    src = PARASURF_ROOT / "dms"
    # The repo ships prebuilt .o files from another toolchain; drop them first.
    for obj in src.glob("*.o"):
        obj.unlink()
    _run(["make", "install",
          f"BINDIR={PREFIX / 'bin'}",
          f"LIBDIR={PREFIX / 'lib'}",
          f"MANDIR={PREFIX / 'man/man1'}"], cwd=src)
    if not dms_bin.exists():
        raise RuntimeError("DMS build reported success but no binary was produced")


def _fix_pdb2pqr() -> None:
    """ParaSurf locates pdb2pqr by walking directories; it only needs +x."""
    for p in PARASURF_ROOT.rglob("pdb2pqr-linux-bin64-*/pdb2pqr"):
        p.chmod(p.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
        print(f"[bootstrap] chmod +x {p}", flush=True)


def _fetch_parasurf_weights() -> None:
    if PARASURF_WEIGHTS.exists() and PARASURF_WEIGHTS.stat().st_size > 0:
        print("[bootstrap] ParaSurf weights already present", flush=True)
        return
    import gdown

    print("[bootstrap] downloading ParaSurf weights (186 MB)…", flush=True)
    gdown.download(f"https://drive.google.com/uc?id={PARASURF_GDRIVE_ID}",
                   str(PARASURF_WEIGHTS), quiet=False)
    if not PARASURF_WEIGHTS.exists() or PARASURF_WEIGHTS.stat().st_size == 0:
        raise RuntimeError("ParaSurf weight download failed or produced an empty file")


def ensure_runtime() -> Path:
    """Assemble everything and put AntiSite on sys.path. Returns the repo root."""
    WORK.mkdir(parents=True, exist_ok=True)

    _clone(ANTISITE_REPO, ANTISITE_ROOT, ANTISITE_REF)
    _clone(PARASURF_REPO, PARASURF_ROOT)

    if shutil.which("make"):
        _build_dms()
    else:
        print("[bootstrap] WARNING: no 'make' — add build-essential to packages.txt; "
              "3D mode will fail", flush=True)

    _fix_pdb2pqr()

    # DMS is invoked as a bare `dms` via os.system, so it must be on PATH.
    bin_dir = str(PREFIX / "bin")
    if bin_dir not in os.environ.get("PATH", ""):
        os.environ["PATH"] = f"{bin_dir}:{os.environ.get('PATH', '')}"

    try:
        _fetch_parasurf_weights()
    except Exception as exc:  # noqa: BLE001 - sequence-only mode still works without it
        print(f"[bootstrap] ParaSurf weights unavailable ({exc}); 3D mode disabled",
              flush=True)

    # test_antisite/ holds infer_3d.py, which is not part of the installed package.
    for p in (ANTISITE_ROOT, ANTISITE_ROOT / "test_antisite"):
        if str(p) not in sys.path:
            sys.path.insert(0, str(p))

    print(f"[bootstrap] ready: {ANTISITE_ROOT}", flush=True)
    return ANTISITE_ROOT


if __name__ == "__main__":
    ensure_runtime()