# app.py
import gradio as gr
import sys, io, traceback, re, subprocess, importlib, multiprocessing, resource, signal

# ----------------------
# Helper: Install missing packages safely
# ----------------------
def _install_missing(packages):
    """Install packages if not already installed. Fail gracefully if not found."""
    for pkg in packages:
        try:
            importlib.import_module(pkg)
        except ImportError:
            try:
                subprocess.check_call(
                    [sys.executable, "-m", "pip", "install", "--disable-pip-version-check", "-q", pkg]
                )
            except subprocess.CalledProcessError:
                # فقط هشدار می‌ده، برنامه نمی‌ترکه
                print(f"[WARN] Package '{pkg}' not found on PyPI or failed to install.", file=sys.stderr)

# ----------------------
# Worker process
# ----------------------
def _worker(code, stdin_input, conn):
    sys.stdin = io.StringIO(stdin_input)
    sys.stdout = io.StringIO()
    sys.stderr = io.StringIO()

    # محدودیت حافظه و زمان
    resource.setrlimit(resource.RLIMIT_AS, (512 * 1024 * 1024, 512 * 1024 * 1024))  # 512MB
    signal.alarm(10)  # 10 ثانیه تایم‌اوت

    try:
        # پیدا کردن importها
        imports = re.findall(r'^\s*import (\w+)|^\s*from (\w+)', code, flags=re.MULTILINE)
        imports = {i[0] or i[1] for i in imports}
        if imports:
            _install_missing(imports)

        safe_globals = {"__builtins__": __builtins__}
        exec(compile(code, "<user_code>", "exec"), safe_globals, None)
        out = sys.stdout.getvalue()
        err = sys.stderr.getvalue()
        conn.send((out, err))
    except Exception:
        conn.send(("", traceback.format_exc()))
    finally:
        conn.close()

# ----------------------
# Runner
# ----------------------
def run_code(code, stdin_input=""):
    parent_conn, child_conn = multiprocessing.Pipe()
    p = multiprocessing.Process(target=_worker, args=(code, stdin_input, child_conn))
    p.start()
    p.join(12)  # یک مقدار بیشتر از تایم‌اوت داخلی
    if p.is_alive():
        p.terminate()
        return "", "Error: Timeout (execution took too long)"
    return parent_conn.recv()

# ----------------------
# Gradio UI
# ----------------------
with gr.Blocks(title="Python Runner (Auto Install)") as demo:
    gr.Markdown("### 🐍 Python Sandbox with Auto-Package Install (USE WITH CAUTION)")

    code = gr.Code(label="Python code", language="python")
    stdin = gr.Textbox(label="stdin", placeholder="Optional input for `input()` calls")
    run_btn = gr.Button("▶ Run")
    stdout = gr.Textbox(label="stdout")
    stderr = gr.Textbox(label="stderr", elem_classes="stderr")

    def execute(c, s):
        return run_code(c, s)

    run_btn.click(execute, [code, stdin], [stdout, stderr])

if __name__ == "__main__":
    demo.launch()