import uvicorn
from fastapi import FastAPI, Request
from env import SOCEnvironment
from models import Action

app = FastAPI()
soc_env = SOCEnvironment()

@app.post("/reset")
async def reset(request: Request):
    body = await request.json() if await request.body() else {}
    task_id = body.get("task_id", "task_1_triage")
    obs = soc_env.reset(task_id)
    return {"observation": obs.model_dump()}

@app.post("/step")
async def step(action: Action):
    obs, reward, done, info = soc_env.step(action)
    return {
        "observation": obs.model_dump(), 
        "reward": reward.score_delta, 
        "done": done, 
        "info": info
    }

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

# Added to satisfy the multi-mode deployment validator
def main():
    uvicorn.run("server.app:app", host="0.0.0.0", port=7860)

if __name__ == "__main__":
    main()