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 "Error: This command is not allowed."
        
        # Capture the standard output
        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
        
        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 noplot()  

    except Exception as e:
        return noplot()  

def noplot():
    fig, ax = plt.subplots()
    ax.text(0.5, 0.5, "Required: matplotlib", 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="code"), 
    outputs=gr.Image(type="numpy", label="Output"), 
    title="matplotlib-graph-creation",
)

interface.launch()