from fastapi import FastAPI, Body
import subprocess

app = FastAPI(title="Agent Sandbox API")

@app.post("/run")
def run_code(payload: dict = Body(...)):
    """
    Execute Python code in a sandboxed subprocess.
    Input JSON: {"code": "print(2+2)"} or {"code": "2+2"}
    Output JSON: {"output": "4"} or {"error": "..."}
    """
    code = payload.get("code", "")
    if not code:
        return {"error": "No code provided"}

    try:
        # Run the code in a subprocess with a 5s timeout
        proc = subprocess.run(
            ["python3", "-c", code],
            capture_output=True,
            text=True,
            timeout=5
        )

        if proc.returncode != 0:
            return {"error": proc.stderr.strip() or "Unknown error"}

        output = proc.stdout.strip()

        # If nothing printed, try evaluating expression directly
        if not output:
            try:
                result = eval(code, {})
                output = str(result)
            except Exception:
                output = "<no output>"

        return {"output": output}

    except subprocess.TimeoutExpired:
        return {"error": "Execution timed out"}
    except Exception as e:
        return {"error": str(e)}