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

# -------------------------
# Sandbox storage
# -------------------------
sandboxes = {}
SANDBOX_TTL = 2*3600  # 2 hours

# -------------------------
# Worker to run bare-metal process
# -------------------------
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, (SANDBOX_TTL, SANDBOX_TTL))  # 2h CPU

    # Save code as a script
    code_file = os.path.join(user_dir, "user_code.sh")
    with open(code_file, "w") as f:
        f.write(code)
    os.chmod(code_file, 0o755)  # make executable

    try:
        result = subprocess.run(
            ["/bin/bash", code_file],
            capture_output=True,
            text=True,
            timeout=SANDBOX_TTL
        )
        # 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).")

# -------------------------
# Auto-cleanup thread
# -------------------------
def cleanup_sandboxes():
    while True:
        now = time.time()
        for sid, info in list(sandboxes.items()):
            if now - info["start_time"] > SANDBOX_TTL:
                try:
                    os.kill(info["pid"], signal.SIGKILL)
                except:
                    pass
                folder = f"/tmp/sandbox_{sid}"
                if os.path.exists(folder):
                    shutil.rmtree(folder)
                del sandboxes[sid]
        time.sleep(60)

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

# -------------------------
# Gradio API 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 = open(os.path.join(user_dir, "stdout.log")).read() if os.path.exists(os.path.join(user_dir, "stdout.log")) else ""
    stderr = open(os.path.join(user_dir, "stderr.log")).read() if os.path.exists(os.path.join(user_dir, "stderr.log")) 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."
    try:
        os.kill(sandboxes[sandbox_id]["pid"], signal.SIGKILL)
        folder = f"/tmp/sandbox_{sandbox_id}"
        if os.path.exists(folder):
            shutil.rmtree(folder)
        del sandboxes[sandbox_id]
        return f"Sandbox {sandbox_id} killed."
    except Exception as e:
        return f"Error: {e}"

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

# -------------------------
# Gradio UI
# -------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🛡 Bare-Metal Sandbox")

    code_input = gr.Textbox(label="Shell Script / Commands", lines=15, placeholder="#!/bin/bash\nls -la")
    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")

    sandbox_state = gr.State(value=None)

    launch_btn.click(
        launch_sandbox,
        inputs=code_input,
        outputs=[sandbox_id_output, sandbox_state]
    )
    logs_btn.click(fetch_logs, inputs=sandbox_state, outputs=logs_output)
    kill_btn.click(kill_sandbox, inputs=sandbox_state, outputs=kill_output)
    status_btn.click(status_sandbox, inputs=sandbox_state, outputs=status_output)

demo.launch()