app.py

# ========= Headless / Offscreen safety (before any VTK import) =========
import os
# Set headless rendering environment variables BEFORE any VTK/pyvista imports
os.environ.setdefault("VTK_DEFAULT_RENDER_WINDOW_OFFSCREEN", "1")
os.environ.setdefault("LIBGL_ALWAYS_SOFTWARE", "1")
os.environ.setdefault("MESA_GL_VERSION_OVERRIDE", "3.3")
# Use dummy DISPLAY value instead of empty string to avoid X server connection attempts
if "DISPLAY" not in os.environ or os.environ.get("DISPLAY") == "":
    # Disable all X11 attempts
    os.environ["DISPLAY"] = ":0"
    os.environ["VTK_USE_X"] = "NO"

os.environ.setdefault("PYVISTA_OFF_SCREEN", "True")

import sys
import tempfile
from pathlib import Path
import traceback
import signal


# Writable base dir (matches alphaLPFM pattern)
# Use APP_DATA_DIR to match alphaLPFM's expectation
DATA_DIR = os.environ.get("APP_DATA_DIR", os.path.join(tempfile.gettempdir(), "appdata_trame"))
os.makedirs(DATA_DIR, exist_ok=True)
os.environ.setdefault("APP_DATA_DIR", DATA_DIR)  # Ensure it's set for the private app
os.environ.setdefault("MPLCONFIGDIR", DATA_DIR)

# Set up shared directories structure (matches alphaLPFM)
SHARED_DIR = os.path.join(DATA_DIR, "shared")
SHARED_HF_DIR = os.path.join(SHARED_DIR, "hf_home")
SHARED_WEIGHTS_DIR = os.path.join(SHARED_DIR, "weights")
for d in (SHARED_DIR, SHARED_HF_DIR, SHARED_WEIGHTS_DIR):
    os.makedirs(d, exist_ok=True)

# Set HF cache environment variables (matches alphaLPFM)
os.environ.setdefault("HF_HOME", SHARED_HF_DIR)
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", SHARED_HF_DIR)
os.environ.setdefault("TRANSFORMERS_CACHE", SHARED_HF_DIR)

def setup_environment():
    cache_dir = Path("cache")
    cache_dir.mkdir(exist_ok=True)
    if str(cache_dir.absolute()) not in sys.path:
        sys.path.insert(0, str(cache_dir.absolute()))
    return cache_dir

def download_private_repo(cache_dir, repo_id, hf_token):
    from huggingface_hub import snapshot_download
    import shutil
    import os
    print(f"🔍 Downloading repo {repo_id}")
    repo_path = snapshot_download(
        repo_id=repo_id,
        token=hf_token,
        repo_type="space",
        local_dir=cache_dir,
        local_dir_use_symlinks=False,
    )
    
    # Copy everything from the repository to current working directory
    print(f"📁 Copying all files and directories from {repo_path} to current directory")
    
    # Get all items in the repository
    for item in os.listdir(repo_path):
        src_path = os.path.join(repo_path, item)
        dst_path = item
        
        # Skip hidden files and directories (like .git, .gitignore, etc.)
        if item.startswith('.'):
            print(f"⏭️ Skipping hidden item: {item}")
            continue
            
        if os.path.isdir(src_path):
            # print(f"📁 Copying directory: {item}")
            if os.path.exists(dst_path):
                shutil.rmtree(dst_path)
            shutil.copytree(src_path, dst_path)
            # print(f"✅ Directory {item} copied successfully")
        elif os.path.isfile(src_path):
            # print(f"📄 Copying file: {item}")
            shutil.copy2(src_path, dst_path)
            # print(f"✅ File {item} copied successfully")
    
    print(f"🎉 All repository contents copied to current working directory")
    
    return repo_path

# Cache for loaded private app
_PRIVATE_APP_CACHE = None
_PRIVATE_APP_CREATE_APP = None

