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

app = FastAPI(
    title="Email Sorting OpenEnv",
    version="1.0.0"
)

env = EmailSortingEnv()

class StepRequest(BaseModel):
    action: str

@app.get("/")
def root():
    return {
        "name": "Email Sorting OpenEnv",
        "version": "1.0.0",
        "description": "Sort emails as spam, important, or promotion",
        "endpoints": ["/reset", "/step", "/state", "/health"]
    }

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/reset")
def reset():
    state = env.reset()
    return {"status": "success", "state": state}

@app.post("/step")
def step(request: StepRequest):
    next_state, reward, done, info = env.step(request.action)
    return {"status": "success", "state": next_state, "reward": reward, "done": done, "info": info}

@app.get("/state")
def get_state():
    return {"status": "success", "state": env.state()}

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

if __name__ == "__main__":
    main()