import gradio as gr
import sys
import io
import traceback
import subprocess
import os
import json
import base64
from pathlib import Path
import tempfile
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# Create workspace directory
WORKSPACE_DIR = Path("/tmp/code_workspace")
WORKSPACE_DIR.mkdir(exist_ok=True)

class CodeInterpreter:
    def __init__(self):
        self.session_vars = {}
        self.installed_packages = set()
        
    def install_package(self, package_name):
        """Install a Python package dynamically"""
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", package_name, "-q"])
            self.installed_packages.add(package_name)
            return f"✅ Successfully installed: {package_name}"
        except Exception as e:
            return f"❌ Failed to install {package_name}: {str(e)}"
    
    def execute_code(self, code, maintain_session=True):
        """Execute Python code with proper output capture"""
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()
        
        # Save current stdout/stderr
        old_stdout = sys.stdout
        old_stderr = sys.stderr
        
        try:
            # Redirect output
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture
            
            # Prepare execution environment
            exec_globals = {
                '__builtins__': __builtins__,
                'WORKSPACE': str(WORKSPACE_DIR),
                'workspace_path': str(WORKSPACE_DIR),
            }
            
            if maintain_session:
                exec_globals.update(self.session_vars)
            
            # Execute code
            exec(code, exec_globals)
            
            # Update session variables
            if maintain_session:
                self.session_vars.update({
                    k: v for k, v in exec_globals.items() 
                    if not k.startswith('__') and k not in ['WORKSPACE', 'workspace_path']
                })
            
            # Get output
            stdout_output = stdout_capture.getvalue()
            stderr_output = stderr_capture.getvalue()
            
            result = {
                'success': True,
                'stdout': stdout_output,
                'stderr': stderr_output,
                'error': None
            }
            
        except Exception as e:
            result = {
                'success': False,
                'stdout': stdout_capture.getvalue(),
                'stderr': stderr_capture.getvalue(),
                'error': f"{type(e).__name__}: {str(e)}\n\n{traceback.format_exc()}"
            }
        finally:
            # Restore stdout/stderr
            sys.stdout = old_stdout
            sys.stderr = old_stderr
            
        return result
    
    def reset_session(self):
        """Clear session variables"""
        self.session_vars.clear()
        plt.close('all')
        return "Session reset successfully"
    
    def list_workspace_files(self):
        """List files in workspace"""
        files = list(WORKSPACE_DIR.glob("*"))
        if not files:
            return "Workspace is empty"
        return "\n".join([f"📄 {f.name}" for f in files if f.is_file()])

# Initialize interpreter
interpreter = CodeInterpreter()

def run_code(code, maintain_session):
    """Main function to run code and return formatted output"""
    if not code.strip():
        return "⚠️ Please enter some code to execute"
    
    result = interpreter.execute_code(code, maintain_session)
    
    output_parts = []
    
    if result['stdout']:
        output_parts.append("📤 OUTPUT:\n" + "=" * 50 + "\n" + result['stdout'])
    
    if result['stderr']:
        output_parts.append("⚠️ WARNINGS:\n" + "=" * 50 + "\n" + result['stderr'])
    
    if result['error']:
        output_parts.append("❌ ERROR:\n" + "=" * 50 + "\n" + result['error'])
    
    if result['success'] and not result['stdout'] and not result['stderr']:
        output_parts.append("✅ Code executed successfully (no output)")
    
    # Check for generated plots
    plot_files = list(WORKSPACE_DIR.glob("*.png")) + list(WORKSPACE_DIR.glob("*.jpg"))
    if plot_files:
        output_parts.append(f"\n📊 Generated {len(plot_files)} plot(s) - check Files tab")
    
    return "\n\n".join(output_parts)

def install_package_handler(package_name):
    """Handle package installation"""
    if not package_name.strip():
        return "⚠️ Please enter a package name"
    return interpreter.install_package(package_name.strip())

def reset_session_handler():
    """Handle session reset"""
    return interpreter.reset_session()

def list_files_handler():
    """Handle file listing"""
    return interpreter.list_workspace_files()

def upload_file_handler(files):
    """Handle file uploads"""
    if not files:
        return "No files uploaded"
    
    uploaded = []
    for file in files:
        try:
            file_path = Path(file.name)
            dest_path = WORKSPACE_DIR / file_path.name
            
            # Copy file to workspace
            with open(file.name, 'rb') as src:
                with open(dest_path, 'wb') as dst:
                    dst.write(src.read())
            
            uploaded.append(f"✅ {file_path.name}")
        except Exception as e:
            uploaded.append(f"❌ {file_path.name}: {str(e)}")
    
    return "Uploaded files:\n" + "\n".join(uploaded)