def _load_private_app_once():
    """Load the private app once and cache it. Called by create_app() or main()."""
    global _PRIVATE_APP_CACHE, _PRIVATE_APP_CREATE_APP
    
    if _PRIVATE_APP_CACHE is not None:
        return _PRIVATE_APP_CACHE, _PRIVATE_APP_CREATE_APP
    
    hf_token = os.environ.get("HF_TOKEN")
    repo_id = os.environ.get("REPO_ID")

    if not hf_token or not repo_id:
        raise ValueError("HF_TOKEN and REPO_ID must be set")

    print(f"🔍 Loading private app from repo_id: {repo_id}")
    cache_dir = setup_environment()
    repo_path = download_private_repo(cache_dir, repo_id, hf_token)
    
    # Import from repo_path directly using importlib to avoid importing wrapper's app.py
    try:
        import importlib.util
        app_file_path = os.path.join(repo_path, "app.py")
        if not os.path.exists(app_file_path):
            raise FileNotFoundError(f"app.py not found in {repo_path}")
        
        spec = importlib.util.spec_from_file_location("private_app_module", app_file_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Could not load spec from {app_file_path}")
        
        private_app = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(private_app)
        # If the private app exposes a thumbnail generator, run it once
        # so example images are ready even when we bypass its main().
        try:
            if hasattr(private_app, "generate_example_thumbnails"):
                print("🖼  Generating example thumbnails in private app...")
                private_app.generate_example_thumbnails()
        except Exception:
            # Thumbnails are optional; don't break startup if this fails.
            traceback.print_exc()

        # ✅ PREFER create_app factory (for multi-user session management)
        if hasattr(private_app, "create_app"):
            print("✅ Found create_app factory in private app (supports multi-user sessions)")
            _PRIVATE_APP_CREATE_APP = private_app.create_app
        # Check for app or demo attribute
        elif hasattr(private_app, "app"):
            _PRIVATE_APP_CREATE_APP = lambda: private_app.app
        elif hasattr(private_app, "demo"):
            _PRIVATE_APP_CREATE_APP = lambda: private_app.demo
        # Check if PFMDemo class exists and instantiate it (fallback, but not ideal for multi-user)
        elif hasattr(private_app, "PFMDemo"):
            print("⚠️  Using PFMDemo directly (may not support multi-user sessions)")
            def _create_pfmdemo_app(server=None):
                from trame.app import get_server
                if server is None:
                    server = get_server()
                app_instance = private_app.PFMDemo(server)
                app_instance.server.controller.add("decimate_again", app_instance.decimate_again)
                app_instance.server.controller.add("reset_mesh", app_instance.reset_mesh)
                return app_instance
            _PRIVATE_APP_CREATE_APP = _create_pfmdemo_app
        else:
            raise RuntimeError("No demo/app/create_app found in private repo")
        
        # Also add repo_path to sys.path for any relative imports the private app might need
        if str(repo_path) not in sys.path:
            sys.path.insert(0, str(repo_path))
        
        _PRIVATE_APP_CACHE = (private_app, repo_path)
        return _PRIVATE_APP_CACHE, _PRIVATE_APP_CREATE_APP
    except Exception as e:
        traceback.print_exc()
        raise RuntimeError(f"Failed to load app from private repo: {e}")

def create_app(server=None, **kwargs):
    """
    Trame application factory (matches alphaLPFM pattern).
    
    This function is called by trame.tools.serve to create the app instance.
    It loads the private app and returns its create_app result (or wraps it).
    """
    _, create_app_fn = _load_private_app_once()
    
    if create_app_fn is None:
        raise RuntimeError("No create_app function found in private app")
    
    # Call the private app's create_app (or wrapper)
    return create_app_fn(server=server, **kwargs)

def main():
    """
    Main entry point (fallback if not using trame.tools.serve).
    This is kept for backward compatibility but trame.tools.serve is preferred.
    """
    # Setup signal handlers for graceful shutdown
    def signal_handler(sig, frame):
        print("\n🛑 Received shutdown signal, cleaning up...")
        sys.exit(0)
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        print(f"🚀 Starting app (using main() - consider using trame.tools.serve for multi-user support)")
        print(f"📂 Current working directory: {os.getcwd()}")
        
        # Load private app (will be cached)
        _, create_app_fn = _load_private_app_once()
        
        if create_app_fn is None:
            raise RuntimeError("No create_app function found in private app")
        
        # Create app instance
        demo = create_app()

        # Check if it's a Trame app (has server.start method)
        if hasattr(demo, 'server') and hasattr(demo.server, 'start'):
            # Trame app
            print(f"🚀 Starting Trame server (host/port from trame.tools.serve or defaults)")
            
            if hasattr(demo.server.state, 'health') and demo.server.state.health == "running":
                print("Trame server already running, not starting again.")
                return

            # Configure server - host/port come from trame.tools.serve or environment
            try:
                # Try with host parameter first (if supported)
                demo.server.start(
                    open_browser=False,
                    show_connection_info=True,
                    backend="aiohttp",
                    exec_mode="main",
                )
            except TypeError:
                # If host parameter is not supported, configure via server options
                # Most Trame versions bind to 0.0.0.0 by default with aiohttp
                demo.server.start(
                    open_browser=False,
                    show_connection_info=True,
                    backend="aiohttp",
                    exec_mode="main",
                )
        elif hasattr(demo, 'launch'):
            # Gradio app
            _, repo_path = _PRIVATE_APP_CACHE if _PRIVATE_APP_CACHE else (None, None)
            demo.launch(
                server_name="0.0.0.0",
                server_port=7860,
                share=False,
                show_error=True,
                allowed_paths=[str(repo_path) if repo_path else ".", DATA_DIR]
            )
        else:
            raise RuntimeError("App does not have launch() or server.start() method")
    except KeyboardInterrupt:
        print("\n🛑 Keyboard interrupt received, shutting down...")
        sys.exit(0)
    except Exception as e:
        print(f"❌ Fatal error: {e}")
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()