"""
app.py — OpenEnv Hybrid AI Email & Calendar Assistant
Streamlit Dashboard
"""

import os, sys, json, time
import streamlit as st
import pandas as pd

sys.path.insert(0, os.path.dirname(__file__))

from env import EmailEnv, Action, SAMPLE_EMAILS
from baseline import HybridAgent

# ─── Page config ──────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="OpenEnv AI Email Assistant",
    page_icon="📬",
    layout="wide",
)

st.markdown("""
<style>
.priority-urgent  { background:#FDEDED; color:#c62828; padding:2px 8px; border-radius:12px; font-size:12px; font-weight:600; }
.priority-high    { background:#FFF3E0; color:#E65100; padding:2px 8px; border-radius:12px; font-size:12px; font-weight:600; }
.priority-normal  { background:#E3F2FD; color:#1565C0; padding:2px 8px; border-radius:12px; font-size:12px; font-weight:600; }
.priority-low     { background:#E8F5E9; color:#2E7D32; padding:2px 8px; border-radius:12px; font-size:12px; font-weight:600; }
.action-box       { background:#F0FFF4; color:#22543D; border-left:4px solid #38A169; padding:12px; border-radius:6px; margin:8px 0; }
.reply-box        { background:#EBF8FF; color:#2B6CB0; border:1px solid #BEE3F8; padding:16px; border-radius:8px; font-family:monospace; }
.reward-pos       { background:#C6F6D5; color:#276749; padding:8px 14px; border-radius:8px; font-weight:700; font-size:18px; }
.reward-neg       { background:#FED7D7; color:#9B2C2C; padding:8px 14px; border-radius:8px; font-weight:700; font-size:18px; }
</style>
""", unsafe_allow_html=True)

# ─── Session state ─────────────────────────────────────────────────────────────

if "env" not in st.session_state:
    st.session_state.env = EmailEnv()
    st.session_state.env.reset()
if "agent" not in st.session_state:
    st.session_state.agent = HybridAgent()
if "results" not in st.session_state:
    st.session_state.results = []
if "selected_email" not in st.session_state:
    st.session_state.selected_email = 0
if "last_action" not in st.session_state:
    st.session_state.last_action = None
if "last_reward" not in st.session_state:
    st.session_state.last_reward = None

env: EmailEnv = st.session_state.env
agent: HybridAgent = st.session_state.agent
emails = env._emails

# ─── Sidebar ──────────────────────────────────────────────────────────────────

with st.sidebar:
    st.title("📊 Performance")
    if st.session_state.results:
        scores = [r["score"] for r in st.session_state.results]
        avg_score = sum(scores) / len(scores)
        correct = sum(1 for r in st.session_state.results if r["action"] == r["expected"])
        c1, c2 = st.columns(2)
        c1.metric("Avg Score", f"{avg_score:.2f}")
        c2.metric("Correct", f"{correct}/{len(scores)}")
    else:
        st.info("Run the agent to see metrics.")

    st.divider()
    st.subheader("🤖 Agent Mode")
    api_key = os.getenv("OPENAI_API_KEY", "")
    if api_key:
        st.success("OpenAI API active")
    else:
        st.warning("No API key — using HF/rules")

    mode_map = {"openai": "OpenAI GPT-4o-mini", "huggingface": "HuggingFace flan-t5", "rules": "Rule-based"}
    if st.session_state.last_action:
        st.info(f"Last mode: **{mode_map.get(agent.mode, agent.mode)}**")

    st.divider()
    if st.session_state.results:
        st.subheader("📈 Task Scores")
        for i, r in enumerate(st.session_state.results):
            color = "green" if r["score"] > 0 else "red"
            st.markdown(f"Task {r['task']}: **:{color}[{r['score']:.2f}]**")

    if st.button("🔄 Reset Environment"):
        env.reset()
        st.session_state.results = []
        st.session_state.last_action = None
        st.session_state.last_reward = None
        st.rerun()

