"""
Quantum Applications - Unified Single-Server App

This app provides both EM Scattering and QLBM experiences in a single Trame server,
avoiding the multi-server/iframe approach that causes issues on HuggingFace Spaces.
"""
import os
import errno
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

# 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"

# --- 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 (must be done after server binding)
em.register_handlers()

# --- Build the Layout ---
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;"
    
    # Custom CSS
    trame_html.Style("""
    :root { --v-theme-primary: 95, 37, 159; }
    .landing-card:hover { transform: translateY(-4px); transition: transform 0.2s ease; }
    """)
    
    with layout.toolbar:
        vuetify3.VSpacer()
        
        # Back button (shown when in a sub-app)
        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 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:
        # === Landing Page ===
        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. Time-domain Maxwell equations are re-cast to Schrödinger-type equations and an equivalent Hamiltonian is computed, and evolved over time. "
                                "Configure geometry, excitation, and visualize field propagation.",
                                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 the classical Lattice Boltzmann method."
                                " Explore advection-diffusion with quantum-enhanced computation."
                                " Configure geometry, initial condition, and visualize field propagation.",
                                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",
        ):
            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()

# --- 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 heartbeat
    _start_hf_heartbeat_thread(interval_s=5)
    
    # Get port from environment (HuggingFace) or use default
    base_port = int(os.environ.get("APP_PORT") or os.environ.get("PORT", 7860))
    host = os.environ.get("APP_HOST", "0.0.0.0")
    max_attempts = 10

    print(f"Starting Quantum Applications server on {host}:{base_port}")
    print("This is a SINGLE SERVER serving both EM and QLBM experiences.")

    for attempt in range(max_attempts):
        port = base_port + attempt
        try:
            server.start(host=host, port=port, open_browser=False)
            break
        except OSError as exc:
            if getattr(exc, "errno", None) == errno.EADDRINUSE and attempt < max_attempts - 1:
                print(f"Port {port} busy, retrying on port {port + 1}...")
                continue
            raise