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):
    plt.close('all')  # Close any old plots

    shell = InteractiveShell.instance()
    img_path = None
    output = ""

    try:
        # Capture output (stdout, stderr) and display anything that gets returned
        with capture_output(display=True) as captured:
            exec(code, {"plt": plt, "pd": __import__('pandas'), "display": display})

        # Get any figures generated by Matplotlib
        figures = [plt.figure(i) for i in plt.get_fignums()]
        if figures:
            img_path = "plot.png"
            figures[-1].savefig(img_path)

        # Join captured output (stdout + rich outputs like HTML or plots)
        output = captured.stdout + ''.join(str(out) for out in captured.outputs)
        return output, img_path

    except Exception as e:
        # If an error occurs, return the traceback as output
        output = traceback.format_exc()
        return output, None

# Create the Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🧪 Python Code Runner with Jupyter-style Terminal")

    # Code input area (multi-line text box simulating terminal)
    code_input = gr.Textbox(label="Python Terminal", placeholder="Enter Python code...", lines=10)
    
    # Output for the console/logs (including rich content)
    output_text = gr.Textbox(label="Console Output", lines=12, interactive=False)
    
    # Display for any generated plots
    plot_output = gr.Image(label="Generated Plot", type="filepath")
    
    # Button to execute the code
    run_button = gr.Button("Run Code")
    
    # Trigger execution on button click
    run_button.click(fn=run_jupyter_style_code, inputs=code_input, outputs=[output_text, plot_output])

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