from __future__ import annotations

from fastapi import FastAPI

from support_triage_openenv.env import SupportTriageEnv
from support_triage_openenv.models import Action

app = FastAPI(title="Support Triage OpenEnv", version="0.1.0")
env = SupportTriageEnv()


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/reset")
def reset(payload: dict | None = None) -> dict:
    task_id = (payload or {}).get("task_id")
    obs = env.reset(task_id=task_id)
    return obs.model_dump()


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


@app.get("/state")
def state() -> dict:
    return env.state()