from flask import Flask, request, jsonify, render_template, stream_with_context, Response
import os
import subprocess
import tempfile
import shutil

app = Flask(__name__)

# Create a temporary directory for operations
temp_dir = tempfile.mkdtemp()

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

def execute_shell_command(command):
    """Executes a shell command and streams output."""
    process = subprocess.Popen(
        command,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        cwd=temp_dir
    )
    for line in iter(process.stdout.readline, ""):
        yield f"{line}<br>"
    for line in iter(process.stderr.readline, ""):
        yield f"Error: {line}<br>"
    process.stdout.close()
    process.stderr.close()
    process.wait()

@app.route("/execute", methods=["POST"])
def execute_code():
    command = request.json.get("code", "").strip()
    if not command:
        return jsonify({"result": "Error: No command provided."})

    try:
        # If the command starts with "!", treat it as a shell command
        if command.startswith("!"):
            shell_command = command[1:]
            return Response(stream_with_context(execute_shell_command(shell_command)))
        else:
            # Treat the command as Python code
            process = subprocess.run(
                ["python3", "-c", command],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                cwd=temp_dir
            )
            return jsonify({"result": process.stdout + process.stderr})
    except Exception as e:
        return jsonify({"result": f"Error: {e}"})

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

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