# trame_timer_full_demo_v2.py
# Trame (Vuetify v2) – robust, reusable ticking timers for any long-running task

import time
import threading
from trame.app import get_server
from trame.ui.vuetify import SinglePageLayout
from trame.widgets import vuetify as v

# -----------------------------------------------------------------------------
# Server / State
# -----------------------------------------------------------------------------

server = get_server()  # Vuetify v2 stack
state = server.state

# Upload stage
state.is_uploading = False
state.upload_elapsed = 0.0
state.upload_progress = 0               # 0..100
state.upload_msg = ""

# Predict stage
state.is_predicting = False
state.predict_elapsed = 0.0
state.predict_progress = 0              # 0..100
state.predict_msg = ""

def _flush_safe():
    try:
        state.flush()
    except Exception:
        # Prefer logging in real apps; avoid crashing worker thread
        pass

# -----------------------------------------------------------------------------
# Generic, reusable timer helpers
# -----------------------------------------------------------------------------

def start_ticker(flag_key: str, elapsed_key: str, interval_s: float = 0.1):
    """
    Start a thread that updates <elapsed_key> while <flag_key> is True.
    """
    state[elapsed_key] = 0.0
    state[flag_key] = True
    _flush_safe()

    def _tick():
        t0 = time.perf_counter()
        while state.get(flag_key, False):
            state[elapsed_key] = time.perf_counter() - t0
            _flush_safe()
            time.sleep(interval_s)

    threading.Thread(target=_tick, daemon=True).start()

def stop_ticker(flag_key: str):
    state[flag_key] = False
    _flush_safe()

def run_with_timer(flag_key: str, elapsed_key: str, work_fn, *args, **kwargs):
    """
    Fire-and-forget: starts ticker, runs work_fn(*args, **kwargs) in a thread,
    and always stops ticker in a finally-block. Returns the worker thread.
    """
    start_ticker(flag_key, elapsed_key)

    def _runner():
        try:
            work_fn(*args, **kwargs)
        finally:
            stop_ticker(flag_key)

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

# -----------------------------------------------------------------------------
# Example “work” functions (replace with your real code)
# -----------------------------------------------------------------------------

def _simulate_upload_work():
    """
    Simulates an upload: increments progress 0..100 and sets a message.
    If cancelled (is_uploading False), it exits cleanly.
    """
    state.upload_msg = "Uploading geometry..."
    state.upload_progress = 0
    _flush_safe()

    for i in range(101):
        if not state.is_uploading:
            state.upload_msg = "Upload cancelled."
            _flush_safe()
            return
        state.upload_progress = i
        _flush_safe()
        time.sleep(0.05)  # Simulate IO chunks

    state.upload_msg = "Upload completed."
    _flush_safe()

def _simulate_predict_work():
    """
    Simulates prediction: increments progress 0..100 and sets a message.
    """
    state.predict_msg = "Running prediction..."
    state.predict_progress = 0
    _flush_safe()

    for i in range(101):
        if not state.is_predicting:
            state.predict_msg = "Prediction cancelled."
            _flush_safe()
            return
        state.predict_progress = i
        _flush_safe()
        time.sleep(0.06)  # Simulate compute steps

    state.predict_msg = "Prediction completed."
    _flush_safe()

# -----------------------------------------------------------------------------
# Controllers (button hooks)
# -----------------------------------------------------------------------------

def start_upload():
    # Reset UI bits for a fresh run
    state.upload_progress = 0
    state.upload_msg = ""
    _flush_safe()
    run_with_timer("is_uploading", "upload_elapsed", _simulate_upload_work)

def cancel_upload():
    # Just drop the flag; worker and ticker will stop on next loop
    stop_ticker("is_uploading")

def start_predict():
    state.predict_progress = 0
    state.predict_msg = ""
    _flush_safe()
    run_with_timer("is_predicting", "predict_elapsed", _simulate_predict_work)

def cancel_predict():
    stop_ticker("is_predicting")

# -----------------------------------------------------------------------------
# UI (Vuetify v2)
# -----------------------------------------------------------------------------

with SinglePageLayout(server) as layout:
    layout.title.set_text("Trame Live Timer – Upload & Predict")

    with layout.toolbar:
        v.VBtn("Start Upload", color="primary", click=start_upload)
        v.VBtn("Cancel Upload", color="error", click=cancel_upload, classes="ml-2")
        v.VDivider(vertical=True, classes="mx-4")
        v.VBtn("Start Predict", color="secondary", click=start_predict)
        v.VBtn("Cancel Predict", color="error", click=cancel_predict, classes="ml-2")

    with layout.content:
        # Upload panel
        v.VSheet(
            class_="pa-4 mx-auto",
            style="max-width: 600px; margin-top: 40px;"
        )(
            v.VSubheader("Upload Stage"),
            # Progress bar (determinate) – remove v_model & set indeterminate=True if duration unknown
            v.VProgressLinear(
                v_model=("upload_progress", 0),
                height=14,
                color="primary",
                rounded=True,
                v_show=("is_uploading",),
            ),
            v.VSpacer(style="min-height:10px;"),
            # Live elapsed time
            v.VLabel(
                "Elapsed: {{ upload_elapsed.toFixed(2) }} s",
                v_show=("is_uploading",),
                style="font-weight:600; font-size:14px; color:#1976d2;",
            ),
            v.VSpacer(style="min-height:6px;"),
            v.VLabel(("{{ upload_msg }}",), style="opacity:0.8;"),
        )

        # Predict panel
        v.VSheet(
            class_="pa-4 mx-auto",
            style="max-width: 600px; margin-top: 20px;"
        )(
            v.VSubheader("Prediction Stage"),
            v.VProgressLinear(
                v_model=("predict_progress", 0),
                height=14,
                color="secondary",
                rounded=True,
                v_show=("is_predicting",),
            ),
            v.VSpacer(style="min-height:10px;"),
            v.VLabel(
                "Elapsed: {{ predict_elapsed.toFixed(2) }} s",
                v_show=("is_predicting",),
                style="font-weight:600; font-size:14px;",
            ),
            v.VSpacer(style="min-height:6px;"),
            v.VLabel(("{{ predict_msg }}",), style="opacity:0.8;"),
        )

# -----------------------------------------------------------------------------
# Run
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    server.start()