app.py

import os
# Ensure OMP_NUM_THREADS is set to avoid libgomp errors
os.environ["OMP_NUM_THREADS"] = "1"

from trame.app import get_server
from trame_vuetify.ui.vuetify3 import SinglePageLayout
from trame_vuetify.widgets import vuetify3
from trame.widgets import html as trame_html
import threading
from concurrent.futures import ThreadPoolExecutor
import time

# Import embedded page wrappers (lazy load inside tabs)
from importlib import import_module
em_page = None
qlbm_page = None

# Force embedded mode for nested pages (avoid iframes/secondary servers)
os.environ.setdefault("TRAME_EMBEDDED", "1")
os.environ.setdefault("TRAME_DISABLE_BROWSER", "1")
os.environ.setdefault("EM_APP_EMBEDDED", "1")
os.environ.setdefault("DISABLE_EM_STANDALONE", "1")

# Create a single server for the multipage app
server = get_server()
state, ctrl = server.state, server.controller

# App state: landing chooser -> specific experience
state.current_page = None  # "EM" or "QLBM" once selected

# -------------------------------------------------------------------------
# Thread pool for long-running jobs to avoid UI freezing by HF timeout
# -------------------------------------------------------------------------
MAX_WORKERS = int(os.environ.get("MAX_WORKERS", "4"))
executor = ThreadPoolExecutor(max_workers=MAX_WORKERS)

def submit_background(fn, *args, **kwargs):
    """
    Run a CPU-heavy / long job without blocking Trame's event loop.

    Usage from pages:
        from trame.app import get_server
        server = get_server()
        server.submit_background(long_fn, arg1, arg2, kw1=...)
    """
    return executor.submit(fn, *args, **kwargs)

# Expose helper on the server so pages.em_page / pages.qlbm_page can use it
server.submit_background = submit_background

# Load Synopsys/Ansys logo as data URI for main toolbar
import base64

def _load_logo_data_uri():
    base_dir = os.path.dirname(__file__)
    candidates = [
        os.path.join(base_dir, "ansys-part-of-synopsys-logo.svg"),
        os.path.join(base_dir, "synopsys-logo-color-rgb.svg"),
        os.path.join(base_dir, "synopsys-logo-color-rgb.png"),
        os.path.join(base_dir, "synopsys-logo-color-rgb.jpg"),
    ]
    for p in candidates:
        if os.path.exists(p):
            ext = os.path.splitext(p)[1].lower()
            mime = "image/svg+xml" if ext == ".svg" else ("image/png" if ext == ".png" else "image/jpeg")
            with open(p, "rb") as f:
                b64 = base64.b64encode(f.read()).decode("ascii")
            return f"data:{mime};base64,{b64}"
    return None

def _stop_subapp(module_name: str, attr_name: str):
    module = globals().get(attr_name)
    if module is None:
        try:
            module = import_module(module_name)
            globals()[attr_name] = module
        except Exception:
            return
    stop_fn = getattr(module, "stop", None)
    if callable(stop_fn):
        try:
            stop_fn()
        except Exception:
            pass


# Safely initialize logo in state (trame state isn't a dict; avoid .get())
try:
    if not hasattr(state, "logo_src") or state.logo_src in (None, ""):
        state.logo_src = _load_logo_data_uri()
except Exception:
    state.logo_src = _load_logo_data_uri()


@state.change("current_page")
def _handle_page_change(current_page, **_):
    """Stop inactive subprocesses as users navigate between experiences."""
    if current_page == "EM":
        _stop_subapp("pages.qlbm_page", "qlbm_page")
    elif current_page == "QLBM":
        _stop_subapp("pages.em_page", "em_page")
    else:
        _stop_subapp("pages.em_page", "em_page")
        _stop_subapp("pages.qlbm_page", "qlbm_page")

