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

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):
    # Mapping task_id to levels 1, 2, 3
    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):
    """
    OpenEnv compliant step logic using Pydantic action models.
    """
    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():
    """State API for debugging and grading."""
    return env.state()

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)