import gradio as gr
import sys
import subprocess
import importlib
import io
import contextlib
from pygments import highlight
from pygments.lexers import PythonLexer
from pygments.formatters import HtmlFormatter
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Function to install packages dynamically
def install_package(package):
    subprocess.check_call([sys.executable, "-m", "pip", "install", package])

# Function to execute user code with output capturing and visualization support
def execute_code(code, libraries):
    # Install libraries
    for library in libraries.splitlines():
        if library.strip():
            install_package(library.strip())
    
    # Prepare output capturing
    output = io.StringIO()
    error_output = io.StringIO()
    
    # Execute the code
    try:
        with contextlib.redirect_stdout(output), contextlib.redirect_stderr(error_output):
            exec_globals = {
                'plt': plt,
                'np': np,
                'pd': pd,
            }
            exec(code, exec_globals)
        
        # Check if a plot was created
        if plt.get_fignums():
            plt.savefig('plot.png')
            plt.close()
            return output.getvalue(), error_output.getvalue(), 'plot.png'
        else:
            return output.getvalue(), error_output.getvalue(), None
    except Exception as e:
        return "", str(e), None

# Function to highlight Python code
def highlight_code(code):
    return highlight(code, PythonLexer(), HtmlFormatter(style="monokai"))

# Define the Gradio interface
with gr.Blocks(css=".gradio-container {background-color: #f0f0f0;}") as demo:
    gr.Markdown("# Advanced Python Code Executor")
    
    with gr.Tabs():
        with gr.TabItem("Code Execution"):
            with gr.Row():
                with gr.Column(scale=2):
                    code_input = gr.Code(label="Python Code", language="python", lines=10)
                    libraries_input = gr.TextArea(label="Libraries (one per line)", placeholder="e.g.\nrequests\nscipy", lines=3)
                    execute_button = gr.Button("Run Code", variant="primary")
                with gr.Column(scale=1):
                    output = gr.Textbox(label="Standard Output", lines=10)
                    error_output = gr.Textbox(label="Error Output", lines=5)
                    plot_output = gr.Image(label="Plot Output")
        
        with gr.TabItem("Code Highlighting"):
            code_to_highlight = gr.Code(label="Enter Python code to highlight", language="python", lines=10)
            highlight_button = gr.Button("Highlight Code")
            highlighted_code = gr.HTML(label="Highlighted Code")
    
    execute_button.click(
        fn=execute_code,
        inputs=[code_input, libraries_input],
        outputs=[output, error_output, plot_output]
    )
    
    highlight_button.click(
        fn=highlight_code,
        inputs=[code_to_highlight],
        outputs=[highlighted_code]
    )

    gr.Markdown("""
    ## Features:
    - Execute Python code with custom library imports
    - Capture and display standard output and error messages separately
    - Support for data visualization with matplotlib
    - Code highlighting for better readability
    - Tab-based interface for different functionalities
    
    ## Tips:
    - Use `plt` for matplotlib, `np` for numpy, and `pd` for pandas (pre-imported)
    - Your plot will automatically be displayed if you use `plt.plot()` or similar functions
    """)

# Launch the interface
demo.launch()