import gradio as gr
from fastapi import FastAPI

app = FastAPI()

class EmailEnv:
    def __init__(self):
        self.reset()

    def reset(self):
        self.emails = [
            {"id": 1, "from": "alice@example.com", "subject": "Meeting Tomorrow", "status": "Unread"},
            {"id": 2, "from": "bob@example.com", "subject": "Lunch Plans", "status": "Archived"}
        ]
        return {"emails": self.emails}

    def step(self, action):
        return {
            "emails": self.emails,
            "reward": 1,
            "done": False
        }

env = EmailEnv()

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

@app.post("/step")
def step(action: dict):
    return env.step(action)

def ui(action):
    return f"Selected: {action}"

iface = gr.Interface(
    fn=ui,
    inputs=gr.Radio(["important", "spam", "normal"]),
    outputs="text",
    title="Smart Email AI Trainer"
)

app = gr.mount_gradio_app(app, iface, path="/ui")