import gradio as gr
import io
import contextlib
import traceback
import matplotlib.pyplot as plt

def run_python_code(code):
    stdout = io.StringIO()
    img_path = "plot.png"

    # Clean up any old plot
    plt.close()

    try:
        with contextlib.redirect_stdout(stdout):
            exec(code, {"plt": plt})  # Provide plt to global scope

        # Check if a figure was created
        figures = [plt.figure(i) for i in plt.get_fignums()]
        if figures:
            # Save the latest figure
            figures[-1].savefig(img_path)
            return stdout.getvalue(), img_path
        else:
            return stdout.getvalue(), None

    except Exception:
        return traceback.format_exc(), None

with gr.Blocks() as demo:
    gr.Markdown("## 🐍 Python Code Runner with Plot Support (`matplotlib`, `pandas`)")
    
    code_input = gr.Code(language="python", label="Enter Python Code")
    output_text = gr.Textbox(label="Text Output", lines=10)
    plot_output = gr.Image(label="Generated Plot")

    run_button = gr.Button("Run Code")
    run_button.click(fn=run_python_code, inputs=code_input, outputs=[output_text, plot_output])

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