import gradio as gr
from gradio_client import Client
import random
import time
import threading

# -------------------------
# Backend HF clients
# -------------------------
backend_engines = [
    Client("ChocoLaboratory/SANDBOX_BACKEND"),
    Client("ChocoLaboratory/SANDBOXBACKEND2")
]

SANDBOX_TTL_SEC = 2 * 3600  # 2 hours

# -------------------------
# Helper functions
# -------------------------
def pick_engine():
    return random.choice(backend_engines)

def launch_sandbox(code):
    engine = pick_engine()
    result = engine.predict(code=code, api_name="/launch_sandbox")
    sandbox_id = result if isinstance(result, str) else result.get("sandbox_id", "unknown")
    start_time = time.time()
    return f"Sandbox launched! ID: {sandbox_id}\n", sandbox_id, start_time, engine

def fetch_logs(sandbox_id, engine):
    if not sandbox_id or not engine:
        return "No sandbox ID. Launch a sandbox first!"
    try:
        logs = engine.predict(api_name="/fetch_logs")
        return logs
    except Exception as e:
        return f"Error fetching logs: {e}"

def kill_sandbox(sandbox_id, engine):
    if not sandbox_id or not engine:
        return "No sandbox ID to kill."
    try:
        result = engine.predict(api_name="/kill_sandbox")
        return result
    except Exception as e:
        return f"Error: {e}"

def status_sandbox(sandbox_id, start_time, engine):
    if not sandbox_id or not engine:
        return "No sandbox ID.", "0:00"
    try:
        status = engine.predict(api_name="/status_sandbox")
    except:
        status = "Unknown"
    elapsed = int(time.time() - start_time)
    remaining = max(0, SANDBOX_TTL_SEC - elapsed)
    minutes, seconds = divmod(remaining, 60)
    time_str = f"{minutes:02d}:{seconds:02d}"
    return status, time_str

def run_command(cmd, sandbox_id, start_time, engine):
    if not sandbox_id or not engine:
        return "No sandbox ID. Launch a sandbox first!", "0:00"
    try:
        engine.predict(code=cmd, api_name="/launch_sandbox")  # ideally /run_command
        logs, time_str = status_and_logs(sandbox_id, start_time, engine)
        return logs, time_str
    except Exception as e:
        return f"Error: {e}", "0:00"

def status_and_logs(sandbox_id, start_time, engine):
    logs = fetch_logs(sandbox_id, engine)
    _, time_str = status_sandbox(sandbox_id, start_time, engine)
    return logs, time_str

# -------------------------
# Auto-update logs thread
# -------------------------
def start_auto_update(terminal_output, sandbox_id_box, start_time_state, sandbox_backend_state, time_remaining):
    def updater():
        while True:
            sandbox_id = sandbox_id_box.value
            start_time = start_time_state.value
            engine = sandbox_backend_state.value
            if sandbox_id and engine:
                logs, time_str = status_and_logs(sandbox_id, start_time, engine)
                terminal_output.value = logs
                time_remaining.value = time_str
            time.sleep(1)
    threading.Thread(target=updater, daemon=True).start()

# -------------------------
# Gradio UI
# -------------------------
with gr.Blocks() as demo:

    with gr.Row():
        # Terminal Column
        with gr.Column(scale=3):
            gr.Markdown("### 🖥 Sandbox Cloud Console")
            terminal_output = gr.Textbox(label="", lines=20, interactive=False)
            command_input = gr.Textbox(label="Command Input", placeholder="echo SandBox!", interactive=True)
            run_btn = gr.Button("Run Command")
        
        # Side Panel
        with gr.Column(scale=1):
            gr.Markdown("### Sandbox Controls")
            launch_btn = gr.Button("Launch Sandbox")
            kill_btn = gr.Button("Kill Sandbox")
            status_btn = gr.Button("Check Status")
            sandbox_id_box = gr.Textbox(label="Sandbox ID", interactive=False)
            time_remaining = gr.Textbox(label="Time Remaining", interactive=False)
            gr.Markdown("💡 Running on a Debian-based HF container!")

    # States
    start_time_state = gr.State(value=0)
    sandbox_backend_state = gr.State(value=None)

    # Callbacks
    launch_btn.click(
        launch_sandbox,
        inputs=command_input,
        outputs=[terminal_output, sandbox_id_box, start_time_state, sandbox_backend_state]
    )

    run_btn.click(
        run_command,
        inputs=[command_input, sandbox_id_box, start_time_state, sandbox_backend_state],
        outputs=[terminal_output, time_remaining]
    )

    kill_btn.click(
        kill_sandbox,
        inputs=[sandbox_id_box, sandbox_backend_state],
        outputs=terminal_output
    )

    status_btn.click(
        status_sandbox,
        inputs=[sandbox_id_box, start_time_state, sandbox_backend_state],
        outputs=[terminal_output, time_remaining]
    )

    # Start auto-update logs thread
    start_auto_update(terminal_output, sandbox_id_box, start_time_state, sandbox_backend_state, time_remaining)

demo.launch()