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

from IPython.display import display, HTML
from contextlib import redirect_stdout
from IPython.core.interactiveshell import InteractiveShell

# Create a function to capture all kinds of Jupyter-like outputs
def run_jupyter_style_code(code):
    # Reset any old figures
    plt.close('all')

    # Setup for capturing output
    shell = InteractiveShell.instance()
    stdout = io.StringIO()

    try:
        with redirect_stdout(stdout):
            result = shell.run_cell(code)

        # Collect rich display outputs
        rich_output = ""
        if result.result is not None:
            try:
                display_out = display(result.result, display_id=False)
                rich_output = str(result.result)
            except Exception:
                rich_output = str(result.result)

        # Show latest matplotlib figure if available
        figures = [plt.figure(i) for i in plt.get_fignums()]
        img_path = None
        if figures:
            img_path = "plot.png"
            figures[-1].savefig(img_path)

        return stdout.getvalue() + "\n" + rich_output, img_path

    except Exception as e:
        return traceback.format_exc(), None

with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Python Code Runner with Jupyter-style Output")

    code_input = gr.Code(language="python", label="Enter Python Code")
    output_text = gr.Textbox(label="Console Output + Rich Result", lines=12)
    plot_output = gr.Image(label="Plot (if any)")

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

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