app.py

"""
Quantum Applications - Unified Single-Server App with Simplified Session Management

This app provides EM scattering and QLBM pages in a single Trame server
with simplified session management based on filesystem directories and persistent indices.
"""
import os
import errno
import uuid
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
import time
import base64

from simple_session_manager import SESSION_MANAGER, get_session_by_alias

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

# --- App state
state.current_page = None  # None = landing, "EM" or "QLBM"
state.session_active = False
state.session_alias = ""
state.current_session_id = None
# Session card UI state
state.session_card_visible = True
state.session_alias_input = ""
state.session_action_success = False
state.session_action_busy = False
state.session_error = ""
state.session_action_trigger = None  # "create" or "load" when user clicks buttons


# --- Logo Loading ---
def _load_logo_data_uri():
    base_dir = os.path.dirname(__file__)
    candidates = [
        os.path.join(base_dir, "ansys-part-of-synopsys-logo.svg")
    ]
    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

state.logo_src = _load_logo_data_uri()

# --- Import Embedded Modules ---
# These are lightweight modules that build UI without creating their own servers
import qlbm_embedded
import em  # Use the modular em package instead of em_embedded

# Set the shared server on both modules
qlbm_embedded.set_server(server)
em.set_server(server)

# Initialize state for both modules
qlbm_embedded.init_state()
em.init_state()

# Register EM handlers
em.register_handlers()

# --- Session Callbacks ---
def on_session_selected(session_id: str):
    """Callback when user selects a session from landing page."""
    try:
        session_info = SESSION_MANAGER.get_session(session_id=session_id)
        if not session_info:
            state.session_error = f"Session not found: {session_id}"
            return
        
        state.current_session_id = session_id
        state.session_alias = session_info.get("alias", "Untitled")
        state.session_active = True
        print(f"[APP] Session selected: {state.session_alias} ({session_id})")
    except Exception as e:
        print(f"[APP] Error selecting session: {e}")
        state.session_error = str(e)
        state.session_active = False
# Watcher to reset navigation when no session is active
@state.change("session_active")
def _on_session_status_change(session_active, **kwargs):
    """Reset landing page if session becomes inactive."""
    if not session_active and state.current_page is not None:
        state.current_page = None
        state.session_card_visible = True
        print("[APP] Session inactive; reset to landing page")


# --- Session action watchers (respond to state changes from UI) ---
@state.change("session_action_trigger")
def _on_session_action(session_action_trigger, **kwargs):
    """Handle session create/load when user clicks buttons."""
    if session_action_trigger == "create":
        try:
            state.session_action_busy = True
            alias = state.session_alias_input.strip() if state.session_alias_input else f"session-{uuid.uuid4().hex[:6]}"
            
            # Create new session
            session_info = SESSION_MANAGER.create_session(alias=alias)
            session_id = session_info["session_id"]
            
            state.session_alias = session_info.get("alias", "Untitled")
            state.session_active = True
            state.current_session_id = session_id
            state.current_page = None
            state.session_alias_input = ""
            state.session_action_success = True
            state.session_error = ""
            print(f"[APP] Session created: {alias} ({session_id})")
            
            # Hide card after success animation
            def _hide():
                time.sleep(1.5)
                state.session_card_visible = False
                state.session_action_success = False
                state.session_action_busy = False
                state.session_action_trigger = None
            
            threading.Thread(target=_hide, daemon=True).start()
        except Exception as e:
            state.session_error = str(e)
            state.session_action_busy = False
            state.session_action_trigger = None
            print(f"[APP] Error creating session: {e}")

    elif session_action_trigger == "load":
        try:
            state.session_action_busy = True
            alias = state.session_alias_input.strip()
            if not alias:
                state.session_error = "Please enter an alias to load."
                state.session_action_busy = False
                state.session_action_trigger = None
                return

            # Load session by alias
            session_info = get_session_by_alias(alias)
            if not session_info:
                state.session_error = f"No session found with alias '{alias}'."
                state.session_action_busy = False
                state.session_action_trigger = None
                return

            session_id = session_info["session_id"]
            state.session_alias = session_info.get("alias", "Untitled")
            state.session_active = True
            state.current_session_id = session_id
            state.current_page = None
            state.session_alias_input = ""
            state.session_action_success = True
            state.session_error = ""
            print(f"[APP] Session loaded: {alias} ({session_id})")

            def _hide_load():
                time.sleep(1.5)
                state.session_card_visible = False
                state.session_action_success = False
                state.session_action_busy = False
                state.session_action_trigger = None

            threading.Thread(target=_hide_load, daemon=True).start()
        except Exception as e:
            state.session_error = str(e)
            state.session_action_busy = False
            state.session_action_trigger = None
            print(f"[APP] Error loading session: {e}")


