from fastapi import FastAPI
from env import TrafficEnv

CONFIG = {
    "max_steps": 20,
    "max_queue": 20,
    "arrival_rate": (0, 2),
    "discharge_rate": (3, 5),
    "emergency_prob": 0.02,
    "switch_penalty": 0.2,
    "starvation_threshold": 10,
    "burst_prob": 0.1,
    "burst_multiplier": 1.2,
}

app = FastAPI()
env = TrafficEnv(CONFIG)

@app.get("/")
def root():
    return {"message": "Traffic Env Running"}

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

@app.post("/step")
def step(action: dict):
    action_value = action.get("action", 0)
    next_state, reward, done, info = env.step(action_value)
    return {
        "next_state": next_state,
        "reward": reward,
        "done": done,
        "info": info,
    }