from fastapi import FastAPI
from env import ThreatChainEnv

app = FastAPI()

env = ThreatChainEnv()

@app.get("/")
def root():
    return {"message": "ThreatChainEnv API running"}

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

@app.get("/step")
@app.post("/step")
def step(action: str):
    state, reward, done, info = env.step(action)

    return {
        "state": state,
        "reward": reward,
        "done": done,
        "info": info
    }

@app.get("/state")
@app.post("/state")
def state():

    if env.task is None:
        env.reset()

    return env.state()