from fastapi import FastAPI
from env import EmergencyEnv

app = FastAPI()

env = EmergencyEnv("medium")


@app.post("/reset")
def reset():
    state = env.reset()
    return state


@app.post("/step")
def step(action: int):
    state, reward, done, _ = env.step(action)
    return {
        "state": state,
        "reward": reward,
        "done": done
    }


@app.get("/state")
def get_state():
    return env.get_state()