from flask import Flask, request, jsonify, render_template
import io
import sys
import subprocess
import tempfile
import shutil
import os

app = Flask(__name__)

# Create a temporary directory for user-installed packages
temp_dir = tempfile.mkdtemp()

@app.route("/")
def index():
    return render_template("index.html")

@app.route("/execute", methods=["POST"])
def execute_code():
    code = request.json.get("code", "").strip()
    output = io.StringIO()
    sys.stdout = output
    sys.stderr = output

    try:
        if code.startswith("!pip install"):
            # Handle pip install commands with temporary directory
            package = code.split("!pip install", 1)[1].strip()
            if package:
                result = subprocess.run(
                    [sys.executable, "-m", "pip", "install", package, "--target", temp_dir],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                )
                return jsonify({"result": result.stdout + result.stderr})
            else:
                return jsonify({"result": "Error: No package specified for installation."})
        else:
            # Adjust sys.path to include temporary package directory
            sys.path.insert(0, temp_dir)
            exec(code)
            result = output.getvalue()
    except Exception as e:
        result = f"Error: {e}"
    finally:
        # Reset stdout and stderr
        sys.stdout = sys.__stdout__
        sys.stderr = sys.__stderr__

    return jsonify({"result": result})

@app.route("/cleanup", methods=["POST"])
def cleanup():
    # Delete the temporary directory
    global temp_dir
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
        temp_dir = tempfile.mkdtemp()  # Recreate for the next session
    return jsonify({"result": "Temporary files cleaned up."})

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