# ─── Main Layout ───────────────────────────────────────────────────────────────

st.title("📬 OpenEnv Hybrid AI Email & Calendar Assistant")
st.caption("Powered by OpenAI GPT-4o-mini • HuggingFace flan-t5-base • Rule-based fallback")

col_inbox, col_main = st.columns([1, 2])

# ─── Inbox Panel ──────────────────────────────────────────────────────────────

with col_inbox:
    st.subheader("📬 Inbox")
    for i, email in enumerate(emails):
        p = email["priority"]
        badge_class = f"priority-{p}"
        is_selected = i == st.session_state.selected_email
        border = "2px solid #4299E1" if is_selected else "1px solid #E2E8F0"
        with st.container():
            clicked = st.button(
                f"{'📌' if p == 'urgent' else '📧'} {email['subject'][:35]}...",
                key=f"email_{i}",
                use_container_width=True,
            )
            if clicked:
                st.session_state.selected_email = i
                st.session_state.last_action = None
                st.session_state.last_reward = None
                st.rerun()
            st.markdown(
                f"<span class='{badge_class}'>{p}</span> · <small>{email['sender']}</small>",
                unsafe_allow_html=True,
            )
            st.markdown("---")

# ─── Main Panel ───────────────────────────────────────────────────────────────

with col_main:
    sel_idx = st.session_state.selected_email
    sel_email = emails[sel_idx]

    # Email Viewer
    st.subheader("📧 Email Viewer")
    with st.container():
        p = sel_email["priority"]
        st.markdown(
            f"**From:** {sel_email['sender']} &nbsp;&nbsp; "
            f"<span class='priority-{p}'>{p}</span>",
            unsafe_allow_html=True,
        )
        st.markdown(f"**Subject:** {sel_email['subject']}")
        st.markdown(f"**Received:** {sel_email['timestamp'][:19]}")
        st.divider()
        st.write(sel_email["body"])

    st.divider()

    # Agent Panel
    st.subheader("🤖 AI Agent Panel")
    run_btn = st.button("▶️ Run AI Assistant", type="primary", use_container_width=True)

    if run_btn:
        with st.spinner("Running hybrid AI agent (this may take a few seconds on CPU)..."):
            action_dict, mode = agent.decide(
                sender=sel_email["sender"],
                subject=sel_email["subject"],
                body=sel_email["body"],
                priority=sel_email["priority"],
            )
            action = Action(**action_dict)

            # Temporarily set env index to current email
            env._index = sel_idx
            _, reward, _, info = env.step(action)

            st.session_state.last_action = {**action_dict, "mode": mode}
            st.session_state.last_reward = {"score": reward.score, "reason": reward.reason}

            result_entry = {
                "task": sel_idx + 1,
                "subject": sel_email["subject"][:40],
                "difficulty": ["easy", "easy", "easy", "medium", "hard"][sel_idx],
                "expected": sel_email["expected_action"],
                "action": action.action_type,
                "score": reward.score,
                "mode": mode,
                "calendar_update": info.get("calendar_update"),
            }
            existing = [r for r in st.session_state.results if r["task"] != sel_idx + 1]
            st.session_state.results = existing + [result_entry]

    if st.session_state.last_action:
        la = st.session_state.last_action
        lr = st.session_state.last_reward
        at = la.get("action_type", "—")
        content = la.get("content", "")
        proposed = la.get("proposed_time", "")

        st.markdown(f"""
<div class='action-box'>
<b>Action:</b> <code>{at}</code> &nbsp;|&nbsp; <b>Mode:</b> {la.get('mode','—')} {f"&nbsp;|&nbsp; <b>Proposed time:</b> {proposed}" if proposed else ""}
</div>""", unsafe_allow_html=True)

        # Drafted reply
        st.markdown("#### ✉️ Drafted Reply")
        st.markdown(f"""<div class='reply-box'>
<b>Subject:</b> Re: {sel_email['subject']}<br><br>
{content}
</div>""", unsafe_allow_html=True)

        # Reward
        st.markdown("#### 🎯 Reward")
        score = lr["score"]
        css_class = "reward-pos" if score >= 0 else "reward-neg"
        st.markdown(
            f"<span class='{css_class}'>Score: {score:.2f}</span>&nbsp;&nbsp; {lr['reason']}",
            unsafe_allow_html=True,
        )

