setup_zeopp.py

"""
setup_zeopp.py - Compile and install the Zeo++ 'network' binary at HF Space startup.

Called from app.py BEFORE ensemble loading:
    from setup_zeopp import ensure_zeopp_binary

This module is a no-op on non-Linux systems (Windows, macOS) and when the binary
already exists. It is safe to call on every app restart; compilation is skipped
if bin/network is already present.

Build strategy
--------------
Zeo++ 0.3 bundles Voro++ as a subdirectory. The build sequence is:
  1. wget the official tarball from zeoplusplus.org into a temp dir.
  2. Untar -> zeo++-0.3/
  3. cd zeo++-0.3/voro++/src && make -> libvoro++.a
  4. cd zeo++-0.3/ && make -> network binary
  5. Copy network -> bin/network  (relative to this file)
  6. chmod +x

packages.txt (apt) must list 'build-essential' so g++ and make are present
on the HF Space base image (Ubuntu 22.04).

Windows / macOS
---------------
On Windows, Zeo++ is not compiled automatically. The code skips the build and
returns False. build_graph.py already falls back to atom-only mode (n_pores=0)
when bin/network does not exist, so the app remains functional; pore nodes are
simply disabled. To enable pore nodes locally on Windows, install WSL2, compile
inside WSL, and copy the resulting 'network' binary to hf_space/bin/network.

Fallback
--------
If compilation fails for any reason (network unavailable, make error, timeout),
a WARNING is logged and the function returns False. build_graph.py detects the
missing binary and falls back to atom-only mode silently.
"""

from __future__ import annotations

import logging
import os
import platform
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import urllib.request

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

# Zeo++ 0.3 — official tarball from zeoplusplus.org
_ZEOPP_URL = 'http://www.zeoplusplus.org/zeo++-0.3.tar.gz'
_ZEOPP_URL_FALLBACK = (
    'https://raw.githubusercontent.com/richardjgowers/zeoplusplus/'
    'master/zeo%2B%2B-0.3.tar.gz'
)

# Where this file lives (hf_space/)
_HERE = os.path.dirname(os.path.abspath(__file__))

# Target binary location consumed by config.py / build_graph.py
_BIN_DIR = os.path.join(_HERE, 'bin')
_BINARY = os.path.join(_BIN_DIR, 'network')

# Compile timeout (seconds). Cold-start on HF free CPU should finish in ~3 min.
_COMPILE_TIMEOUT = 300


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