def get_workspace_files():
    """Get list of files in workspace for download"""
    files = list(WORKSPACE_DIR.glob("*"))
    return [str(f) for f in files if f.is_file()]

# Example code snippets
EXAMPLES = {
    "Hello World": """print("Hello from Code Interpreter!")
print("Python version:", sys.version)""",
    
    "Data Analysis": """import pandas as pd
import numpy as np

# Create sample data
data = {
    'Name': ['Alice', 'Bob', 'Charlie', 'David'],
    'Age': [25, 30, 35, 28],
    'Score': [85, 92, 78, 95]
}

df = pd.DataFrame(data)
print("DataFrame:")
print(df)
print(f"\\nAverage Score: {df['Score'].mean():.2f}")""",
    
    "Visualization": """import matplotlib.pyplot as plt
import numpy as np

# Create data
x = np.linspace(0, 10, 100)
y1 = np.sin(x)
y2 = np.cos(x)

# Create plot
plt.figure(figsize=(10, 6))
plt.plot(x, y1, label='sin(x)', linewidth=2)
plt.plot(x, y2, label='cos(x)', linewidth=2)
plt.xlabel('x')
plt.ylabel('y')
plt.title('Sine and Cosine Functions')
plt.legend()
plt.grid(True, alpha=0.3)

# Save plot
plt.savefig(f'{WORKSPACE}/trigonometric.png', dpi=150, bbox_inches='tight')
print("✅ Plot saved to workspace as 'trigonometric.png'")
plt.close()""",
    
    "File Operations": """import os

# Create a text file
with open(f'{WORKSPACE}/sample.txt', 'w') as f:
    f.write("Hello from Code Interpreter!\\n")
    f.write("This is a sample file.\\n")

# Read it back
with open(f'{WORKSPACE}/sample.txt', 'r') as f:
    content = f.read()

print("File content:")
print(content)

# List workspace files
files = os.listdir(WORKSPACE)
print(f"\\nWorkspace files: {files}")""",

    "Machine Learning": """from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

# Load dataset
iris = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)

# Train model
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Evaluate
predictions = model.predict(X_test)
accuracy = accuracy_score(y_test, predictions)

print(f"Model Accuracy: {accuracy:.2%}")
print(f"\\nFeature Importances:")
for name, importance in zip(iris.feature_names, model.feature_importances_):
    print(f"  {name}: {importance:.3f}")"""
}

def load_example(example_name):
    """Load example code"""
    return EXAMPLES.get(example_name, "")

# Custom CSS for better UI
custom_css = """
#code-input textarea {
    font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace !important;
    font-size: 14px !important;
}
#output-display textarea {
    font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', monospace !important;
    font-size: 13px !important;
}
.gradio-container {
    max-width: 1400px !important;
}
"""

