app.py

"""
app.py — InboxOps Demo UI (Gradio, HuggingFace Spaces).

Provides a human-usable interface that lets anyone:
  1. Pick a scenario and difficulty
  2. See the inbox
  3. Manually triage emails or let the AI agent run
  4. See live scoring feedback
"""

from __future__ import annotations

import json
import os
import time
from typing import Optional

import gradio as gr
from openai import OpenAI

from env import InboxOpsEnv
from models import Action, ActionType, EmailCategory, Priority, SuggestedAction, TaskDifficulty


# ─────────────────────────────────────────────
# Globals
# ─────────────────────────────────────────────

API_BASE_URL = os.getenv("API_BASE_URL", "https://api.anthropic.com/v1")
MODEL_NAME   = os.getenv("MODEL_NAME",   "claude-sonnet-4-20250514")
HF_TOKEN     = os.getenv("HF_TOKEN",     "")

_env = InboxOpsEnv()
_current_obs = None
_ai_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "placeholder")

SYSTEM_PROMPT = """You are InboxOps Agent — an expert operations assistant managing
a startup founder's inbox. Triage each email and respond ONLY with valid JSON.

ACTION SCHEMA:
{
  "action_type": "classify",
  "email_id": "email_001",
  "category": "<investor|customer_support|partnership|personal|newsletter|notification|spam|press|internal|operational|customer_feedback|sales|unknown>",
  "priority": "<critical|high|medium|low>",
  "escalation_team": "<team or null>",
  "suggested_action": "<reply|reply_later|reply_immediately|escalate|escalate_immediately|route|archive|flag>",
  "draft_body": "<reply text or null>",
  "reply_tone": "<tone or null>",
  "notes": "<reasoning>"
}
When done: {"action_type": "complete"}
"""


# ─────────────────────────────────────────────
# Helper functions
# ─────────────────────────────────────────────

def format_inbox_html(obs) -> str:
    """Render the inbox as an HTML table."""
    if obs is None:
        return "<p style='color:#888'>No active episode. Click Start Episode.</p>"

    rows = ""
    for email in obs.emails:
        status_icon = "⏳" if email.id in obs.pending_email_ids else "✅"
        status_cls  = "pending" if email.id in obs.pending_email_ids else "done"
        rows += f"""
        <tr class="{status_cls}">
          <td class="status-cell">{status_icon}</td>
          <td class="id-cell">{email.id}</td>
          <td class="from-cell">{email.from_address}</td>
          <td class="subject-cell">{email.subject}</td>
          <td class="body-cell">{email.body[:120]}{'…' if len(email.body)>120 else ''}</td>
        </tr>"""

    return f"""
    <style>
      .inbox-table {{ width:100%; border-collapse:collapse; font-family:'IBM Plex Mono',monospace; font-size:12px; }}
      .inbox-table th {{ background:#1a1a2e; color:#e0e0ff; padding:8px 12px; text-align:left; }}
      .inbox-table td {{ padding:7px 12px; border-bottom:1px solid #2a2a3e; }}
      .inbox-table tr.pending {{ background:#0d1117; color:#e0e0e0; }}
      .inbox-table tr.done    {{ background:#0a1a0a; color:#5a8a5a; }}
      .inbox-table tr:hover   {{ background:#1e1e3a; }}
      .status-cell {{ width:30px; font-size:16px; }}
      .id-cell     {{ width:100px; color:#7878ff; font-weight:600; }}
      .from-cell   {{ width:200px; }}
      .subject-cell {{ width:240px; font-weight:500; }}
      .body-cell   {{ color:#aaa; }}
    </style>
    <table class="inbox-table">
      <thead><tr>
        <th></th><th>ID</th><th>From</th><th>Subject</th><th>Body Preview</th>
      </tr></thead>
      <tbody>{rows}</tbody>
    </table>
    """


def format_score_html(obs) -> str:
    if obs is None:
        return ""
    score = obs.score_so_far
    pct = int(score * 100)
    colour = "#22c55e" if score >= 0.7 else "#f59e0b" if score >= 0.4 else "#ef4444"
    pending = len(obs.pending_email_ids)
    processed = len(obs.processed_email_ids)
    return f"""
    <div style="font-family:'IBM Plex Mono',monospace;background:#111;border-radius:8px;padding:16px;margin-top:8px;">
      <div style="font-size:11px;color:#888;margin-bottom:4px;">RUNNING SCORE</div>
      <div style="font-size:36px;font-weight:700;color:{colour};">{pct}<span style="font-size:18px;">%</span></div>
      <div style="background:#222;border-radius:4px;height:8px;margin:8px 0;">
        <div style="background:{colour};border-radius:4px;height:8px;width:{pct}%;transition:width 0.4s;"></div>
      </div>
      <div style="font-size:11px;color:#888;">
        📧 {processed} processed · ⏳ {pending} pending · 🪜 {obs.step_budget_remaining} steps left
      </div>
    </div>
    """


def format_last_result(obs) -> str:
    if obs is None or not obs.last_action_result:
        return ""
    return obs.last_action_result


