main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, Optional, List
from env import EmailEnv
from models import Action, Observation
import uvicorn
import os

app = FastAPI(title="EmailTriage OpenEnv", description="A real-world email management and scheduling environment.")

# Global instance
env = EmailEnv()

class ResetRequest(BaseModel):
    task_id: int = 1

@app.post("/reset")
async def reset_env(req: ResetRequest):
    obs, info = env.reset(options={"level": req.task_id})
    return {
        "observation": obs,
        "info": info,
        "status": "Ready"
    }

@app.post("/step")
async def step_env(action: Dict):
    obs, reward, terminated, truncated, info = env.step(action)
    return {
        "observation": obs,
        "reward": float(reward),
        "terminated": bool(terminated),
        "truncated": bool(truncated),
        "info": info
    }

@app.get("/state")
async def get_state():
    return env.state()

# Health check for HF Spaces
@app.get("/")
async def root():
    return {"status": "running", "environment": "OpenEnv"}

if __name__ == "__main__":
    # HF Spaces default port is 7860
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)