import gradio as gr
import subprocess
import matplotlib.pyplot as plt
import traceback
import sys
import io
import numpy as np
from PIL import Image

def run_code(code):
    try:
        if any(bad_cmd in code for bad_cmd in ['rm', 'os', 'shutil']):
            return create_image("")

        old_stdout = sys.stdout
        sys.stdout = io.StringIO()

        local_env = {'plt': plt, 'subprocess': subprocess}
        exec(code, {}, local_env)

        output = sys.stdout.getvalue()
        sys.stdout = old_stdout
        
        if any('plt' in line for line in code.splitlines()):
            buf = io.BytesIO()
            plt.savefig(buf, format='png')
            buf.seek(0)
            img = Image.open(buf)
            img_array = np.array(img)
            plt.close()
            if np.array(img_array).sum() > 0:
                return img_array

        return create_image("")

    except Exception as e:
        return create_image("")

def create_image(message):
    fig, ax = plt.subplots()
    ax.text(0.5, 0.5, message, fontsize=20, ha='center', va='center')
    ax.axis('off')
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = Image.open(buf)
    plt.close()
    return np.array(img)

interface = gr.Interface(
    fn=run_code, 
    inputs=gr.Code(language="python", label="Enter Python Code"), 
    outputs=gr.Image(type="numpy", label="Output"), 
    title="Python Interpreter",
    description="Run Python code. Requires matplotlib for plots."
)

interface.launch()