from fastapi import FastAPI
from pydantic import BaseModel

from server.disruption_environment import DisruptionEnv
from models import DisruptionAction

app = FastAPI()
env = DisruptionEnv(seed=0)

class StepRequest(BaseModel):
    action: DisruptionAction

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

@app.post("/step")
def step(request: StepRequest):
    obs, reward, done, info = env.step(request.action)
    return {
        "observation": obs.model_dump(),
        "reward": reward,
        "done": done,
        "info": info
    }