"""
MimiReady Python Executor Backend
Exécute du code Python avec TensorFlow/Keras sur Hugging Face Spaces
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional
import sys
import io
import base64
import traceback
from contextlib import redirect_stdout, redirect_stderr

# Configuration matplotlib avant import
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

app = FastAPI(title="MimiReady Python Executor")

# CORS pour permettre les requêtes depuis le frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class CodeRequest(BaseModel):
    code: str
    timeout: int = 30

class CodeResponse(BaseModel):
    stdout: str
    stderr: str
    plots: List[str]
    success: bool
    error: Optional[str] = None

# Variables globales pour capturer les plots
_plot_images: List[str] = []

def capture_plots():
    """Capture tous les plots matplotlib en base64"""
    global _plot_images
    _plot_images = []
    
    original_show = plt.show
    
    def patched_show(*args, **kwargs):
        global _plot_images
        buf = io.BytesIO()
        plt.savefig(buf, format='png', dpi=100, bbox_inches='tight', facecolor='white')
        buf.seek(0)
        img_str = base64.b64encode(buf.read()).decode('utf-8')
        _plot_images.append(img_str)
        plt.clf()
        plt.close('all')
    
    plt.show = patched_show
    return original_show

def restore_show(original_show):
    """Restaure plt.show original"""
    plt.show = original_show

@app.get("/")
def read_root():
    return {
        "service": "MimiReady Python Executor",
        "status": "running",
        "capabilities": ["numpy", "pandas", "matplotlib", "scikit-learn", "tensorflow", "keras"]
    }

@app.get("/health")
def health_check():
    return {"status": "healthy"}

@app.post("/execute", response_model=CodeResponse)
async def execute_code(request: CodeRequest):
    """Exécute du code Python et retourne les résultats"""
    global _plot_images
    
    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()
    
    # Préparer l'environnement d'exécution
    exec_globals = {
        "__builtins__": __builtins__,
        "__name__": "__main__",
    }
    
    # Pre-import des bibliothèques courantes
    try:
        import numpy as np
        exec_globals["np"] = np
        exec_globals["numpy"] = np
    except ImportError:
        pass
    
    try:
        import pandas as pd
        exec_globals["pd"] = pd
        exec_globals["pandas"] = pd
    except ImportError:
        pass
    
    try:
        import sklearn
        exec_globals["sklearn"] = sklearn
    except ImportError:
        pass
    
    # TensorFlow/Keras
    try:
        import tensorflow as tf
        from tensorflow import keras
        exec_globals["tf"] = tf
        exec_globals["tensorflow"] = tf
        exec_globals["keras"] = keras
    except ImportError as e:
        print(f"TensorFlow import warning: {e}")
    
    # Matplotlib
    exec_globals["plt"] = plt
    exec_globals["matplotlib"] = matplotlib
    
    success = True
    error_msg = None
    
    original_show = capture_plots()
    
    try:
        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
            exec(request.code, exec_globals)
    except Exception as e:
        success = False
        error_msg = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
    finally:
        restore_show(original_show)
    
    return CodeResponse(
        stdout=stdout_capture.getvalue(),
        stderr=stderr_capture.getvalue(),
        plots=_plot_images.copy(),
        success=success,
        error=error_msg
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)