streamlit_app.py

import streamlit as st
import os
import sys
import time
import traceback
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

# ─────────────────────────────────────────────────────────────────────────────
# Page Config  (MUST be the very first Streamlit call)
# ─────────────────────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="BlogForge AI – Agent Studio",
    page_icon="✍️",
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={"About": "BlogForge AI – Powered by LangGraph & Gemini"},
)

# ─────────────────────────────────────────────────────────────────────────────
# CSS
# ─────────────────────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&family=JetBrains+Mono:wght@400;500&display=swap');

html, body, [data-testid="stApp"] {
    background: #090b10 !important;
    font-family: 'Inter', sans-serif !important;
    color: #e2e8f0 !important;
}
[data-testid="stAppViewContainer"] {
    background:
        radial-gradient(ellipse 80% 55% at 50% -5%,  rgba(99,102,241,.16) 0%, transparent 65%),
        radial-gradient(ellipse 50% 40% at 85% 100%, rgba(16,185,129,.09) 0%, transparent 55%),
        #090b10;
}

/* ── Sidebar ── */
[data-testid="stSidebar"] {
    background: linear-gradient(180deg, #0d1017 0%, #111827 100%) !important;
    border-right: 1px solid rgba(99,102,241,.18) !important;
}
[data-testid="stSidebar"] * { color: #e2e8f0 !important; }
[data-testid="stSidebar"] .stTextArea textarea {
    background: #1a1f2e !important;
    border: 1px solid rgba(99,102,241,.35) !important;
    border-radius: 10px !important;
    color: #e2e8f0 !important;
    font-size: .9rem !important;
    font-family: 'Inter', sans-serif !important;
}
[data-testid="stSidebar"] .stTextArea textarea:focus {
    border-color: #6366f1 !important;
    box-shadow: 0 0 0 3px rgba(99,102,241,.2) !important;
}

/* ── Global text ── */
h1,h2,h3,h4,h5,h6,p,span,label,li { color: #e2e8f0 !important; }
code { background:#1a1f2e !important; color:#a5f3fc !important;
       border-radius:4px !important; padding:1px 5px !important;
       font-family:'JetBrains Mono',monospace !important; }
pre  { background:#0f172a !important; border:1px solid rgba(99,102,241,.2) !important;
       border-radius:10px !important; }

/* ── Run Button ── */
.stButton > button {
    background: linear-gradient(135deg,#6366f1,#8b5cf6) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 10px !important;
    font-weight: 600 !important;
    font-size: .95rem !important;
    padding: .65rem 1.6rem !important;
    width: 100% !important;
    transition: all .25s ease !important;
    box-shadow: 0 4px 22px rgba(99,102,241,.35) !important;
    font-family: 'Inter', sans-serif !important;
}
.stButton > button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 32px rgba(99,102,241,.55) !important;
    background: linear-gradient(135deg,#7c3aed,#6366f1) !important;
}
.stButton > button:active { transform: translateY(0) !important; }

/* ── Tabs ── */
.stTabs [data-baseweb="tab-list"] {
    background: #111827 !important;
    border-radius: 12px !important;
    padding: 4px !important;
    gap: 4px !important;
    border: 1px solid rgba(99,102,241,.2) !important;
}
.stTabs [data-baseweb="tab"] {
    background: transparent !important;
    border-radius: 9px !important;
    color: #94a3b8 !important;
    font-weight: 500 !important;
    padding: 8px 22px !important;
    border: none !important;
    font-family: 'Inter', sans-serif !important;
    transition: color .2s !important;
}
.stTabs [aria-selected="true"] {
    background: linear-gradient(135deg,#6366f1,#8b5cf6) !important;
    color: #fff !important;
    box-shadow: 0 2px 14px rgba(99,102,241,.4) !important;
}
.stTabs [data-baseweb="tab-panel"] { padding-top: 1.5rem !important; }

/* ── Metrics ── */
[data-testid="stMetric"] {
    background: #111827 !important;
    border: 1px solid rgba(99,102,241,.2) !important;
    border-radius: 12px !important;
    padding: 1rem 1.3rem !important;
}
[data-testid="stMetricLabel"] { color:#94a3b8 !important; font-size:.8rem !important; }
[data-testid="stMetricValue"] { color:#e2e8f0 !important; font-weight:700 !important; }

/* ── Scrollbar ── */
::-webkit-scrollbar { width:5px; height:5px; }
::-webkit-scrollbar-track { background:#090b10; }
::-webkit-scrollbar-thumb { background:rgba(99,102,241,.4); border-radius:4px; }

/* ── Hero ── */
.hero-header { text-align:center; padding:2.5rem 1rem 1rem; }
.hero-title {
    font-size:3rem; font-weight:900; letter-spacing:-1.5px;
    background:linear-gradient(135deg,#6366f1 0%,#a78bfa 45%,#10b981 100%);
    -webkit-background-clip:text; -webkit-text-fill-color:transparent; background-clip:text;
    margin:0; line-height:1.1;
}
.hero-sub { font-size:1rem; color:#64748b !important; margin-top:.55rem; font-weight:400; }

/* ── Node pipeline ── */
.pipeline-flow {
    display:flex; align-items:center; justify-content:center;
    gap:.35rem; flex-wrap:wrap; padding:1rem 0 1.5rem;
}
.pipeline-node { display:flex; flex-direction:column; align-items:center; gap:.3rem; }
.pipeline-bubble {
    width:46px; height:46px; border-radius:50%;
    display:flex; align-items:center; justify-content:center;
    font-size:1.15rem;
    border:2px solid rgba(99,102,241,.25);
    background:#1a1f2e;
    transition:all .4s ease;
}
.pipeline-bubble.active {
    border-color:#6366f1; background:rgba(99,102,241,.22);
    box-shadow:0 0 20px rgba(99,102,241,.55);
    animation:glow-pulse 1.4s ease infinite;
}
.pipeline-bubble.done {
    border-color:#10b981; background:rgba(16,185,129,.15);
    box-shadow:0 0 10px rgba(16,185,129,.25);
}
.pipeline-bubble.error {
    border-color:#ef4444; background:rgba(239,68,68,.12);
}
@keyframes glow-pulse {
    0%,100% { box-shadow:0 0 20px rgba(99,102,241,.55); }
    50%      { box-shadow:0 0 35px rgba(99,102,241,.85); }
}
.pipeline-label {
    font-size:.58rem; color:#64748b !important; font-weight:700;
    text-transform:uppercase; letter-spacing:.05em;
    text-align:center; max-width:58px;
}
.pipeline-arrow { font-size:.9rem; color:rgba(99,102,241,.35) !important; margin-bottom:16px; }

/* ── Node cards ── */
.node-card {
    background:#111827;
    border:1px solid rgba(99,102,241,.18);
    border-radius:14px;
    padding:1.1rem 1.4rem;
    margin-bottom:.7rem;
    transition:all .35s ease;
}
.node-card.active {
    border-color:#6366f1;
    box-shadow:0 0 22px rgba(99,102,241,.22), inset 0 0 22px rgba(99,102,241,.03);
    animation:card-pulse 1.6s ease infinite;
}
.node-card.done  { border-color:rgba(16,185,129,.45); box-shadow:0 0 10px rgba(16,185,129,.1); }
.node-card.error { border-color:rgba(239,68,68,.45);  box-shadow:0 0 10px rgba(239,68,68,.1); }
@keyframes card-pulse {
    0%,100% { box-shadow:0 0 22px rgba(99,102,241,.22); }
    50%      { box-shadow:0 0 35px rgba(99,102,241,.42); }
}
.node-header { display:flex; align-items:center; gap:.6rem; margin-bottom:.4rem; }
.node-icon { font-size:1.25rem; }
.node-name { font-weight:700; font-size:.95rem; color:#e2e8f0 !important; }
.node-status-badge {
    margin-left:auto; font-size:.68rem; font-weight:700;
    padding:2px 10px; border-radius:20px;
}
.badge-waiting { background:rgba(100,116,139,.18); color:#64748b !important; }
.badge-running { background:rgba(99,102,241,.18);  color:#a78bfa !important; }
.badge-done    { background:rgba(16,185,129,.15);  color:#34d399 !important; }
.badge-error   { background:rgba(239,68,68,.15);   color:#f87171 !important; }
.badge-skipped { background:rgba(234,179,8,.15);   color:#fbbf24 !important; }
.node-detail {
    font-size:.78rem; color:#64748b !important; margin:0;
    font-family:'JetBrains Mono',monospace; line-height:1.65;
    max-height:100px; overflow-y:auto; white-space:pre-wrap; word-break:break-word;
}

/* ── Details panel cards ── */
.detail-card {
    background:#111827;
    border:1px solid rgba(99,102,241,.18);
    border-radius:12px;
    padding:.9rem 1.1rem;
    margin-bottom:.65rem;
}
.detail-card-title {
    font-size:.68rem; font-weight:700; color:#6366f1 !important;
    font-family:'JetBrains Mono',monospace;
    text-transform:uppercase; letter-spacing:.06em; margin-bottom:.4rem;
}

/* ── Log box ── */
.log-box {
    background:#090b10; border:1px solid rgba(99,102,241,.15);
    border-radius:12px; padding:1rem;
    max-height:520px; overflow-y:auto;
    font-family:'JetBrains Mono',monospace;
}
.log-line { font-size:.77rem; padding:1px 0; border-bottom:1px solid rgba(255,255,255,.025); line-height:1.75; }
.log-ts   { color:#334155; margin-right:.55rem; }
.log-info { color:#94a3b8 !important; }
.log-success { color:#34d399 !important; }
.log-warn    { color:#fbbf24 !important; }
.log-error   { color:#f87171 !important; }
.log-node    { color:#a78bfa !important; font-weight:600; }
.log-llm     { color:#38bdf8 !important; }
.log-tool    { color:#fb923c !important; }

/* ── API badges ── */
.api-badge {
    display:inline-flex; align-items:center; gap:.35rem;
    padding:.28rem .65rem; border-radius:20px;
    font-size:.73rem; font-weight:600; margin:.18rem;
}
.api-ok  { background:rgba(16,185,129,.13); color:#34d399 !important; border:1px solid rgba(16,185,129,.28); }
.api-err { background:rgba(239,68,68,.12);  color:#f87171 !important; border:1px solid rgba(239,68,68,.28); }

/* ── Model pills ── */
.model-pill {
    display:inline-flex; align-items:center; gap:.32rem;
    background:rgba(99,102,241,.1); border:1px solid rgba(99,102,241,.28);
    color:#a78bfa !important; padding:.22rem .6rem; border-radius:20px;
    font-size:.7rem; font-weight:600; margin:.14rem;
    font-family:'JetBrains Mono',monospace;
}

/* ── Blog preview ── */
.blog-preview {
    background:#111827; border:1px solid rgba(99,102,241,.18);
    border-radius:16px; padding:2rem 2.5rem; line-height:1.82;
}
.blog-preview h1 { font-size:2.1rem !important; font-weight:900 !important;
    background:linear-gradient(135deg,#6366f1,#a78bfa); -webkit-background-clip:text;
    -webkit-text-fill-color:transparent; background-clip:text; }
.blog-preview h2 { font-size:1.3rem !important; font-weight:700 !important;
    color:#c4b5fd !important; margin-top:2rem !important;
    border-bottom:1px solid rgba(99,102,241,.2); padding-bottom:.3rem; }
.blog-preview h3 { font-size:1.05rem !important; font-weight:600 !important; color:#a78bfa !important; }
.blog-preview p  { color:#cbd5e1 !important; font-size:.97rem; }
.blog-preview pre { background:#0f172a !important; border:1px solid rgba(99,102,241,.18) !important;
    border-radius:10px !important; padding:1rem !important; overflow-x:auto; }
.blog-preview blockquote { border-left:3px solid #6366f1 !important; padding-left:1rem;
    color:#94a3b8 !important; margin:1rem 0; background:rgba(99,102,241,.05); border-radius:0 8px 8px 0; }
.blog-preview img { max-width:100%; border-radius:12px;
    border:1px solid rgba(99,102,241,.2); margin:1rem 0; }
.blog-preview ul,.blog-preview ol { color:#cbd5e1 !important; padding-left:1.5rem; }
.blog-preview li { margin-bottom:.3rem; }
.blog-preview table { border-collapse:collapse; width:100%; margin:1rem 0; }
.blog-preview th { background:rgba(99,102,241,.15); color:#c4b5fd !important;
    padding:.5rem .8rem; border:1px solid rgba(99,102,241,.2); font-size:.85rem; }
.blog-preview td { padding:.45rem .8rem; border:1px solid rgba(255,255,255,.06);
    color:#cbd5e1 !important; font-size:.85rem; }

/* ── Stat boxes ── */
.stat-box {
    background:linear-gradient(135deg,rgba(99,102,241,.1),rgba(139,92,246,.05));
    border:1px solid rgba(99,102,241,.22); border-radius:12px;
    padding:1rem 1.2rem; text-align:center;
}
.stat-val   { font-size:1.8rem; font-weight:800;
    background:linear-gradient(135deg,#6366f1,#a78bfa);
    -webkit-background-clip:text; -webkit-text-fill-color:transparent; background-clip:text; }
.stat-label { font-size:.72rem; color:#64748b !important; font-weight:700;
    text-transform:uppercase; letter-spacing:.06em; margin-top:.18rem; }

.divider { border:none; border-top:1px solid rgba(99,102,241,.12); margin:1.3rem 0; }

/* ── Progress bar override ── */
[data-testid="stProgress"] > div > div { background:linear-gradient(90deg,#6366f1,#8b5cf6) !important; border-radius:10px !important; }

/* ── Download button ── */
[data-testid="stDownloadButton"] > button {
    background:rgba(16,185,129,.15) !important;
    color:#34d399 !important;
    border:1px solid rgba(16,185,129,.35) !important;
    border-radius:10px !important;
    font-weight:600 !important;
}
[data-testid="stDownloadButton"] > button:hover {
    background:rgba(16,185,129,.28) !important;
    transform:translateY(-1px) !important;
}
</style>
""", unsafe_allow_html=True)

# ─────────────────────────────────────────────────────────────────────────────
# Session State
# ─────────────────────────────────────────────────────────────────────────────
_DEFAULTS = {
    "running": False,
    "done": False,
    "error": None,
    "logs": [],
    "node_states": {
        "router":       {"status": "waiting", "detail": ""},
        "research":     {"status": "waiting", "detail": ""},
        "orchestrator": {"status": "waiting", "detail": ""},
        "worker":       {"status": "waiting", "detail": ""},
        "reducer":      {"status": "waiting", "detail": ""},
    },
    "router_decision": None,
    "evidence": [],
    "plan": None,
    "sections_done": 0,
    "sections_total": 0,
    "image_specs": [],
    "final_md": "",
    "elapsed": 0.0,
    "start_time": None,
    "last_topic": "",
    "current_model": "",
}
for _k, _v in _DEFAULTS.items():
    if _k not in st.session_state:
        st.session_state[_k] = _v

# ─────────────────────────────────────────────────────────────────────────────
# Constants
# ─────────────────────────────────────────────────────────────────────────────
PROVIDER_ICON = {"google": "✦", "groq": "⚡", "openrouter": "🌐"}
NODE_META = {
    "router":       {"icon": "🔀", "label": "Router",       "desc": "Classifies topic & decides if web research is needed"},
    "research":     {"icon": "🔍", "label": "Research",     "desc": "Runs Tavily API searches and synthesizes evidence"},
    "orchestrator": {"icon": "🗂️", "label": "Orchestrator", "desc": "Plans sections, goals, and word counts"},
    "worker":       {"icon": "✍️", "label": "Workers",      "desc": "Writes each section in parallel fan-out"},
    "reducer":      {"icon": "🔧", "label": "Reducer",      "desc": "Merges sections → plans images → generates with Gemini"},
}
LOG_CSS = {"info":"log-info","success":"log-success","warn":"log-warn","error":"log-error","node":"log-node","llm":"log-llm","tool":"log-tool"}

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────
def _add_log(msg: str, kind: str = "info"):
    if "logs" not in st.session_state:
        st.session_state.logs = []
    st.session_state.logs.append({"ts": time.strftime("%H:%M:%S"), "msg": msg, "kind": kind})

def _set_node(name: str, status: str, detail: str = ""):
    st.session_state.node_states[name] = {"status": status, "detail": detail}

def _api_badge(env_key: str, label: str, icon: str) -> str:
    v = os.getenv(env_key, "")
    ok = bool(v and len(v) > 8)
    masked = v[:6] + "••••" + v[-3:] if ok else "Not set"
    cls = "api-ok" if ok else "api-err"
    dot = "🟢" if ok else "🔴"
    return f'<span class="api-badge {cls}">{dot} {icon} {label}: <code style="font-size:.68rem">{masked}</code></span>'

def _node_card_html(name: str) -> str:
    meta = NODE_META[name]
    ns   = st.session_state.node_states[name]
    status, detail = ns["status"], ns["detail"]
    badge_map = {
        "waiting": ("WAITING",   "badge-waiting"),
        "running": ("RUNNING ●", "badge-running"),
        "done":    ("DONE ✓",   "badge-done"),
        "error":   ("ERROR ✗",  "badge-error"),
        "skipped": ("SKIPPED",  "badge-skipped"),
    }
    bt, bc = badge_map.get(status, ("–","badge-waiting"))
    card_cls = {"running":"active","done":"done","error":"error"}.get(status,"")
    det_html = f'<p class="node-detail">{detail}</p>' if detail else ""
    return f"""
<div class="node-card {card_cls}">
  <div class="node-header">
    <span class="node-icon">{meta['icon']}</span>
    <span class="node-name">{meta['label']}</span>
    <span class="node-status-badge {bc}">{bt}</span>
  </div>
  <div style="font-size:.75rem;color:#475569;">{meta['desc']}</div>
  {det_html}
</div>"""

def _pipeline_html() -> str:
    nodes = ["router","research","orchestrator","worker","reducer"]
    parts = []
    for i, n in enumerate(nodes):
        meta = NODE_META[n]
        s = st.session_state.node_states[n]["status"]
        bcls = "active" if s=="running" else ("done" if s=="done" else ("error" if s=="error" else ""))
        parts.append(f"""
<div class="pipeline-node">
  <div class="pipeline-bubble {bcls}">{meta['icon']}</div>
  <div class="pipeline-label">{meta['label']}</div>
</div>""")
        if i < len(nodes)-1:
            parts.append('<div class="pipeline-arrow">›</div>')
    return f'<div class="pipeline-flow">{"".join(parts)}</div>'

def _details_html() -> str:
    parts = []

    # Router decision
    rd = st.session_state.router_decision
    if rd:
        mc = {"closed_book":"#10b981","hybrid":"#f59e0b","open_book":"#6366f1"}.get(rd.mode,"#94a3b8")
        q_rows = "".join(
            f'<div style="font-size:.75rem;color:#64748b;font-family:JetBrains Mono,monospace;margin-top:.25rem;">› {q}</div>'
            for q in rd.queries
        )
        parts.append(f"""
<div class="detail-card">
  <div class="detail-card-title">🔀 Router Decision</div>
  <div style="display:flex;gap:.5rem;flex-wrap:wrap;align-items:center;margin-bottom:.35rem;">
    <span style="background:rgba(99,102,241,.12);border-radius:7px;padding:.15rem .55rem;font-size:.78rem;color:{mc};font-weight:700;border:1px solid {mc}33;">{rd.mode.upper()}</span>
    <span style="font-size:.78rem;color:#94a3b8;">Research: {'✅ Yes' if rd.needs_research else '❌ No'}</span>
    <span style="font-size:.78rem;color:#94a3b8;">Queries: {len(rd.queries)}</span>
  </div>
  {q_rows}
</div>""")

    # Evidence
    ev = st.session_state.evidence
    if ev:
        rows = "".join(
            f'<div style="margin-bottom:.4rem;padding:.4rem 0;border-bottom:1px solid rgba(255,255,255,.04);">'
            f'<div style="font-size:.78rem;font-weight:600;color:#c4b5fd;">{e.title[:65]}{"…" if len(e.title)>65 else ""}</div>'
            f'<div style="font-size:.7rem;color:#64748b;">{e.source or ""} · {e.published_at or "n/a"}</div>'
            f'<div style="font-size:.7rem;color:#475569;margin-top:.12rem;">{(e.snippet or "")[:100]}{"…" if len(e.snippet or "")>100 else ""}</div>'
            f'</div>'
            for e in ev[:6]
        )
        parts.append(f"""
<div class="detail-card">
  <div class="detail-card-title">🔍 Research Evidence ({len(ev)} items)</div>
  <div style="max-height:190px;overflow-y:auto;">{rows}</div>
</div>""")

    # Plan
    plan = st.session_state.plan
    if plan:
        ki = {"explainer":"📖","tutorial":"🛠️","news_roundup":"📰","comparison":"⚖️","system_design":"🏗️"}.get(plan.blog_kind,"📝")
        sec_rows = "".join(
            f'<div style="margin-bottom:.4rem;padding:.5rem .7rem;background:#0f172a;border-radius:8px;border:1px solid rgba(99,102,241,.12);">'
            f'<div style="font-size:.65rem;font-weight:700;color:#6366f1;font-family:JetBrains Mono,monospace;">§{t.id}'
            f' {"🔍" if t.requires_research else ""}{"💻" if t.requires_code else ""}{"🔗" if t.requires_citations else ""}</div>'
            f'<div style="font-size:.82rem;font-weight:600;color:#e2e8f0;">{t.title}</div>'
            f'<div style="font-size:.7rem;color:#64748b;margin-top:.1rem;">{t.target_words}w · {len(t.bullets)} bullets</div>'
            f'</div>'
            for t in plan.tasks
        )
        parts.append(f"""
<div class="detail-card">
  <div class="detail-card-title">{ki} Blog Plan · {plan.blog_kind.replace("_"," ").title()}</div>
  <div style="font-size:1rem;font-weight:800;color:#e2e8f0;margin:.2rem 0 .12rem;">{plan.blog_title}</div>
  <div style="font-size:.78rem;color:#64748b;margin-bottom:.6rem;">👥 {plan.audience} · 🗣️ {plan.tone}</div>
  <div style="max-height:250px;overflow-y:auto;">{sec_rows}</div>
</div>""")

    # Section progress
    total = st.session_state.sections_total
    done_n = st.session_state.sections_done
    if total > 0:
        pct = int((done_n/total)*100)
        bc = "#10b981" if done_n==total else "#6366f1"
        parts.append(f"""
<div class="detail-card">
  <div class="detail-card-title">✍️ Section Writers ({done_n}/{total})</div>
  <div style="display:flex;justify-content:space-between;font-size:.75rem;color:#94a3b8;margin:.3rem 0 .25rem;">
    <span>{done_n} of {total} written</span><span>{pct}%</span>
  </div>
  <div style="background:#1a1f2e;border-radius:20px;height:7px;overflow:hidden;">
    <div style="background:linear-gradient(90deg,{bc},{bc}bb);width:{pct}%;height:100%;border-radius:20px;transition:width .5s ease;"></div>
  </div>
</div>""")

    # Image specs
    specs = st.session_state.image_specs
    if specs:
        sp_rows = "".join(
            f'<div style="margin-bottom:.45rem;padding:.45rem .65rem;background:#0f172a;border-radius:8px;border:1px solid rgba(99,102,241,.12);">'
            f'<div style="font-size:.72rem;font-weight:700;color:#6366f1;font-family:JetBrains Mono,monospace;">{s.get("placeholder","")}</div>'
            f'<div style="font-size:.76rem;color:#94a3b8;margin-top:.15rem;">{s.get("alt_text","")[:70]}</div>'
            f'<div style="font-size:.68rem;color:#475569;margin-top:.1rem;font-style:italic;">{(s.get("caption",""))[:80]}</div>'
            f'</div>'
            for s in specs
        )
        parts.append(f"""
<div class="detail-card">
  <div class="detail-card-title">🖼️ Image Plan ({len(specs)} images)</div>
  {sp_rows}
</div>""")

    if not parts:
        return '<div style="color:#475569;font-size:.85rem;padding:1.5rem;text-align:center;">Run the agent to see live details here&hellip;</div>'
    return "".join(parts)

def _logs_html() -> str:
    if not st.session_state.logs:
        return '<div style="color:#475569;font-size:.82rem;padding:.5rem;">No logs yet…</div>'
    lines = []
    for e in st.session_state.logs[-250:]:
        cls = LOG_CSS.get(e["kind"], "log-info")
        lines.append(f'<div class="log-line {cls}"><span class="log-ts">[{e["ts"]}]</span>{e["msg"]}</div>')
    return f'<div class="log-box">{"".join(lines)}</div>'

# ─────────────────────────────────────────────────────────────────────────────
# Sidebar
# ─────────────────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("""
<div style="text-align:center;padding:1.1rem 0 .6rem;">
  <div style="font-size:3rem;line-height:1">✍️</div>
  <div style="font-size:1.25rem;font-weight:900;background:linear-gradient(135deg,#6366f1,#a78bfa);
       -webkit-background-clip:text;-webkit-text-fill-color:transparent;background-clip:text;letter-spacing:-.5px;">BlogForge AI</div>
  <div style="font-size:.7rem;color:#475569;margin-top:.2rem;">Powered by LangGraph & Gemini</div>
</div>
""", unsafe_allow_html=True)

    st.markdown('<hr class="divider">', unsafe_allow_html=True)

    # API status
    st.markdown("**🔑 API Keys**")
    st.markdown(
        _api_badge("GOOGLE_API_KEY",      "Gemini",     "✦") + "<br>" +
        _api_badge("GROQ_API_KEY",        "Groq",       "⚡") + "<br>" +
        _api_badge("OPENROUTER_API_KEY",  "OpenRouter",  "🌐") + "<br>" +
        _api_badge("TAVILY_API_KEY",      "Tavily",     "🔍"),
        unsafe_allow_html=True
    )

    st.markdown('<hr class="divider">', unsafe_allow_html=True)

    # Topic
    st.markdown("**📝 Blog Topic**")
    topic_input = st.text_area(
        "topic",
        value=st.session_state.last_topic,
        height=115,
        placeholder="e.g. How Kubernetes handles pod scheduling\n\nBe specific for best results!",
        label_visibility="collapsed",
    )

    # Model fallback chain
    st.markdown("**🤖 Model Fallback Chain**")
    from llm_manager import FALLBACK_MODELS
    for i, m in enumerate(FALLBACK_MODELS):
        icon = PROVIDER_ICON.get(m["provider"], "🤖")
        active = st.session_state.current_model == f"{m['provider']}/{m['model']}"
        active_style = "border-color:#6366f1;background:rgba(99,102,241,.22);" if active else ""
        st.markdown(
            f'<span class="model-pill" style="{active_style}">{"▶ " if active else ""}{icon} #{i+1} {m["model"]}</span>',
            unsafe_allow_html=True,
        )

    st.markdown('<hr class="divider">', unsafe_allow_html=True)

    # Run button
    run_disabled = st.session_state.running or not topic_input.strip()
    label = "⏳  Generating…" if st.session_state.running else "🚀  Generate Blog Post"
    run_clicked = st.button(label, disabled=run_disabled, use_container_width=True)

    if st.session_state.done and not st.session_state.running:
        if st.button("🔄  Start New Blog", use_container_width=True):
            for k in list(st.session_state.keys()):
                del st.session_state[k]
            st.rerun()

    # Live stats
    if st.session_state.running or st.session_state.done:
        st.markdown('<hr class="divider">', unsafe_allow_html=True)
        el = st.session_state.elapsed
        st.markdown(f"""
<div style="display:flex;gap:.55rem;flex-wrap:wrap;">
  <div class="stat-box" style="flex:1">
    <div class="stat-val">{int(el)}s</div>
    <div class="stat-label">Elapsed</div>
  </div>
  <div class="stat-box" style="flex:1">
    <div class="stat-val">{st.session_state.sections_done}</div>
    <div class="stat-label">Sections</div>
  </div>
</div>
<div style="height:.45rem"></div>
<div class="stat-box">
  <div class="stat-val">{len(st.session_state.image_specs)}</div>
  <div class="stat-label">Images</div>
</div>
""", unsafe_allow_html=True)

# ─────────────────────────────────────────────────────────────────────────────
# Main Content
# ─────────────────────────────────────────────────────────────────────────────
st.markdown("""
<div class="hero-header">
  <h1 class="hero-title">BlogForge AI Agent Studio</h1>
  <p class="hero-sub">LangGraph Multi-Agent · Gemini + Groq + OpenRouter Fallback · Tavily Research · AI Image Generation</p>
</div>
""", unsafe_allow_html=True)

# Pipeline visual
pipeline_ph = st.empty()
pipeline_ph.markdown(_pipeline_html(), unsafe_allow_html=True)

# Tabs
tab_exec, tab_blog, tab_logs = st.tabs([
    "⚡  Live Execution",
    "📄  Blog Output",
    "📋  Agent Logs",
])

# ── Live Execution tab layout ──────────────────────────────────────────────
with tab_exec:
    col_nodes, col_details = st.columns([1, 1.45])
    with col_nodes:
        st.markdown("### 🔷 Agent Pipeline")
        node_phs = {n: st.empty() for n in ["router","research","orchestrator","worker","reducer"]}
    with col_details:
        st.markdown("### 📊 Live Details")
        details_ph = st.empty()

# ── Blog Output tab ────────────────────────────────────────────────────────
with tab_blog:
    blog_ph = st.empty()

# ── Logs tab ───────────────────────────────────────────────────────────────
with tab_logs:
    logs_ph = st.empty()

# ─────────────────────────────────────────────────────────────────────────────
# Render helpers
# ─────────────────────────────────────────────────────────────────────────────
def _render_all():
    pipeline_ph.markdown(_pipeline_html(), unsafe_allow_html=True)
    for n, ph in node_phs.items():
        ph.markdown(_node_card_html(n), unsafe_allow_html=True)
    details_ph.markdown(_details_html(), unsafe_allow_html=True)
    logs_ph.markdown(_logs_html(), unsafe_allow_html=True)

def _render_blog():
    fm = st.session_state.final_md
    if not fm:
        with blog_ph.container():
            st.markdown("""
<div style="text-align:center;padding:3.5rem 1rem;color:#475569;">
  <div style="font-size:4rem;margin-bottom:1rem">✍️</div>
  <div style="font-size:1.15rem;font-weight:600;color:#64748b;">Your blog will appear here</div>
  <div style="font-size:.88rem;margin-top:.4rem;">Enter a topic in the sidebar and click <strong>Generate Blog Post</strong></div>
</div>
""", unsafe_allow_html=True)
        return

    words    = len(fm.split())
    img_cnt  = fm.count("![")
    sec_cnt  = fm.count("\n## ")
    el       = st.session_state.elapsed
    plan     = st.session_state.plan

    with blog_ph.container():
        m1, m2, m3, m4 = st.columns(4)
        m1.metric("📝 Words",    f"{words:,}")
        m2.metric("📸 Images",   img_cnt)
        m3.metric("📑 Sections", sec_cnt)
        m4.metric("⏱️ Time",     f"{el:.0f}s")

        st.markdown('<hr class="divider">', unsafe_allow_html=True)

        fname = (plan.blog_title if plan else "blog").replace(" ","_").replace(":","").replace("?","")
        st.download_button(
            "⬇️  Download Markdown",
            data=fm.encode("utf-8"),
            file_name=f"{fname}.md",
            mime="text/markdown",
        )
        st.markdown('<div style="height:.6rem"></div>', unsafe_allow_html=True)
        st.markdown('<div class="blog-preview">', unsafe_allow_html=True)
        st.markdown(fm, unsafe_allow_html=True)
        st.markdown('</div>', unsafe_allow_html=True)

        # Generated images gallery
        images_dir = Path("images")
        if images_dir.exists():
            img_files = (
                list(images_dir.glob("*.png"))
                + list(images_dir.glob("*.jpg"))
                + list(images_dir.glob("*.jpeg"))
            )
            if img_files:
                st.markdown('<hr class="divider">', unsafe_allow_html=True)
                st.markdown("#### 🖼️ Generated Images")
                gcols = st.columns(min(len(img_files), 3))
                for i, p in enumerate(img_files[:3]):
                    with gcols[i]:
                        st.image(str(p), caption=p.stem.replace("_"," ").title(), use_container_width=True)

# ─────────────────────────────────────────────────────────────────────────────
# Execution logic
# ─────────────────────────────────────────────────────────────────────────────
if run_clicked and topic_input.strip():
    # Reset
    st.session_state.update({
        "running": True, "done": False, "error": None,
        "logs": [], "final_md": "", "image_specs": [],
        "evidence": [], "plan": None, "router_decision": None,
        "sections_done": 0, "sections_total": 0,
        "elapsed": 0.0, "start_time": time.time(),
        "last_topic": topic_input.strip(), "current_model": "",
        "node_states": {n: {"status":"waiting","detail":""} for n in NODE_META},
    })

    _add_log("🚀 BlogForge AI Agent started", "node")
    _add_log(f"📝 Topic: {topic_input.strip()}", "info")
    _render_all()

    # Monkey-patch LLM manager to capture model events
    import llm_manager as _lm
    _orig_inv = _lm.FallbackLLMManager.invoke
    _orig_str = _lm.StructuredFallbackLLM.invoke

    def _p_inv(self, messages, **kw):
        mi = _lm.FALLBACK_MODELS[self.current_model_idx]
        mn = f"{mi['provider']}/{mi['model']}"
        st.session_state.current_model = mn
        _add_log(f"🤖 LLM invoke → {mn}", "llm")
        return _orig_inv(self, messages, **kw)

    def _p_str(self, messages, **kw):
        mi = _lm.FALLBACK_MODELS[self.manager.current_model_idx]
        mn = f"{mi['provider']}/{mi['model']}"
        _add_log(f"📐 Structured({self.schema.__name__}) → {mn}", "llm")
        return _orig_str(self, messages, **kw)

    _lm.FallbackLLMManager.invoke = _p_inv
    _lm.StructuredFallbackLLM.invoke = _p_str

    try:
        from blog_agent import app

        _add_log("📡 Connecting to LangGraph…", "info")

        for event in app.stream({"topic": topic_input.strip(), "sections": []}, stream_mode="updates"):
            st.session_state.elapsed = time.time() - st.session_state.start_time

            for node_name, upd in event.items():

                # ── Router ──────────────────────────────────────────────────
                if node_name == "router":
                    _set_node("router", "running", "Analyzing topic…")
                    _add_log("🔀 Router node executing…", "node")
                    _render_all()
                    rd = upd.get("router_decision")
                    if rd:
                        st.session_state.router_decision = rd
                        det = f"Mode: {rd.mode} | Research: {rd.needs_research}"
                        if rd.queries:
                            det += "\n" + "\n".join(f"› {q}" for q in rd.queries[:5])
                        _set_node("router", "done", det)
                        _add_log(f"✅ Router → mode={rd.mode}, research={rd.needs_research}", "success")
                        if rd.needs_research:
                            _add_log(f"🔍 {len(rd.queries)} search queries queued", "tool")
                        else:
                            _set_node("research", "skipped", "No research needed (closed_book)")
                            _add_log("⏭️ Research skipped (closed_book)", "warn")

                # ── Research ────────────────────────────────────────────────
                elif node_name == "research":
                    _set_node("research", "running", "Querying Tavily API…")
                    _add_log("🔍 Research node: fetching evidence via Tavily…", "node")
                    _render_all()
                    ev = upd.get("evidence", [])
                    st.session_state.evidence = ev
                    _set_node("research", "done", f"Retrieved {len(ev)} evidence items")
                    _add_log(f"✅ Research done → {len(ev)} items", "success")
                    for e in ev[:3]:
                        _add_log(f"  📄 {e.title[:50]} ({e.source})", "tool")

                # ── Orchestrator ─────────────────────────────────────────────
                elif node_name == "orchestrator":
                    _set_node("orchestrator", "running", "Building content plan…")
                    _add_log("🗂️ Orchestrator: planning blog…", "node")
                    _render_all()
                    plan = upd.get("plan")
                    if plan:
                        st.session_state.plan = plan
                        st.session_state.sections_total = len(plan.tasks)
                        det = f"'{plan.blog_title}'\n{len(plan.tasks)} sections · {plan.blog_kind}"
                        _set_node("orchestrator", "done", det)
                        _add_log(f"✅ Plan → '{plan.blog_title}' ({len(plan.tasks)} sections)", "success")
                        for t in plan.tasks:
                            _add_log(f"  §{t.id} {t.title} ({t.target_words}w)", "info")

                # ── Worker ───────────────────────────────────────────────────
                elif node_name == "worker":
                    sections = upd.get("sections", [])
                    if sections:
                        st.session_state.sections_done += len(sections)
                        d = st.session_state.sections_done
                        t = st.session_state.sections_total
                        _set_node("worker", "done" if d>=t else "running", f"{d}/{t} sections written")
                        for s in sections:
                            sid = s.split(":::",1)[0] if ":::" in s else "?"
                            wc = len(s.split())
                            _add_log(f"✍️ Section §{sid} written ({wc} words)", "success")

                # ── Reducer / subgraph nodes ─────────────────────────────────
                elif node_name in ("reducer","merge_content","decide_images","generate_and_place_images"):
                    if node_name in ("reducer","merge_content"):
                        _set_node("reducer","running","Merging sections…")
                        _add_log("🔧 Reducer: merging section output…", "node")
                    if node_name == "decide_images":
                        _set_node("reducer","running","Planning image placements (LLM)…")
                        _add_log("🖼️ decide_images: asking LLM for image placement plan…", "tool")
                    if node_name == "generate_and_place_images":
                        _set_node("reducer","running","Generating images with Gemini 2.5 Flash…")
                        _add_log("🎨 generate_and_place_images: calling gemini-2.5-flash-image…", "tool")

                    specs = upd.get("image_specs")
                    if specs:
                        st.session_state.image_specs = specs
                        _add_log(f"🖼️ Image plan ready: {len(specs)} images", "success")
                        for sp in specs:
                            _add_log(f"  🎨 {sp.get('placeholder','')} → {sp.get('filename','')}", "tool")

                    final = upd.get("final")
                    if final and len(final) > 80:
                        st.session_state.final_md = final
                        wc = len(final.split())
                        ic = final.count("![")
                        _set_node("reducer","done", f"Complete · {wc} words · {ic} images")
                        _add_log(f"📄 Final blog: {wc} words, {ic} inline images", "success")

                _render_all()

        # ── Complete ──────────────────────────────────────────────────────
        st.session_state.elapsed = time.time() - st.session_state.start_time
        st.session_state.done    = True
        st.session_state.running = False
        _add_log(f"🎉 Done! Blog generated in {st.session_state.elapsed:.1f}s", "success")
        _render_all()

    except Exception as ex:
        st.session_state.running = False
        st.session_state.error   = str(ex)
        tb = traceback.format_exc()
        _add_log(f"❌ Error: {ex}", "error")
        for line in tb.splitlines()[-12:]:
            _add_log(f"  {line}", "error")
        for n in st.session_state.node_states:
            if st.session_state.node_states[n]["status"] == "running":
                _set_node(n, "error", str(ex)[:120])
        _render_all()

    finally:
        _lm.FallbackLLMManager.invoke = _orig_inv
        _lm.StructuredFallbackLLM.invoke = _orig_str

# ─────────────────────────────────────────────────────────────────────────────
# Always render current state
# ─────────────────────────────────────────────────────────────────────────────
_render_all()
_render_blog()

if st.session_state.running:
    with tab_blog:
        blog_ph.info("⏳ Blog is being written… check the **Live Execution** tab to follow along!")

if st.session_state.error:
    st.error(f"❌ **Agent Error:** {st.session_state.error}")