import gradio as gr
import io
import traceback
import matplotlib.pyplot as plt
from IPython.display import display
from IPython.core.interactiveshell import InteractiveShell
from IPython.utils.capture import capture_output

def run_jupyter_style_code(code):
    # Close any old plots
    plt.close('all')

    # Set up an interactive shell instance
    shell = InteractiveShell.instance()
    img_path = None

    try:
        # Capture output (both text and rich outputs like HTML, plots)
        with capture_output(display=True) as captured:
            exec(code, {"plt": plt, "pd": __import__('pandas'), "display": display})

        # Check for any figures and save the last one as an image
        figures = [plt.figure(i) for i in plt.get_fignums()]
        if figures:
            img_path = "plot.png"
            figures[-1].savefig(img_path)

        # Return the captured HTML/text outputs
        html_output = captured.stdout + ''.join(str(out) for out in captured.outputs)
        return html_output, img_path

    except Exception as e:
        # Return any errors as string
        return traceback.format_exc(), None


# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Python Code Runner with Jupyter-style Output")
    
    # Code input area
    code_input = gr.Code(language="python", label="Enter Python Code")
    
    # Outputs (HTML + Plot)
    output_text = gr.HTML(label="Console Output + Rich Result")
    plot_output = gr.Image(label="Generated Plot", type="filepath")

    # Run button
    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()