from pydantic import BaseModel
from typing import List, Literal
import gradio as gr

class EmailAction(BaseModel):
    action_type: Literal['send', 'reply', 'archive', 'delete']
    message: str = ""
    email_id: int = None

class EmailObservation(BaseModel):
    emails: List[dict]
    ai_feedback: str

class EmailEnv:
    def __init__(self):
        self.emails = [
            {"id": 1, "from": "alice@example.com", "subject": "Meeting Tomorrow", "label": "Work", "status": "Unread"},
            {"id": 2, "from": "bob@example.com", "subject": "Lunch Plans", "label": "Personal", "status": "Archived"},
            {"id": 3, "from": "carol@example.com", "subject": "Project Update", "label": "Work", "status": "Sent"},
            {"id": 4, "from": "dave@spam.com", "subject": "Win a Prize", "label": "Spam", "status": "Deleted"},
        ]

    def step(self, action: EmailAction) -> EmailObservation:
        feedback = ""
        target = next((e for e in self.emails if e["id"] == action.email_id), None) if action.email_id else None
        if action.action_type == "send":
            new_id = max([e["id"] for e in self.emails]+[0]) + 1
            self.emails.append({
                "id": new_id,
                "from": "me@example.com",
                "subject": "Project Update",
                "label": "Work",
                "status": "Sent",
                "body": action.message
            })
            feedback = "Action Executed: Send Email ✅"
        elif action.action_type == "reply" and target:
            new_id = max([e["id"] for e in self.emails]+[0]) + 1
            self.emails.append({
                "id": new_id,
                "from": "me@example.com",
                "subject": f"Re: {target['subject']}",
                "label": target["label"],
                "status": "Sent",
                "body": action.message
            })
            feedback = f"Action Executed: Reply ✅ to email ID {action.email_id}"
        elif action.action_type == "archive" and target:
            target["status"] = "Archived"
            feedback = f"Action Executed: Archive ✅ email ID {action.email_id}"
        elif action.action_type == "delete" and target:
            target["status"] = "Deleted"
            feedback = f"Action Executed: Delete ✅ email ID {action.email_id}"
        else:
            feedback = "Invalid action or email ID."
        return EmailObservation(emails=self.emails, ai_feedback=feedback)

env = EmailEnv()

status_colors = {
    "Unread": "orange",
    "Sent": "green",
    "Archived": "blue",
    "Deleted": "red"
}

def render_table_html(emails):
    table_html = "<table style='border-collapse: collapse; width: 100%;'>"
    table_html += "<tr><th>ID</th><th>From</th><th>Subject</th><th>Label</th><th>Status</th></tr>"
    for e in emails:
        color = status_colors.get(e["status"], "white")
        table_html += f"<tr style='border:1px solid #ddd;'>"
        table_html += f"<td>{e['id']}</td>"
        table_html += f"<td>{e['from']}</td>"
        table_html += f"<td>{e['subject']}</td>"
        table_html += f"<td>{e['label']}</td>"
        table_html += f"<td style='background-color:{color}; color:white;'>{e['status']}</td>"
        table_html += "</tr>"
    table_html += "</table>"
    return table_html

def gradio_step(action_type, message, email_id):
    if not action_type:
        return "Invalid action selected.", ""
    action_type_safe = str(action_type).lower()
    email_id_val = None
    if action_type_safe != "send" and email_id:
        email_id_str = str(email_id).strip()
        if email_id_str != "":
            try:
                email_id_val = int(email_id_str)
            except ValueError:
                return "Invalid Email ID. Must be a number.", ""
    obs = env.step(EmailAction(action_type=action_type_safe, message=message, email_id=email_id_val))
    table_html = render_table_html(obs.emails)
    return obs.ai_feedback, table_html

iface = gr.Interface(
    fn=gradio_step,
    inputs=[
        gr.Dropdown(['Send', 'Reply', 'Archive', 'Delete'], label="Choose Action"),
        gr.Textbox(label="Message"),
        gr.Textbox(label="Email ID (for reply/archive/delete)")
    ],
    outputs=[
        gr.Textbox(label="AI Feedback"),
        gr.HTML(label="Emails Dashboard")
    ],
    live=False,
    title="Smart Email Management AI Environment (OpenEnv)",
    description="Simulated email environment with AI-powered actions: send, reply, archive, delete."
)

iface.launch(server_name="0.0.0.0", server_port=7860)