#!/usr/bin/env python3
"""
LibreChat Pyodide Code Interpreter - Client-Side Python Execution
"""

import gradio as gr

def create_pyodide_interface():
    """Create a Gradio interface that uses Pyodide for client-side Python execution"""
    
    # Enhanced HTML/JS with better error handling and status updates
    pyodide_html = """
    <div id="pyodide-container" style="border: 1px solid #ddd; padding: 15px; border-radius: 5px; margin: 10px 0;">
        <div id="pyodide-status" style="font-weight: bold; padding: 10px; background: #f0f0f0; border-radius: 3px;">
            🔄 Loading Pyodide... This may take 10-30 seconds on first load.
        </div>
        <div id="pyodide-output" style="display:none; margin-top: 10px;">
            <h4>Execution Results:</h4>
            <pre id="output-text" style="background: #f8f8f8; padding: 10px; border-radius: 3px; max-height: 300px; overflow-y: auto;"></pre>
            <div id="plot-container" style="text-align: center; margin-top: 10px;"></div>
        </div>
    </div>
    
    <script src="https://cdn.jsdelivr.net/pyodide/v0.24.1/full/pyodide.js"></script>
    <script>
        let pyodide = null;
        let pyodideReady = false;
        let initializationStarted = false;
        
        function updateStatus(message, color = 'black') {
            const statusDiv = document.getElementById('pyodide-status');
            statusDiv.innerHTML = message;
            statusDiv.style.color = color;
            console.log('Pyodide Status:', message);
        }
        
        async function initPyodide() {
            if (initializationStarted) return;
            initializationStarted = true;
            
            try {
                updateStatus('🔄 Loading Pyodide core...', 'blue');
                
                pyodide = await loadPyodide({
                    indexURL: "https://cdn.jsdelivr.net/pyodide/v0.24.1/full/"
                });
                
                updateStatus('📦 Loading Python packages (numpy, matplotlib, pandas, plotly)...', 'blue');
                
                await pyodide.loadPackage(["numpy", "matplotlib", "pandas", "plotly"]);
                updateStatus('📦 All packages loaded, setting up backends...', 'blue');
                
                // Set up matplotlib backend for web - IMPROVED VERSION
                pyodide.runPython(`
                    import matplotlib
                    matplotlib.use('AGG')
                    import matplotlib.pyplot as plt
                    import numpy as np
                    import pandas as pd
                    import plotly.graph_objects as go
                    import plotly.express as px
                    import plotly
                    import io
                    import base64
                    import json
                    
                    # Global variables for plot data
                    _current_plot_data = None
                    _current_plotly_data = None
                    
                    def capture_matplotlib():
                        global _current_plot_data
                        try:
                            if len(plt.get_fignums()) > 0:
                                buffer = io.BytesIO()
                                plt.savefig(buffer, format='png', bbox_inches='tight', dpi=100)
                                buffer.seek(0)
                                plot_data = buffer.getvalue()
                                buffer.close()
                                _current_plot_data = base64.b64encode(plot_data).decode()
                                plt.close('all')
                                return _current_plot_data
                            return None
                        except Exception as e:
                            print(f"Matplotlib capture error: {e}")
                            return None
                    
                    def capture_plotly(fig):
                        global _current_plotly_data
                        try:
                            _current_plotly_data = fig.to_html(include_plotlyjs='cdn', div_id='plotly-div')
                            return _current_plotly_data
                        except Exception as e:
                            print(f"Plotly capture error: {e}")
                            return None
                    
                    # Override functions
                    original_show = plt.show
                    def custom_show(*args, **kwargs):
                        return capture_matplotlib()
                    plt.show = custom_show
                    
                    # Override plotly show
                    original_plotly_show = go.Figure.show
                    def custom_plotly_show(self, *args, **kwargs):
                        return capture_plotly(self)
                    go.Figure.show = custom_plotly_show
                    
                    def get_matplotlib_data():
                        return _current_plot_data
                    
                    def get_plotly_data():
                        return _current_plotly_data
                    
                    def clear_plot_data():
                        global _current_plot_data, _current_plotly_data
                        _current_plot_data = None
                        _current_plotly_data = None
                    
                    print("Pyodide ready with matplotlib AND plotly support!")
                `);
                
                pyodideReady = true;
                updateStatus('✅ Pyodide ready! You can now execute Python code.', 'green');
                
                document.getElementById('pyodide-output').style.display = 'block';
                document.getElementById('output-text').textContent = 'Pyodide initialization complete! Ready to execute Python code.';
                
            } catch (error) {
                console.error('Pyodide initialization error:', error);
                updateStatus('❌ Failed to initialize Pyodide: ' + error.toString(), 'red');
                pyodideReady = false;
            }
        }
        
        async function executePyodideCode(code) {
            console.log('Execute called, pyodideReady:', pyodideReady);
            
            if (!pyodideReady) {
                updateStatus('⏳ Pyodide is still initializing. Please wait...', 'orange');
                return "Pyodide is still loading. Please wait for the green 'ready' status above.";
            }
            
            if (!pyodide) {
                return "Error: Pyodide not available.";
            }
            
            try {
                updateStatus('▶️ Executing Python code...', 'blue');
                
                // Clear any previous plot data
                pyodide.runPython('clear_plot_data()');
                
                // Capture stdout
                pyodide.runPython(`
                    import sys
                    from io import StringIO
                    old_stdout = sys.stdout
                    sys.stdout = captured_output = StringIO()
                `);
                
                // Execute user code
                let result = pyodide.runPython(code);
                
                // Get captured output (this should NOT contain base64 data)
                let stdout = pyodide.runPython(`
                    sys.stdout = old_stdout
                    captured_output.getvalue()
                `);
                
                // Get both types of plot data
                let matplotlibData = pyodide.runPython('get_matplotlib_data()');
                let plotlyData = pyodide.runPython('get_plotly_data()');
                
                console.log('Stdout length:', stdout ? stdout.length : 0);
                console.log('Matplotlib data length:', matplotlibData ? matplotlibData.length : 0);
                console.log('Plotly data length:', plotlyData ? plotlyData.length : 0);
                console.log('Matplotlib data preview:', matplotlibData ? matplotlibData.substring(0, 50) + '...' : 'None');
                console.log('Plotly data preview:', plotlyData ? plotlyData.substring(0, 50) + '...' : 'None');
                
                // Display results
                let outputDiv = document.getElementById('pyodide-output');
                let outputText = document.getElementById('output-text');
                let plotContainer = document.getElementById('plot-container');
                
                outputDiv.style.display = 'block';
                
                let textOutput = stdout || (result !== undefined ? String(result) : '');
                outputText.textContent = textOutput || 'Code executed successfully (no output)';
                
                // Display plots
                let plotHtml = '';
                
                if (matplotlibData && matplotlibData.length > 100) {
                    plotHtml += `
                        <div style="margin: 10px 0;">
                            <h5>📊 Matplotlib Plot:</h5>
                            <img src="data:image/png;base64,${matplotlibData}" 
                                 style="max-width: 100%; height: auto; border: 1px solid #ddd;" 
                                 alt="Matplotlib Plot">
                        </div>
                    `;
                }
                
                if (plotlyData) {
                    plotHtml += `
                        <div style="margin: 10px 0;">
                            <h5>📊 Interactive Plotly Plot:</h5>
                            <div style="border: 1px solid #ddd; border-radius: 3px;">
                                ${plotlyData}
                            </div>
                        </div>
                    `;
                }
                
                plotContainer.innerHTML = plotHtml;
                
                if (plotHtml) {
                    if (textOutput.trim()) {
                        outputText.textContent += '\\n\\n📊 Plot(s) generated and displayed below!';
                    } else {
                        outputText.textContent = '📊 Plot(s) generated and displayed below!';
                    }
                    updateStatus('✅ Code executed with plot(s) generated!', 'green');
                } else {
                    updateStatus('✅ Code executed successfully!', 'green');
                }
                
                return textOutput || 'Code executed successfully';
                
            } catch (error) {
                console.error('Execution error:', error);
                document.getElementById('output-text').textContent = 'Error: ' + error.toString();
                updateStatus('❌ Execution error: ' + error.toString(), 'red');
                return 'Error: ' + error.toString();
            }
        }
        
        // Auto-retry initialization if it fails
        async function safeInitPyodide() {
            try {
                await initPyodide();
            } catch (error) {
                console.error('Init failed, retrying in 3 seconds:', error);
                updateStatus('🔄 Initialization failed, retrying in 3 seconds...', 'orange');
                setTimeout(safeInitPyodide, 3000);
            }
        }
        
        // Initialize when page loads
        if (document.readyState === 'loading') {
            document.addEventListener('DOMContentLoaded', safeInitPyodide);
        } else {
            safeInitPyodide();
        }
        
        // Make functions globally available
        window.executePyodideCode = executePyodideCode;
        window.checkPyodideStatus = () => pyodideReady;
    </script>
    """
    
    return pyodide_html