# --- Build the Layout ---
with SinglePageLayout(server) as layout:
    layout.title.set_text("dummyquantum - Experiment")
    layout.toolbar.classes = "pl-2 pr-1 py-1 elevation-0"
    layout.toolbar.style = "background-color: #ffffff; border-bottom: 3px solid #5f259f;"
    
    # Custom CSS
    trame_html.Style("""
    :root { --v-theme-primary: 95, 37, 159; }
    .landing-card:hover { transform: translateY(-4px); transition: transform 0.2s ease; }
        @keyframes successBounce {
            0% { transform: scale(0.85); opacity: 0; }
            50% { transform: scale(1.12); opacity: 1; }
            100% { transform: scale(1); opacity: 1; }
        }
        .success-bounce { animation: successBounce 0.6s ease; color: #1e9e3e; }
    """)
    
    with layout.toolbar:
        vuetify3.VSpacer()
        
        # Back button (shown when in a sub-app)
        @ctrl.add("return_to_landing")
        def return_to_landing():
            """Return to landing page."""
            state.current_page = None
        
        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",
        )
        
        # Current session indicator
        vuetify3.VChip(
            v_if="session_active",
            label=True,
            color="info",
            text_color="white",
            children=["Session: {{ session_alias }}"],
            size="small",
            classes="mr-2",
        )
        
        # Current page indicator
        vuetify3.VChip(
            v_if="current_page",
            label=True,
            color="primary",
            text_color="white",
            children=["{{ current_page === 'EM' ? 'Electromagnetic Scattering' : 'Quantum LBM' }}"],
            classes="mr-2",
        )
        
        # Logo
        vuetify3.VImg(
            v_if="logo_src",
            src=("logo_src", None),
            style="height: 40px; width: auto;",
            classes="ml-2",
        )

    
    with layout.content:
        # === Simple Landing Page (compact) ===
        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(
                    "Quantum Applications Hub",
                    classes="text-h4 text-primary font-weight-bold mb-2 text-center",
                )
                vuetify3.VCardSubtitle(
                    "Choose a quantum CAE app",
                    classes="text-body-1 text-center mb-6",
                )
                
                with vuetify3.VRow(justify="center", align="stretch", class_="text-left"):
                    # EM Card
                    with vuetify3.VCol(cols=12, md=5, class_="d-flex"):
                        with vuetify3.VCard(elevation=4, classes="pa-6 flex-grow-1 landing-card"):
                            vuetify3.VIcon("mdi-radar", size=52, color="primary", classes="mb-4")
                            vuetify3.VCardTitle("Electromagnetic Scattering", classes="text-h5 mb-2")
                            vuetify3.VCardText(
                                    "Simulate electromagnetic wave scattering using quantum Hamiltonian Simulation.",
                                    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'",
                            )
                    
                    # QLBM Card
                    with vuetify3.VCol(cols=12, md=5, class_="d-flex"):
                        with vuetify3.VCard(elevation=4, classes="pa-6 flex-grow-1 landing-card"):
                            vuetify3.VIcon("mdi-water", size=52, color="secondary", classes="mb-4")
                            vuetify3.VCardTitle("Fluids", classes="text-h5 mb-2")
                            vuetify3.VCardText(
                                "3D fluid simulation using a quantum analog of the classical Lattice Boltzmann method.",
                                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'",
                            )

            # Small floating corner card for session management (create / load)
            with vuetify3.VContainer(v_if="session_card_visible", classes="pa-0", style="position: absolute; right: 16px; bottom: 16px; max-width: 320px;"):
                with vuetify3.VCard(elevation=8, classes="pa-3", style="background: rgba(255,255,255,0.98); width: 100%;"):
                    vuetify3.VCardTitle("Session", classes="text-subtitle-1 mb-1")
                    vuetify3.VCardText("Load an existing session by alias or create a new one.", classes="text-body-2 mb-3")

                    vuetify3.VTextField(label="Session Alias", v_model=("session_alias_input", None), placeholder="my-session")

                    # App selection moved to main landing page; sessions may hold data for both apps.

                    with vuetify3.VRow(gutter="8"):
                        with vuetify3.VCol(cols=6):
                            vuetify3.VBtn(text="Create", color="primary", block=True, disabled=("session_action_busy", None), click="session_action_trigger = 'create'")
                        with vuetify3.VCol(cols=6):
                            vuetify3.VBtn(text="Load", color="secondary", block=True, disabled=("session_action_busy", None), click="session_action_trigger = 'load'")

                    # Success / error feedback
                    vuetify3.VRow()
                    with vuetify3.VCol():
                        vuetify3.VIcon(v_if="session_action_success", children=["mdi-check-circle"], class_="success-bounce", size=36)
                        vuetify3.VAlert(v_if="session_error", type="error", dense=True, text=True, children=["{{ session_error }}"]) 

        # === EM Experience ===
        with vuetify3.VContainer(
            v_if="current_page === 'EM'",
            fluid=True,
            classes="pa-0 fill-height",
        ):
            em.build_ui()
        
        # === QLBM Experience ===
        with vuetify3.VContainer(
            v_if="current_page === 'QLBM'",
            fluid=True,
            classes="pa-0 fill-height",
        ):
            qlbm_embedded.build_ui()

# Enable point picking after UI is built (prevents KeyError with Trame state)
em.enable_point_picking_on_plotter()

# Reset to landing page on every server startup (so fresh browser loads start at landing)
state.current_page = None
state.session_card_visible = True
print("[OK] Landing page state initialized on startup")

# --- Heartbeat for HuggingFace ---
def _start_hf_heartbeat_thread(interval_s: int = 5):
    """Keep the WebSocket alive for HuggingFace Spaces."""
    import asyncio
    
    def _loop():
        while True:
            time.sleep(interval_s)
            try:
                # Create a new event loop for this thread if needed
                try:
                    loop = asyncio.get_event_loop()
                except RuntimeError:
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                
                # Try to flush, but don't crash if it fails
                if hasattr(server, 'controller') and hasattr(server.controller, 'flush'):
                    server.controller.flush()
            except Exception:
                # Silently continue - heartbeat is optional
                pass
    
    t = threading.Thread(target=_loop, daemon=True, name="HeartbeatThread")
    t.start()


# --- Entry Point ---
if __name__ == "__main__":
    _start_hf_heartbeat_thread(interval_s=5)

    # Check if running on HF Spaces with launcher
    import sys
    if "--launcher" in sys.argv or os.environ.get("HF_SPACE_ID"):
        # Running on HF Spaces with multi-user launcher
        # The launcher will manage user isolation and sessions
        server.start(open_browser=False)
    else:
        # Local development mode
        host = "0.0.0.0"
        port = int(os.environ.get("PORT", "7860"))
        server.start(host=host, port=port, open_browser=False)