import os
import uuid
import subprocess
import resource
import multiprocessing
import signal
import time
import gradio as gr

# Store running sandboxes
sandboxes = {}

# -------------------------
# Sandbox worker function
# -------------------------
def sandbox_worker(code: str, sandbox_id: str):
    user_dir = f"/tmp/sandbox_{sandbox_id}"
    os.makedirs(user_dir, exist_ok=True)
    os.chdir(user_dir)

    # Resource limits
    resource.setrlimit(resource.RLIMIT_AS, (4 * 1024**3, 4 * 1024**3))  # 4GB RAM
    resource.setrlimit(resource.RLIMIT_CPU, (7200, 7200))  # 2h CPU

    # Save user code
    code_file = os.path.join(user_dir, "user_code.py")
    with open(code_file, "w") as f:
        f.write(code)

    # Run code
    try:
        result = subprocess.run(
            ["python3", code_file],
            capture_output=True,
            text=True,
            timeout=7200  # 2h wall-clock
        )
        # Save logs
        with open(os.path.join(user_dir, "stdout.log"), "w") as f:
            f.write(result.stdout)
        with open(os.path.join(user_dir, "stderr.log"), "w") as f:
            f.write(result.stderr)
    except subprocess.TimeoutExpired:
        with open(os.path.join(user_dir, "stderr.log"), "w") as f:
            f.write("Execution timed out (2h limit reached).")

# -------------------------
# Gradio interface functions
# -------------------------
def launch_sandbox(code):
    sandbox_id = str(uuid.uuid4())[:8]
    p = multiprocessing.Process(target=sandbox_worker, args=(code, sandbox_id))
    p.start()
    sandboxes[sandbox_id] = {"pid": p.pid, "start_time": time.time()}
    return f"Sandbox launched! ID: {sandbox_id}", sandbox_id

def fetch_logs(sandbox_id):
    user_dir = f"/tmp/sandbox_{sandbox_id}"
    if not os.path.exists(user_dir):
        return "Sandbox not found."
    stdout_file = os.path.join(user_dir, "stdout.log")
    stderr_file = os.path.join(user_dir, "stderr.log")
    stdout = open(stdout_file).read() if os.path.exists(stdout_file) else ""
    stderr = open(stderr_file).read() if os.path.exists(stderr_file) else ""
    return f"STDOUT:\n{stdout}\n\nSTDERR:\n{stderr}"

def kill_sandbox(sandbox_id):
    if sandbox_id not in sandboxes:
        return "Sandbox not found."
    pid = sandboxes[sandbox_id]["pid"]
    try:
        os.kill(pid, signal.SIGKILL)
        del sandboxes[sandbox_id]
        return f"Sandbox {sandbox_id} killed."
    except Exception as e:
        return f"Error killing sandbox: {e}"

def status_sandbox(sandbox_id):
    if sandbox_id not in sandboxes:
        return "Sandbox not found."
    elapsed = time.time() - sandboxes[sandbox_id]["start_time"]
    return f"Running for {int(elapsed)} seconds"

# -------------------------
# Gradio UI
# -------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Sandbox Backend")
    
    code_input = gr.Textbox(label="Python Code", lines=15, placeholder="Write your code here...")
    launch_btn = gr.Button("Launch Sandbox")
    sandbox_id_output = gr.Textbox(label="Sandbox ID")
    
    logs_btn = gr.Button("Fetch Logs")
    logs_output = gr.Textbox(label="Logs", lines=10)
    
    kill_btn = gr.Button("Kill Sandbox")
    kill_output = gr.Textbox(label="Kill Status")
    
    status_btn = gr.Button("Check Status")
    status_output = gr.Textbox(label="Sandbox Status")
    
    # Temporary storage for sandbox ID
    sandbox_id_storage = gr.State(value=None)
    
    # -------------------------
    # Button actions
    # -------------------------
    launch_btn.click(
        launch_sandbox,
        inputs=code_input,
        outputs=[sandbox_id_output, sandbox_id_storage]
    )
    logs_btn.click(
        fetch_logs,
        inputs=sandbox_id_storage,
        outputs=logs_output
    )
    kill_btn.click(
        kill_sandbox,
        inputs=sandbox_id_storage,
        outputs=kill_output
    )
    status_btn.click(
        status_sandbox,
        inputs=sandbox_id_storage,
        outputs=status_output
    )

# -------------------------
# Launch Gradio app
# -------------------------
demo.launch()