def ensure_zeopp_binary() -> bool:
    """
    Ensure bin/network exists and is executable.

    Returns True  if the binary is ready (either pre-existing or freshly compiled).
    Returns False if compilation was skipped (non-Linux) or failed.
    """
    # --- Already present ----------------------------------------------------
    if os.path.isfile(_BINARY) and os.access(_BINARY, os.X_OK):
        logger.info('Zeo++ binary already present at %s. Skipping build.', _BINARY)
        return True

    # --- Non-Linux: skip (no g++/make available without WSL) ----------------
    if platform.system() != 'Linux':
        logger.warning(
            'Zeo++ auto-compile is Linux-only. '
            'System is %s. Running in atom-only mode. '
            'To enable pore nodes locally, compile Zeo++ in WSL2 and copy '
            'the resulting "network" binary to hf_space/bin/network.',
            platform.system(),
        )
        return False

    # --- Check g++ and make are available ------------------------------------
    if shutil.which('g++') is None or shutil.which('make') is None:
        logger.warning(
            'g++ or make not found. Cannot compile Zeo++. '
            'Ensure packages.txt lists "build-essential". Atom-only mode.'
        )
        return False

    # --- Attempt compilation -------------------------------------------------
    logger.info('Compiling Zeo++ 0.3 from source. This takes ~1-3 minutes on first start.')

    os.makedirs(_BIN_DIR, exist_ok=True)

    with tempfile.TemporaryDirectory() as tmp:
        tarball = os.path.join(tmp, 'zeo.tar.gz')

        # Download
        if not _download(tarball):
            return False

        # Unpack
        try:
            with tarfile.open(tarball, 'r:gz') as tf:
                tf.extractall(tmp)
        except Exception as exc:
            logger.warning('Failed to unpack Zeo++ tarball: %s. Atom-only mode.', exc)
            return False

        # Locate unpacked directory (should be zeo++-0.3/)
        src_dir = os.path.join(tmp, 'zeo++-0.3')
        if not os.path.isdir(src_dir):
            # Fallback: search for any directory containing Makefile
            candidates = [
                d for d in os.listdir(tmp)
                if os.path.isdir(os.path.join(tmp, d))
                and os.path.isfile(os.path.join(tmp, d, 'Makefile'))
            ]
            if not candidates:
                logger.warning(
                    'Could not locate Zeo++ source directory after unpack. '
                    'Atom-only mode.'
                )
                return False
            src_dir = os.path.join(tmp, candidates[0])
            logger.info('Using Zeo++ source directory: %s', src_dir)

        # Step 1: compile voro++ (bundled dependency)
        voro_src = os.path.join(src_dir, 'voro++', 'src')
        if os.path.isdir(voro_src):
            if not _run_make(voro_src, label='voro++'):
                return False
        else:
            # Older tarballs put voro++ directly in the root
            logger.info('voro++/src not found; attempting top-level make directly.')

        # Step 2: compile zeo++ (produces 'network')
        if not _run_make(src_dir, label='zeo++'):
            return False

        # Step 3: copy binary
        built_binary = os.path.join(src_dir, 'network')
        if not os.path.isfile(built_binary):
            logger.warning(
                'Compilation appeared to succeed but "network" binary not found at %s. '
                'Atom-only mode.',
                built_binary,
            )
            return False

        shutil.copy2(built_binary, _BINARY)
        _make_executable(_BINARY)
        logger.info('Zeo++ binary installed at %s.', _BINARY)
        return True


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _download(dest: str) -> bool:
    """Locate Zeo++ tarball.

    Priority:
      1. Bundled tarball at hf_space/zeo++-0.3.tar.gz (preferred — no network).
      2. Primary URL on zeoplusplus.org.
      3. Fallback URL on github.

    HF Spaces outbound network has been observed to corrupt the tarball download
    on some runtimes (the file arrives larger than 10 KB but is not valid gzip),
    which is why bundling the tarball in the repo is the primary path.
    """
    # 1. Bundled local tarball
    local = os.path.join(_HERE, 'zeo++-0.3.tar.gz')
    if os.path.isfile(local) and os.path.getsize(local) > 1_000_000:
        try:
            shutil.copy2(local, dest)
            # Validate gzip magic
            with open(dest, 'rb') as f:
                head = f.read(2)
            if head == b'\x1f\x8b':
                logger.info('Using bundled Zeo++ tarball (%d bytes).', os.path.getsize(dest))
                return True
            logger.warning('Bundled tarball is not a valid gzip. Falling back to download.')
        except Exception as exc:
            logger.warning('Bundled tarball copy failed: %s. Falling back to download.', exc)

    # 2 & 3. Network fallback
    for url in (_ZEOPP_URL, _ZEOPP_URL_FALLBACK):
        try:
            logger.info('Downloading Zeo++ from %s ...', url)
            urllib.request.urlretrieve(url, dest)
            with open(dest, 'rb') as f:
                head = f.read(2)
            if head == b'\x1f\x8b' and os.path.getsize(dest) > 1_000_000:
                logger.info('Download complete: %d bytes.', os.path.getsize(dest))
                return True
            logger.warning('Downloaded file from %s is not a valid gzip. Trying next URL.', url)
        except Exception as exc:
            logger.warning('Download failed from %s: %s. Trying next URL.', url, exc)
    logger.warning('All Zeo++ sources failed. Atom-only mode.')
    return False


def _run_make(directory: str, label: str = '') -> bool:
    """Run 'make' in directory. Returns True on success."""
    try:
        result = subprocess.run(
            ['make'],
            cwd=directory,
            check=True,
            capture_output=True,
            text=True,
            timeout=_COMPILE_TIMEOUT,
        )
        logger.info('make %s succeeded.', label)
        if result.stdout:
            logger.debug('make %s stdout: %s', label, result.stdout[-500:])
        return True
    except subprocess.TimeoutExpired:
        logger.warning('make %s timed out after %ds. Atom-only mode.', label, _COMPILE_TIMEOUT)
    except subprocess.CalledProcessError as exc:
        logger.warning(
            'make %s failed (exit %d):\n%s\nAtom-only mode.',
            label, exc.returncode, (exc.stderr or '')[-1000:],
        )
    except Exception as exc:
        logger.warning('make %s unexpected error: %s. Atom-only mode.', label, exc)
    return False


def _make_executable(path: str) -> None:
    """chmod +x path."""
    st = os.stat(path)
    os.chmod(path, st.st_mode | stat.S_IEXEC | stat.S_IXGRP | stat.S_IXOTH)


# ---------------------------------------------------------------------------
# CLI usage (for debugging on HPC/WSL)
# ---------------------------------------------------------------------------
if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s')
    ok = ensure_zeopp_binary()
    sys.exit(0 if ok else 1)