with SinglePageLayout(server) as layout:
    layout.title.set_text("Quantum Applications")
    layout.toolbar.classes = "pl-2 pr-1 py-1 elevation-0"
    layout.toolbar.style = "background-color: #ffffff; border-bottom: 3px solid #5f259f;"

    with layout.toolbar:
        vuetify3.VSpacer()
        vuetify3.VBtn(
            v_if="current_page",
            text="Main Page",
            variant="text",
            color="primary",
            prepend_icon="mdi-arrow-left",
            click="current_page = null",
            classes="mr-2",
        )
        vuetify3.VChip(
            v_if="current_page",
            label=True,
            color="primary",
            text_color="white",
            children=["{{ current_page === 'EM' ? 'Electromagnetic Scattering' : 'Quantum LBM' }}"],
            classes="mr-2",
        )
        vuetify3.VImg(
            v_if="logo_src",
            src=("logo_src", None),
            style="height: 40px; width: auto;",
            classes="ml-2",
        )

    with layout.content:
        # Landing screen
        with vuetify3.VContainer(
            v_if="!current_page",
            fluid=True,
            classes="fill-height d-flex align-center justify-center pa-6",
        ):
            with vuetify3.VSheet(
                elevation=6,
                rounded=True,
                style="max-width: 1080px; width: 100%; background: linear-gradient(135deg, #fdfbff, #f3ecff);",
                classes="pa-8",
            ):
                vuetify3.VCardTitle(
                    "Pick a quantum experience",
                    classes="text-h4 text-primary font-weight-bold mb-2 text-center",
                )
                vuetify3.VCardSubtitle(
                    "Choose one workflow. We'll spin up only that server until you switch back.",
                    classes="text-body-1 text-center mb-6",
                )
                with vuetify3.VRow(justify="center", align="stretch", class_="text-left"):
                    with vuetify3.VCol(cols=12, md=5, class_="d-flex"):
                        with vuetify3.VCard(elevation=4, classes="pa-6 flex-grow-1"):
                            vuetify3.VIcon("mdi-radar", size=52, color="primary", classes="mb-4")
                            vuetify3.VCardTitle("Electromagnetic Scattering", classes="text-h5 mb-2")
                            vuetify3.VCardText(
                                "Placeholder",
                                classes="text-body-2 mb-6",
                            )
                            vuetify3.VBtn(
                                text="Launch EM",
                                color="primary",
                                block=True,
                                prepend_icon="mdi-play-circle",
                                size="large",
                                click="current_page = 'EM'",
                            )
                    with vuetify3.VCol(cols=12, md=5, class_="d-flex"):
                        with vuetify3.VCard(elevation=4, classes="pa-6 flex-grow-1"):
                            vuetify3.VIcon("mdi-water", size=52, color="secondary", classes="mb-4")
                            vuetify3.VCardTitle("Fluids", classes="text-h5 mb-2")
                            vuetify3.VCardText(
                                "Placeholder",
                                classes="text-body-2 mb-6",
                            )
                            vuetify3.VBtn(
                                text="Launch QLBM",
                                color="secondary",
                                block=True,
                                prepend_icon="mdi-play-circle",
                                size="large",
                                click="current_page = 'QLBM'",
                            )

        # EM experience
        with vuetify3.VContainer(v_if="current_page === 'EM'", fluid=True, classes="pa-0 fill-height"):
            try:
                if not globals().get("em_page"):
                    globals()["em_page"] = import_module("pages.em_page")
                em_page.build(server)
            except Exception as e:
                trame_html.Div(f"EM embed failed: {e}", style="padding:8px;color:#b00020;")

        # QLBM experience
        with vuetify3.VContainer(v_if="current_page === 'QLBM'", fluid=True, classes="pa-0 fill-height"):
            try:
                if not globals().get("qlbm_page"):
                    globals()["qlbm_page"] = import_module("pages.qlbm_page")
                qlbm_page.build(server)
            except Exception as e:
                trame_html.Div(f"QLBM embed failed: {e}", style="padding:8px;color:#b00020;")

# -------------------------------------------------------------------------
# Heartbeat: keep HuggingFace WebSocket alive during idle periods
# -------------------------------------------------------------------------
def _start_hf_heartbeat_thread(interval_s: int = 5):
    """
    Start a background thread that periodically flushes the server,
    keeping the WebSocket "active" in the eyes of Hugging Face.
    """
    def _loop():
        while True:
            time.sleep(interval_s)
            try:
                server.controller.flush()
            except Exception:
                # If server is shutting down or flush fails, exit the thread
                break

    t = threading.Thread(target=_loop, daemon=True)
    t.start()


if __name__ == "__main__":

    # Start HF heartbeat to prevent timeout
    _start_hf_heartbeat_thread(interval_s=5)

    # Allow reverse-proxy setups to pin the internal host/port independently from the public PORT
    port = int(os.environ.get("APP_PORT") or os.environ.get("PORT", 7860))
    host = os.environ.get("APP_HOST", "0.0.0.0")
    server.start(host=host, port=port, open_browser=False)