from fastapi import FastAPI
from pydantic import BaseModel
from env import EmailEnv
import uvicorn

app = FastAPI()
env = EmailEnv()

class Action(BaseModel):
    label: str
    reply: str

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

@app.post("/step")
def step(action: Action):
    reward = env.step(action.dict())
    return {
        "reward": reward,
        "done": True
    }

@app.get("/")
def home():
    return {"message": "Email Triage AI Environment Running"}


def main():
    uvicorn.run(app, host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()