# Build Gradio Interface
with gr.Blocks(css=custom_css, title="🚀 Code Interpreter Pro", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 Code Interpreter Pro - HuggingFace Edition
    
    A powerful Python code execution sandbox with:
    - ✨ **Real-time code execution** with session management
    - 📦 **Dynamic package installation** 
    - 📊 **Data visualization** (matplotlib, plotly, seaborn)
    - 📁 **File upload/download** capabilities
    - 🔄 **Persistent session** (variables maintained across runs)
    - 🎯 **Better than basic sandboxes** - Full-featured development environment
    """)
    
    with gr.Tabs():
        # Main Code Editor Tab
        with gr.Tab("💻 Code Editor"):
            with gr.Row():
                with gr.Column(scale=2):
                    code_input = gr.Code(
                        label="Python Code Editor",
                        language="python",
                        lines=20,
                        elem_id="code-input"
                    )
                    
                    with gr.Row():
                        run_btn = gr.Button("▶️ Run Code", variant="primary", size="lg")
                        maintain_session_check = gr.Checkbox(
                            label="Maintain Session (keep variables)", 
                            value=True
                        )
                    
                    with gr.Row():
                        example_dropdown = gr.Dropdown(
                            choices=list(EXAMPLES.keys()),
                            label="📚 Load Example",
                            value=None
                        )
                        load_example_btn = gr.Button("Load", size="sm")
                
                with gr.Column(scale=2):
                    output_display = gr.Textbox(
                        label="Output",
                        lines=20,
                        max_lines=30,
                        elem_id="output-display"
                    )
        
        # Package Manager Tab
        with gr.Tab("📦 Package Manager"):
            gr.Markdown("### Install Python Packages")
            gr.Markdown("Install any package from PyPI (e.g., pandas, numpy, requests)")
            
            with gr.Row():
                package_input = gr.Textbox(
                    label="Package Name",
                    placeholder="e.g., pandas, requests, beautifulsoup4"
                )
                install_btn = gr.Button("Install", variant="primary")
            
            install_output = gr.Textbox(label="Installation Output", lines=5)
            
            gr.Markdown("""
            ### Pre-installed Packages
            - numpy, pandas, matplotlib
            - scikit-learn, scipy
            - requests, beautifulsoup4
            - And more...
            """)
        
        # File Manager Tab
        with gr.Tab("📁 File Manager"):
            gr.Markdown("### Workspace File Manager")
            gr.Markdown(f"Files are stored in: `{WORKSPACE_DIR}`")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### Upload Files")
                    file_upload = gr.File(
                        label="Select files to upload",
                        file_count="multiple"
                    )
                    upload_btn = gr.Button("Upload to Workspace", variant="primary")
                    upload_output = gr.Textbox(label="Upload Status", lines=3)
                
                with gr.Column():
                    gr.Markdown("#### Workspace Files")
                    list_files_btn = gr.Button("🔄 Refresh File List", variant="secondary")
                    files_list = gr.Textbox(label="Files in Workspace", lines=10)
                    
                    gr.Markdown("#### Download Files")
                    workspace_files = gr.File(
                        label="Download files from workspace",
                        file_count="multiple",
                        type="filepath",
                        interactive=False
                    )
        
        # Session Manager Tab
        with gr.Tab("🔄 Session Manager"):
            gr.Markdown("### Session Management")
            gr.Markdown("Control your execution session and view session state")
            
            reset_btn = gr.Button("🗑️ Reset Session", variant="stop")
            reset_output = gr.Textbox(label="Status", lines=2)
            
            gr.Markdown("""
            ### Session Info
            - **Maintain Session ON**: Variables persist between code executions
            - **Maintain Session OFF**: Fresh environment for each execution
            - **Reset Session**: Clear all variables and start fresh
            """)
        
        # Help Tab
        with gr.Tab("❓ Help & Tips"):
            gr.Markdown("""
            ## 🎯 Quick Start Guide
            
            ### Basic Usage
            1. Write or paste Python code in the editor
            2. Click "▶️ Run Code" to execute
            3. View output in the right panel
            
            ### Advanced Features
            
            #### 📊 Creating Visualizations
            ```python
            import matplotlib.pyplot as plt
            
            plt.figure(figsize=(10, 6))
            plt.plot([1, 2, 3, 4], [1, 4, 9, 16])
            plt.savefig(f'{WORKSPACE}/my_plot.png')
            print("Plot saved!")
            ```
            
            #### 📁 Working with Files
            ```python
            # Access workspace directory
            print(f"Workspace: {WORKSPACE}")
            
            # Save data
            with open(f'{WORKSPACE}/data.txt', 'w') as f:
                f.write("Hello World!")
            
            # Read data
            with open(f'{WORKSPACE}/data.txt', 'r') as f:
                print(f.read())
            ```
            
            #### 📦 Installing Packages
            Go to "Package Manager" tab and install any PyPI package!
            
            ### Tips
            - Use `maintain_session=True` to keep variables between runs
            - Save plots to workspace to download them
            - Upload CSV/JSON files to analyze data
            - Install packages dynamically as needed
            
            ### Advantages over e2b.dev
            - ✅ Free and open-source
            - ✅ Integrated file management
            - ✅ Visual package manager
            - ✅ Session persistence
            - ✅ Better UI/UX
            - ✅ Hosted on HuggingFace (reliable infrastructure)
            """)
    
    # Event handlers
    run_btn.click(
        fn=run_code,
        inputs=[code_input, maintain_session_check],
        outputs=output_display
    )
    
    load_example_btn.click(
        fn=load_example,
        inputs=example_dropdown,
        outputs=code_input
    )
    
    install_btn.click(
        fn=install_package_handler,
        inputs=package_input,
        outputs=install_output
    )
    
    upload_btn.click(
        fn=upload_file_handler,
        inputs=file_upload,
        outputs=upload_output
    )
    
    list_files_btn.click(
        fn=list_files_handler,
        outputs=files_list
    )
    
    reset_btn.click(
        fn=reset_session_handler,
        outputs=reset_output
    )

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