# ─── Calendar ─────────────────────────────────────────────────────────────────

st.divider()
st.subheader("📅 Calendar — Today")
slots = env.calendar.slots
cols = st.columns(10)
for i, (time_str, status) in enumerate(slots.items()):
    with cols[i % 10]:
        is_free = status == "free"
        st.markdown(
            f"<div style='text-align:center;padding:6px 4px;border-radius:8px;"
            f"background:{'#C6F6D5' if is_free else '#FED7D7'};"
            f"color:{'#276749' if is_free else '#9B2C2C'};"
            f"font-size:12px;font-weight:600;margin:2px;'>"
            f"{time_str}<br>{'FREE' if is_free else 'BUSY'}</div>",
            unsafe_allow_html=True,
        )

# ─── Inference Table + Charts ─────────────────────────────────────────────────

st.divider()
st.subheader("📊 Evaluation Results")

run_all = st.button("🚀 Run All Tasks (Full Inference)", use_container_width=True)
if run_all:
    with st.spinner("Running all 5 tasks..."):
        env2 = EmailEnv()
        obs = env2.reset()
        agent2 = HybridAgent()
        all_results = []
        idx = 0
        while obs is not None:
            e = env2._emails[idx]
            ad, mode = agent2.decide(obs.sender, obs.subject, obs.body, obs.priority)
            action = Action(**ad)
            _, reward, done, info = env2.step(action)
            all_results.append({
                "task": idx + 1,
                "subject": e["subject"][:40],
                "difficulty": ["easy", "easy", "easy", "medium", "hard"][idx],
                "expected": e["expected_action"],
                "action": action.action_type,
                "score": reward.score,
                "mode": mode,
            })
            obs = None if done else env2._make_observation()
            idx += 1
        st.session_state.results = all_results

if st.session_state.results:
    df = pd.DataFrame(st.session_state.results)
    df["match"] = df["action"] == df["expected"]
    st.dataframe(df[["task","subject","difficulty","expected","action","score","mode","match"]], use_container_width=True)

    avg = df["score"].mean()
    st.success(f"📈 Overall Average Score: **{avg:.3f}** | Correct: **{df['match'].sum()}/{len(df)}**")

    col_chart1, col_chart2 = st.columns(2)
    with col_chart1:
        st.markdown("**Score per task**")
        chart_df = df.set_index("task")[["score"]]
        st.line_chart(chart_df)
    with col_chart2:
        st.markdown("**Avg score by difficulty**")
        diff_df = df.groupby("difficulty")["score"].mean().reset_index()
        st.bar_chart(diff_df.set_index("difficulty"))

# ─── Workflow Panel ────────────────────────────────────────────────────────────

st.divider()
st.subheader("🧠 Agent Workflow")
steps = [
    ("📩", "Email received", f"From: {sel_email['sender']}"),
    ("🔍", "Intent detected", f"Priority: {sel_email['priority']} — analyzing content"),
    ("⚡", "Action selected", st.session_state.last_action.get("action_type", "Pending...") if st.session_state.last_action else "Pending..."),
    ("✍️", "Response generated", f"via {st.session_state.last_action.get('mode', '—')}" if st.session_state.last_action else "—"),
    ("🎯", "Reward assigned", f"{st.session_state.last_reward['score']:.2f}" if st.session_state.last_reward else "—"),
]
wf_cols = st.columns(5)
for col, (icon, title, detail) in zip(wf_cols, steps):
    with col:
        st.markdown(f"### {icon}")
        st.markdown(f"**{title}**")
        st.caption(detail)