# ─────────────────────────────────────────────
# State actions
# ─────────────────────────────────────────────

def start_episode(scenario_id: str, difficulty: str):
    global _current_obs
    diff = {"Easy": TaskDifficulty.EASY, "Medium": TaskDifficulty.MEDIUM, "Hard": TaskDifficulty.HARD}[difficulty]
    _current_obs = _env.reset(scenario_id=scenario_id, task_difficulty=diff, seed=42)
    return (
        format_inbox_html(_current_obs),
        format_score_html(_current_obs),
        f"Episode started: {scenario_id} | Difficulty: {difficulty}",
        gr.update(interactive=True),
        gr.update(interactive=True),
    )


def manual_triage(email_id, category, priority, routing, action, draft):
    global _current_obs
    if _current_obs is None:
        return format_inbox_html(None), format_score_html(None), "No active episode.", ""

    act = Action(
        action_type=ActionType.CLASSIFY,
        email_id=email_id,
        category=EmailCategory(category) if category else None,
        priority=Priority(priority) if priority else None,
        escalation_team=routing if routing else None,
        suggested_action=SuggestedAction(action) if action else None,
        draft_body=draft if draft else None,
        notes="Manual triage",
    )

    _current_obs, reward, done, info = _env.step(act)
    result = (
        f"Step reward: {reward.total:.3f}\n"
        f"{reward.breakdown_notes}"
    )
    summary_txt = ""
    if done:
        summary = _env.episode_summary()
        summary_txt = f"\n\n🏁 EPISODE COMPLETE!\nFinal score: {summary['final_score']:.3f} (Grade {summary['letter_grade']})"

    return (
        format_inbox_html(_current_obs),
        format_score_html(_current_obs),
        result + summary_txt,
    )


def ai_triage_one():
    """Run one AI triage step."""
    global _current_obs
    if _current_obs is None:
        return format_inbox_html(None), format_score_html(None), "No active episode.", ""
    if _current_obs.done:
        return format_inbox_html(_current_obs), format_score_html(_current_obs), "Episode already complete.", ""
    if not _current_obs.pending_email_ids:
        return format_inbox_html(_current_obs), format_score_html(_current_obs), "No pending emails.", ""

    obs_text = _current_obs.to_text()
    user_msg = f"""OBSERVATION:\n{obs_text}\n\nProduce a JSON action for one pending email. Respond with ONLY valid JSON."""

    try:
        resp = _ai_client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_msg},
            ],
            max_tokens=1024,
            temperature=0.0,
        )
        raw = resp.choices[0].message.content or ""
        raw = raw.strip()
        if raw.startswith("```"):
            raw = "\n".join(l for l in raw.splitlines() if not l.strip().startswith("```"))
        data = json.loads(raw)
        action = Action.from_dict(data)
    except Exception as exc:
        return (
            format_inbox_html(_current_obs),
            format_score_html(_current_obs),
            f"AI error: {exc}",
        )

    _current_obs, reward, done, info = _env.step(action)
    result = (
        f"🤖 AI action: {action.action_type.value} on {action.email_id}\n"
        f"Category: {action.category} | Priority: {action.priority}\n"
        f"Routing: {action.escalation_team}\n"
        f"Step reward: {reward.total:.3f}\n"
        f"{reward.breakdown_notes}"
    )
    summary_txt = ""
    if done:
        summary = _env.episode_summary()
        summary_txt = f"\n\n🏁 EPISODE COMPLETE!\nFinal score: {summary['final_score']:.3f} (Grade {summary['letter_grade']})"

    return (
        format_inbox_html(_current_obs),
        format_score_html(_current_obs),
        result + summary_txt,
    )


def ai_triage_all():
    """Run AI agent to completion."""
    global _current_obs
    if _current_obs is None:
        return format_inbox_html(None), format_score_html(None), "No active episode."

    log_lines = []
    step = 0
    while not _current_obs.done:
        inbox_html, score_html, result = ai_triage_one()
        step += 1
        log_lines.append(f"Step {step}: {result.splitlines()[0]}")
        if _current_obs.done:
            break

    log_lines.append("\n🏁 All steps complete.")
    return (
        format_inbox_html(_current_obs),
        format_score_html(_current_obs),
        "\n".join(log_lines),
    )


# ─────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600&family=Space+Grotesk:wght@400;600;700&display=swap');