# Create the Gradio interface
with gr.Blocks(title="LibreChat Pyodide Code Interpreter") as demo:
    gr.Markdown("# LibreChat Pyodide Code Interpreter")
    gr.Markdown("**Client-side Python execution** using Pyodide - Python runs directly in your browser!")
    gr.Markdown("⏳ **Please wait for Pyodide to fully initialize before executing code (watch the status above).**")
    
    # Pyodide interface
    pyodide_interface = gr.HTML(create_pyodide_interface())
    
    with gr.Row():
        with gr.Column():
            code_input = gr.Textbox(
                value="""# Simple test - try this first
print("Hello from browser-based Python!")
print("2 + 2 =", 2 + 2)

# Then try this matplotlib example:
# import matplotlib.pyplot as plt
# import numpy as np
# x = np.linspace(0, 10, 50)
# plt.plot(x, np.sin(x))
# plt.title('Sine Wave')
# plt.show()""",
                lines=15,
                label="Python Code (executes in your browser)"
            )
            
            execute_btn = gr.Button("Execute with Pyodide", variant="primary")
            check_status_btn = gr.Button("Check Pyodide Status", variant="secondary")
        
        with gr.Column():
            gr.Markdown("### Execution Status")
            status_display = gr.Textbox(
                label="Last Execution Result",
                interactive=False,
                lines=3
            )
    
    # JavaScript execution
    execute_btn.click(
        fn=None,
        inputs=[code_input],
        outputs=[status_display],
        js="""
        function(code) {
            if (window.executePyodideCode) {
                return window.executePyodideCode(code);
            } else {
                return "Pyodide functions not available. Check browser console for errors.";
            }
        }
        """
    )
    
    # Status check button
    check_status_btn.click(
        fn=None,
        inputs=[],
        outputs=[status_display],
        js="""
        function() {
            if (window.checkPyodideStatus) {
                return window.checkPyodideStatus() ? "✅ Pyodide is ready!" : "⏳ Pyodide is still loading...";
            } else {
                return "❌ Pyodide not available.";
            }
        }
        """
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )