from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import pandas as pd
import traceback
import sys
import io

app = FastAPI()

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

class CodeExecutionRequest(BaseModel):
    code: str

@app.post("/execute")
async def execute_code(request: CodeExecutionRequest):
    # Redirect stdout to capture print statements
    old_stdout = sys.stdout
    new_stdout = io.StringIO()
    sys.stdout = new_stdout
        
    code = request.code
    try:
        # Define a safe local environment for code execution
        local_env = {}

        # Execute the provided code
        exec(code)
        output = new_stdout.getvalue()
    
    except Exception as e:
        return {"error": str(e), "traceback": traceback.format_exc()}
    finally:
        # Reset stdout to the original value
        sys.stdout = old_stdout

    return {"output": output}