body, .gradio-container { background: #0a0a12 !important; color: #e0e0ff !important; }

.panel-box {
  background: #0f0f1e;
  border: 1px solid #2a2a4a;
  border-radius: 10px;
  padding: 16px;
}

h1.title {
  font-family: 'Space Grotesk', sans-serif;
  font-size: 2rem;
  font-weight: 700;
  color: #7878ff;
  letter-spacing: -0.5px;
  margin-bottom: 4px;
}

.subtitle {
  font-family: 'IBM Plex Mono', monospace;
  font-size: 12px;
  color: #5a5a8a;
  margin-bottom: 24px;
}

label { font-family: 'IBM Plex Mono', monospace !important; color: #a0a0d0 !important; }
.gr-button { font-family: 'IBM Plex Mono', monospace !important; }

.primary-btn {
  background: linear-gradient(135deg, #4a4aff 0%, #7878ff 100%) !important;
  color: white !important;
  border: none !important;
  font-weight: 600 !important;
}
.ai-btn {
  background: linear-gradient(135deg, #ff6b6b 0%, #ff9a9a 100%) !important;
  color: white !important;
  border: none !important;
}
.secondary-btn {
  background: #1a1a2e !important;
  color: #a0a0d0 !important;
  border: 1px solid #3a3a5a !important;
}

textarea, input, select {
  background: #0d0d1e !important;
  border: 1px solid #2a2a4a !important;
  color: #e0e0ff !important;
  font-family: 'IBM Plex Mono', monospace !important;
}
"""

with gr.Blocks(css=CSS, title="InboxOps — Founder Inbox AI Agent") as demo:

    gr.HTML("""
    <h1 class="title">📬 InboxOps</h1>
    <p class="subtitle">OpenEnv · Founder Operations Inbox · AI Agent Benchmark</p>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML('<div style="font-family:IBM Plex Mono;font-size:11px;color:#5a5a8a;margin-bottom:8px;">EPISODE SETUP</div>')
            scenario_dd = gr.Dropdown(
                choices=["scenario_001", "scenario_002"],
                value="scenario_001",
                label="Scenario",
            )
            difficulty_dd = gr.Dropdown(
                choices=["Easy", "Medium", "Hard"],
                value="Hard",
                label="Task Difficulty",
            )
            start_btn = gr.Button("▶ Start Episode", elem_classes=["primary-btn"])

        with gr.Column(scale=3):
            score_html = gr.HTML(label="Score")

    inbox_html = gr.HTML(label="Inbox", value="<p style='color:#5a5a8a;font-family:IBM Plex Mono;'>Click Start Episode to load the inbox.</p>")

    with gr.Row():
        with gr.Column(scale=2):
            gr.HTML('<div style="font-family:IBM Plex Mono;font-size:11px;color:#5a5a8a;margin:12px 0 8px;">MANUAL TRIAGE</div>')
            email_id_input  = gr.Textbox(label="Email ID", placeholder="email_001", interactive=False)
            category_input  = gr.Dropdown(
                choices=[c.value for c in EmailCategory],
                label="Category",
                interactive=False,
            )
            priority_input  = gr.Dropdown(
                choices=[p.value for p in Priority],
                label="Priority",
                interactive=False,
            )
            routing_input   = gr.Textbox(label="Escalation Team (or leave blank)", interactive=False)
            action_input    = gr.Dropdown(
                choices=[a.value for a in SuggestedAction],
                label="Suggested Action",
                interactive=False,
            )
            draft_input     = gr.Textbox(label="Draft Reply (Task 3 only)", lines=4, interactive=False)
            triage_btn      = gr.Button("Submit Triage", elem_classes=["secondary-btn"], interactive=False)

        with gr.Column(scale=1):
            gr.HTML('<div style="font-family:IBM Plex Mono;font-size:11px;color:#5a5a8a;margin:12px 0 8px;">AI AGENT</div>')
            ai_step_btn = gr.Button("🤖 AI: One Step", elem_classes=["ai-btn"], interactive=False)
            ai_all_btn  = gr.Button("🚀 AI: Run to End", elem_classes=["ai-btn"], interactive=False)
            gr.HTML('<div style="font-family:IBM Plex Mono;font-size:10px;color:#3a3a5a;margin-top:8px;">Uses the configured MODEL_NAME env var.</div>')

    result_box = gr.Textbox(label="Last Action Result", lines=6, interactive=False)

    # ── Event wiring ──
    start_btn.click(
        fn=start_episode,
        inputs=[scenario_dd, difficulty_dd],
        outputs=[inbox_html, score_html, result_box, triage_btn, ai_step_btn],
    ).then(
        fn=lambda: (gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True),
                    gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)),
        outputs=[email_id_input, category_input, priority_input, routing_input, action_input, draft_input],
    ).then(
        fn=lambda: gr.update(interactive=True),
        outputs=[ai_all_btn],
    )

    triage_btn.click(
        fn=manual_triage,
        inputs=[email_id_input, category_input, priority_input, routing_input, action_input, draft_input],
        outputs=[inbox_html, score_html, result_box],
    )

    ai_step_btn.click(
        fn=ai_triage_one,
        outputs=[inbox_html, score_html, result_box],
    )

    ai_all_btn.click(
        fn=ai_triage_all,
        outputs=[inbox_html, score_html, result_box],
    )

    gr.HTML("""
    <div style="font-family:'IBM Plex Mono',monospace;font-size:10px;color:#3a3a3a;margin-top:24px;border-top:1px solid #1a1a2e;padding-top:12px;">
      InboxOps · OpenEnv · Task difficulties: Easy (classify) | Medium (+priority+routing) | Hard (+draft replies)
    </div>
    """)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)