app.py

"""
Grasping Gooning — analysis agent UI
Run: streamlit run app.py  (from /Users/binx/Desktop/Goon/)
"""

from __future__ import annotations

import inspect
import json
import os
import random
import sys
import threading
import time
from pathlib import Path

import pandas as pd
import plotly.io as pio
import streamlit as st
from dotenv import load_dotenv

# ── paths ──────────────────────────────────────────────────────────────────
ROOT = Path(__file__).parent
sys.path.insert(0, str(ROOT / "agent"))
load_dotenv(ROOT / "agent" / ".env")

from analysis import run_agent, list_datasets

# ── page config ────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Grasping Gooning",
    layout="wide",
    initial_sidebar_state="expanded",
)

@st.cache_data(show_spinner=False)
def load_post_samples(n: int = 120) -> list[dict]:
    """Random sample of real post titles for the loading slideshow."""
    try:
        import pyarrow.dataset as _ds
        _path = ROOT / "data" / "posts.parquet"
        if not _path.exists():
            return []
        d = _ds.dataset(str(_path), format="parquet")
        t = d.scanner(columns=["subreddit", "title"]).head(40000).to_pandas()
        mask = (
            t["title"].str.len() > 30
        ) & (
            t["title"].str.len() < 180
        ) & (
            ~t["title"].str.lower().str.startswith("[")
        )
        sample = t[mask].sample(min(n, mask.sum()), random_state=None)
        return sample[["subreddit", "title"]].to_dict(orient="records")
    except Exception:
        return []


LOADING_HINTS = [
    "you're so close…",
    "keep going…",
    "deeper…",
    "almost there…",
    "don't stop now…",
    "just a bit more…",
    "stay with it…",
    "right there…",
    "edge of something…",
    "hold on…",
    "so close…",
    "don't stop…",
]

# ── global CSS ─────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/icon?family=Material+Icons');

/* ---- tokens ---- */
:root {
  --bg:      #ffffff;
  --surface: #f5f5f5;
  --border:  #e0e0e0;
  --divider: #ebebeb;
  --ink:     #000000;
  --body:    #222222;
  --mid:     #555555;
  --muted:   #888888;
  --faint:   #aaaaaa;
}

/* ---- base ---- */
html, body, .stApp,
[data-testid="stAppViewContainer"],
[data-testid="stMain"],
[data-testid="stHeader"],
[data-testid="stToolbar"],
[data-testid="stBottom"],
[data-testid="stBottomBlockContainer"] {
  background: var(--bg) !important;
  color: var(--ink) !important;
  font-family: Arial, Helvetica, sans-serif !important;
}

/* header bar */
[data-testid="stHeader"] {
  border-bottom: 1px solid var(--border) !important;
  box-shadow: none !important;
}

/* bottom chat bar */
[data-testid="stBottom"] {
  border-top: 1px solid var(--border) !important;
  box-shadow: none !important;
}

#MainMenu, footer { visibility: hidden; }
/* hide deploy button (confirmed testid from Streamlit 1.50 bundle) */
[data-testid="stAppDeployButton"] {
  display: none !important;
}

/* ---- layout ---- */
.block-container {
  max-width: 1060px !important;
  padding: 56px 32px 140px !important;
}

/* ---- sidebar (dark) ---- */
section[data-testid="stSidebar"] {
  background: #0f0f0f !important;
  border-right: 1px solid #1e1e1e !important;
}
section[data-testid="stSidebar"] .block-container {
  padding: 28px 16px 48px !important;
}
/* all text inside sidebar goes light */
section[data-testid="stSidebar"] p,
section[data-testid="stSidebar"] span,
section[data-testid="stSidebar"] label,
section[data-testid="stSidebar"] div,
section[data-testid="stSidebar"] .stMarkdown {
  color: #cccccc !important;
  font-family: Arial, Helvetica, sans-serif !important;
}
/* sidebar collapse button */
[data-testid="stSidebarCollapseButton"] button::after { color: #555 !important; }
[data-testid="stExpandSidebarButton"] button::after   { color: #555 !important; }

/* ---- sidebar heading animation ---- */
.sb-title {
  font-size: 13px;
  font-weight: 700;
  letter-spacing: 0.18em;
  text-transform: uppercase;
  color: #ffffff !important;
  margin-bottom: 2px;
  animation: sbFadeDown 500ms cubic-bezier(0.22,1,0.36,1) both;
}
.sb-tagline {
  font-size: 10px;
  letter-spacing: 0.08em;
  color: #555 !important;
  overflow: hidden;
  white-space: nowrap;
  width: 0;
  animation: sbTypewriter 1.4s steps(32, end) 300ms forwards,
             sbBlinkCursor 600ms step-end 300ms 3;
  border-right: 1px solid #444;
}
@keyframes sbFadeDown {
  from { opacity: 0; transform: translateY(-6px); }
  to   { opacity: 1; transform: none; }
}
@keyframes sbTypewriter {
  to { width: 100%; border-right-color: transparent; }
}
@keyframes sbBlinkCursor {
  50% { border-right-color: transparent; }
}

/* ---- sidebar labels ---- */
.sidebar-label {
  font-size: 9px;
  font-weight: 700;
  letter-spacing: 0.16em;
  text-transform: uppercase;
  color: #444444 !important;
  animation: sbFadeDown 400ms ease both;
}
.sidebar-box {
  border-top: 1px solid #1e1e1e;
  padding-top: 14px;
  margin-top: 14px;
}
.sidebar-copy {
  font-size: 10px;
  line-height: 1.7;
  color: #777777 !important;
}
.sidebar-stat {
  font-size: 11px;
  color: #aaaaaa !important;
  font-weight: 600;
}

/* ---- sidebar buttons ---- */
section[data-testid="stSidebar"] .stButton > button {
  width: 100% !important;
  background: transparent !important;
  border: 1px solid #2a2a2a !important;
  border-radius: 0 !important;
  color: #666666 !important;
  padding: 8px 10px !important;
  text-align: left !important;
  font-size: 10px !important;
  letter-spacing: 0.08em !important;
  text-transform: uppercase !important;
  box-shadow: none !important;
  transition: background 150ms, border-color 150ms, color 150ms !important;
}
section[data-testid="stSidebar"] .stButton > button:hover {
  background: #1a1a1a !important;
  border-color: #555555 !important;
  color: #eeeeee !important;
}

/* ---- sidebar inputs ---- */
section[data-testid="stSidebar"] .stTextInput input {
  background: #1a1a1a !important;
  border: 1px solid #2a2a2a !important;
  border-radius: 0 !important;
  color: #cccccc !important;
  box-shadow: none !important;
  font-size: 12px !important;
  transition: border-color 150ms !important;
}
section[data-testid="stSidebar"] .stTextInput input:focus {
  border-color: #555555 !important;
}
section[data-testid="stSidebar"] .stTextInput input::placeholder {
  color: #444444 !important;
}

/* ---- sidebar expander (dark) ---- */
section[data-testid="stSidebar"] [data-testid="stExpander"] {
  background: transparent !important;
  border: 1px solid #1e1e1e !important;
}
section[data-testid="stSidebar"] [data-testid="stExpander"]:hover {
  border-color: #333333 !important;
}
section[data-testid="stSidebar"] [data-testid="stExpander"] summary p,
section[data-testid="stSidebar"] [data-testid="stExpander"] summary span {
  color: #888888 !important;
  font-size: 10px !important;
  letter-spacing: 0.1em !important;
  text-transform: uppercase !important;
}

/* ---- main area buttons ---- */
.block-container .stButton > button {
  width: 100% !important;
  background: transparent !important;
  border: 1px solid #cccccc !important;
  border-radius: 0 !important;
  color: var(--mid) !important;
  padding: 8px 10px !important;
  text-align: left !important;
  font-size: 10px !important;
  letter-spacing: 0.08em !important;
  text-transform: uppercase !important;
  box-shadow: none !important;
  transition: background 150ms, border-color 150ms, color 150ms !important;
}
.block-container .stButton > button:hover {
  background: var(--surface) !important;
  border-color: var(--ink) !important;
  color: var(--ink) !important;
}

/* ---- inputs ---- */
.stTextInput input {
  background: var(--bg) !important;
  border: 1px solid var(--border) !important;
  border-radius: 0 !important;
  color: var(--ink) !important;
  box-shadow: none !important;
  font-size: 12px !important;
  transition: border-color 150ms !important;
}
.stTextInput input:focus { border-color: var(--ink) !important; }

/* chat input container */
[data-testid="stChatInput"],
[data-testid="stChatInput"] > div,
[data-testid="stChatInputContainer"] {
  background: #ffffff !important;
  border: 1px solid #d0d0d0 !important;
  border-radius: 0 !important;
  box-shadow: none !important;
  transition: border-color 150ms !important;
}
[data-testid="stChatInput"]:focus-within,
[data-testid="stChatInputContainer"]:focus-within {
  border-color: #000000 !important;
  box-shadow: none !important;
}
[data-testid="stChatInput"] textarea {
  background: #ffffff !important;
  color: #000000 !important;
  font-size: 14px !important;
  font-family: Arial, Helvetica, sans-serif !important;
}
/* send button */
[data-testid="stChatInput"] button,
[data-testid="stChatInputContainer"] button {
  background: #000000 !important;
  border: none !important;
  border-radius: 0 !important;
  color: #ffffff !important;
  box-shadow: none !important;
  min-width: 44px !important;
  width: 44px !important;
  height: 100% !important;
  font-family: Arial, Helvetica, sans-serif !important;
  font-size: 13px !important;
  font-weight: 700 !important;
  letter-spacing: 0.04em !important;
}
[data-testid="stChatInput"] button:hover,
[data-testid="stChatInputContainer"] button:hover {
  background: #333333 !important;
  opacity: 1 !important;
}
/* replace SVG arrow with text "->" */
[data-testid="stChatInput"] button svg,
[data-testid="stChatInputContainer"] button svg { display: none !important; }
[data-testid="stChatInput"] button::after,
[data-testid="stChatInputContainer"] button::after {
  content: "->";
  font-family: Arial, Helvetica, sans-serif !important;
  font-size: 13px;
  font-weight: 700;
  color: #ffffff;
}
/* kill any rounded wrapper Streamlit adds around the whole bar */
[data-testid="stBottom"] > div,
[data-testid="stBottomBlockContainer"] > div {
  background: #ffffff !important;
  border-radius: 0 !important;
  box-shadow: none !important;
}

/* ---- expander ---- */
[data-testid="stExpander"] {
  background: transparent !important;
  border: 1px solid var(--border) !important;
  border-radius: 0 !important;
  overflow: hidden !important;
  transition: border-color 150ms !important;
}
[data-testid="stExpander"]:hover { border-color: var(--ink) !important; }
[data-testid="stExpander"] summary { padding: 10px 14px !important; }
/* kill the expander toggle icon (data-testid confirmed from Streamlit 1.50 bundle) */
[data-testid="stExpander"] summary [data-testid="stIconMaterial"],
[data-testid="stExpander"] summary [data-testid="stImageIcon"] {
  display: none !important;
}

/* sidebar collapse / expand toggle icons -> replace with < > */
[data-testid="stExpandSidebarButton"] [data-testid="stIconMaterial"],
[data-testid="stSidebarCollapseButton"] [data-testid="stIconMaterial"] {
  display: none !important;
}
[data-testid="stExpandSidebarButton"] button::after {
  content: ">";
  font-size: 15px; font-weight: 700;
  font-family: Arial, Helvetica, sans-serif !important;
  color: #555555;
}
[data-testid="stSidebarCollapseButton"] button::after {
  content: "<";
  font-size: 15px; font-weight: 700;
  font-family: Arial, Helvetica, sans-serif !important;
  color: #555555;
}

/* ---- progress bar ---- */
.prog-wrap { padding: 20px 0 12px; }
.prog-hint {
  font-size: 11px; color: var(--muted);
  letter-spacing: 0.1em; margin-bottom: 10px;
  font-style: italic;
  animation: progPulse 1.8s ease-in-out infinite;
}
@keyframes progPulse {
  0%, 100% { opacity: 0.5; }
  50%       { opacity: 1; }
}
.prog-bg {
  background: var(--divider); height: 1px; width: 100%; margin-bottom: 6px;
  position: relative; overflow: hidden;
}
.prog-fill {
  background: var(--ink); height: 1px;
  transition: width 0.35s ease;
  position: absolute; top: 0; left: 0;
}
/* shimmer on the fill bar */
.prog-fill::after {
  content: "";
  position: absolute; top: 0; right: 0;
  width: 40px; height: 1px;
  background: linear-gradient(to right, transparent, #fff, transparent);
  animation: shimmer 1.2s ease-in-out infinite;
}
@keyframes shimmer {
  0%   { opacity: 0; transform: translateX(-40px); }
  50%  { opacity: 1; }
  100% { opacity: 0; transform: translateX(40px); }
}
.prog-pct {
  font-size: 9px; color: var(--faint);
  letter-spacing: 0.14em; text-transform: uppercase;
  font-family: "Courier New", monospace !important;
}
.prog-stuck {
  margin-top: 8px;
  font-size: 10px; color: var(--muted);
  letter-spacing: 0.08em; font-style: italic;
  animation: stuckFadeIn 400ms ease both;
}
.prog-stuck-0 { color: var(--muted); }
.prog-stuck-1 { color: var(--faint); }
.prog-stuck-2 { color: #cccccc; font-size: 9px; }
.prog-stuck-3 { color: #dddddd; font-size: 9px; }
.prog-stuck-4 { color: #e0e0e0; font-size: 9px; }
.prog-stuck-5 { color: #e8e8e8; font-size: 9px; }
@keyframes stuckFadeIn {
  from { opacity: 0; transform: translateY(4px); }
  to   { opacity: 1; transform: none; }
}

/* ---- loading post slideshow ---- */
.post-slide {
  margin-top: 20px;
  padding: 14px 18px;
  background: var(--surface);
  border-left: 2px solid var(--divider);
  animation: slideIn 300ms cubic-bezier(0.22,1,0.36,1) both;
}
.post-slide-sub {
  font-size: 9px; letter-spacing: 0.14em; text-transform: uppercase;
  color: var(--faint); margin-bottom: 6px;
  font-family: "Courier New", monospace !important;
}
.post-slide-title {
  font-size: 13px; line-height: 1.5; color: var(--body);
  font-style: italic;
}
@keyframes slideIn {
  from { opacity: 0; transform: translateY(6px); }
  to   { opacity: 1; transform: none; }
}

/* ---- chat ---- */
[data-testid="stChatMessage"] {
  background: transparent !important;
  border: none !important;
  padding: 0 !important;
  margin: 0 0 28px !important;
  gap: 12px !important;
  animation: fadeUp 240ms cubic-bezier(0.22,1,0.36,1) both;
}
/* user avatar: kaomoji */
[data-testid="stChatMessageAvatarUser"] {
  background: #000000 !important;
  border: none !important;
  width: 34px !important; height: 34px !important;
  min-width: 34px !important;
  border-radius: 0 !important;
  position: relative !important;
  overflow: visible !important;
}
[data-testid="stChatMessageAvatarUser"] svg,
[data-testid="stChatMessageAvatarUser"] img { display: none !important; }
[data-testid="stChatMessageAvatarUser"]::after {
  content: "( ˘▾˘)";
  position: absolute; top: 50%; left: 50%;
  transform: translate(-50%, -50%);
  font-size: 11px; line-height: 1; color: #ffffff;
  white-space: nowrap;
  font-family: Arial, Helvetica, sans-serif !important;
}
/* assistant avatar: kaomoji */
[data-testid="stChatMessageAvatarAssistant"] {
  background: #ffffff !important;
  border: 1px solid var(--border) !important;
  width: 34px !important; height: 34px !important;
  min-width: 34px !important;
  border-radius: 0 !important;
  position: relative !important;
  overflow: visible !important;
}
[data-testid="stChatMessageAvatarAssistant"] svg,
[data-testid="stChatMessageAvatarAssistant"] img { display: none !important; }
[data-testid="stChatMessageAvatarAssistant"]::after {
  content: "ʕ•ᴥ•ʔ";
  position: absolute; top: 50%; left: 50%;
  transform: translate(-50%, -50%);
  font-size: 11px; line-height: 1; color: #000000;
  white-space: nowrap;
  font-family: Arial, Helvetica, sans-serif !important;
}

.msg-meta {
  display: flex; align-items: center; gap: 12px;
  margin-bottom: 8px;
}
.msg-label {
  font-size: 10px; font-weight: 700;
  letter-spacing: 0.14em; text-transform: uppercase;
  color: var(--muted);
}
.route-tag {
  font-size: 10px; letter-spacing: 0.1em; text-transform: uppercase;
  color: var(--faint); border-left: 1px solid var(--border); padding-left: 10px;
}
.msg-body {
  border-top: 1px solid var(--divider);
  padding-top: 12px;
  animation: fadeIn 220ms 60ms ease both;
}
.msg-body p {
  font-size: 14px !important; line-height: 1.7 !important;
  color: var(--body) !important; max-width: 72ch !important;
}

/* ---- cost bar ---- */
.cost-row {
  display: flex; align-items: center; gap: 16px;
  margin-top: 22px; margin-bottom: 18px;
  padding: 14px 18px;
  background: #000000;
}
.cost-label {
  font-size: 10px; letter-spacing: 0.18em; text-transform: uppercase;
  color: #666666; white-space: nowrap; font-family: "Courier New", monospace !important;
}
.cost-track {
  flex: 1; height: 2px; background: #2a2a2a; position: relative;
}
.cost-fill {
  position: absolute; top: 0; left: 0; height: 2px;
  background: #ffffff;
  transition: width 600ms cubic-bezier(0.22,1,0.36,1);
}
.cost-val {
  font-size: 12px; letter-spacing: 0.06em; color: #ffffff;
  font-family: "Courier New", monospace !important; white-space: nowrap;
}
.cost-tok {
  font-size: 10px; letter-spacing: 0.04em; color: #666666;
  font-family: "Courier New", monospace !important; white-space: nowrap;
}

/* ---- step trace ---- */
.step-title {
  display: block; font-size: 12px; font-weight: 700;
  margin-bottom: 4px; color: var(--ink);
}
.spath {
  display: block; margin-top: 4px;
  font-size: 11px; color: var(--muted);
  font-family: "Courier New", monospace !important;
}
[data-testid="stDataFrame"] table {
  font-size: 12px !important;
  font-family: Arial, Helvetica, sans-serif !important;
}

/* ---- keyframes ---- */
@keyframes fadeUp {
  from { opacity:0; transform:translateY(6px); }
  to   { opacity:1; transform:none; }
}
@keyframes fadeIn {
  from { opacity:0; } to { opacity:1; }
}

/* ---- responsive ---- */
@media (max-width:640px) {
  .block-container { padding: 32px 16px 100px !important; }
}
</style>
""", unsafe_allow_html=True)


# ── helpers ────────────────────────────────────────────────────────────────
def fmt(v: int | None) -> str:
    return "n/a" if v is None else f"{v:,}"


def dataset_snapshot() -> dict:
    try:
        return list_datasets()
    except Exception:
        return {}


def render_plot(plotly_json: str) -> None:
    try:
        fig = pio.from_json(plotly_json)
        fig.update_layout(
            paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)",
            font=dict(family="Arial, Helvetica, sans-serif", color="#222", size=12),
            margin=dict(l=0, r=0, t=32, b=0),
            colorway=["#000", "#555", "#888", "#bbb"],
            xaxis=dict(gridcolor="#ebebeb", linecolor="#e0e0e0"),
            yaxis=dict(gridcolor="#ebebeb", linecolor="#e0e0e0"),
        )
        st.plotly_chart(fig, use_container_width=True)
    except Exception as exc:
        st.warning(f"Chart could not be rendered: {exc}")


def compact_tool_result(result: object) -> dict:
    if not isinstance(result, dict):
        return {"value": result}
    compact: dict = {"keys": sorted(result.keys())}
    for key in ("saved_csv", "saved_png", "plotly_json", "error", "analysis", "dataset", "filters"):
        if key in result and result.get(key) is not None:
            compact[key] = result[key]
    table = result.get("table")
    if isinstance(table, list):
        compact["table_rows"] = len(table)
        compact["table_preview"] = table[:5]
    return compact


def extract_artifacts(tool_calls: list[dict]) -> list[dict]:
    artifacts: list[dict] = []
    for tc in tool_calls:
        result = tc.get("result") or {}
        if not isinstance(result, dict):
            continue
        for key, atype in (("saved_csv", "csv"), ("saved_png", "png")):
            if result.get(key):
                artifacts.append({"type": atype, "tool": tc.get("tool", "?"), "path": result[key]})
        if result.get("plotly_json"):
            artifacts.append({"type": "plotly_json", "tool": tc.get("tool", "?"), "present": True})
    return artifacts


def build_backend_history(turns: list[dict]) -> list[dict]:
    history: list[dict] = []
    for turn in turns:
        history.append({"role": "user", "content": turn["question"]})
        content = turn["answer"]
        state = {
            "tool_calls": [
                {"tool": tc.get("tool"), "args": tc.get("args") or {},
                 "result": compact_tool_result(tc.get("result"))}
                for tc in turn.get("tool_calls", [])
            ],
            "artifacts": turn.get("artifacts", []),
            "plotly_json": bool(turn.get("plotly_json")),
            "route": turn.get("route"),
        }
        if state["tool_calls"] or state["artifacts"] or state["plotly_json"]:
            content += f"\n\n<analysis_state>\n{json.dumps(state, default=str, indent=2)}\n</analysis_state>"
        history.append({"role": "assistant", "content": content})
    return history


def call_agent(question: str, history: list[dict], turns: list[dict]) -> dict:
    kwargs = {"history": history}
    params = inspect.signature(run_agent).parameters
    for name in ("analysis_context", "conversation_state", "turns"):
        if name in params:
            kwargs[name] = turns
            break
    return run_agent(question, **kwargs)


_POST_SAMPLES: list[dict] = []


def call_agent_with_progress(question: str, backend_history: list[dict], turns: list[dict], slot) -> dict:
    """Run agent in a background thread; update a progress slot from the main thread."""
    result_holder: dict = {}
    exc_holder: dict = {}

    def worker() -> None:
        try:
            result_holder["r"] = call_agent(question, backend_history, turns)
        except Exception as e:
            exc_holder["e"] = e

    t = threading.Thread(target=worker, daemon=True)
    t.start()

    STUCK_MSGS = [
        "i promise i'm still gooning",
        "locked in. fully gooned. cannot stop.",
        "the data is vast. the goon is deep. patience.",
        "i've been edging this query for so long i've lost track of time.",
        "every second is another row scanned. feel it.",
        "this is what a true goon session looks like. no shortcuts.",
        "i am one with the dataset. do not disturb.",
    ]

    global _POST_SAMPLES
    if not _POST_SAMPLES:
        _POST_SAMPLES = load_post_samples()
    posts = _POST_SAMPLES if _POST_SAMPLES else []
    random.shuffle(posts)

    pct = 0
    idx = 0
    post_idx = 0
    start = time.time()
    stuck_since: float | None = None
    last_pct_change = time.time()

    while t.is_alive():
        prev_pct = pct
        pct = min(pct + random.randint(1, 5), 93)

        if pct != prev_pct:
            stuck_since = None
            last_pct_change = time.time()
        else:
            if stuck_since is None:
                stuck_since = time.time()

        elapsed = int(time.time() - start)
        elapsed_str = f"{elapsed}s" if elapsed < 60 else f"{elapsed // 60}m {elapsed % 60}s"
        stuck_sec = int(time.time() - stuck_since) if stuck_since else 0

        hint = LOADING_HINTS[idx % len(LOADING_HINTS)]

        # Build up stuck messages — one new line per 12s window, cleared when pct moves
        n_stuck = min(stuck_sec // 12, len(STUCK_MSGS))
        stuck_html = "".join(
            f'<div class="prog-stuck prog-stuck-{i}">{STUCK_MSGS[i]}</div>'
            for i in range(n_stuck)
        )

        # rotate post every ~11 ticks (~4 seconds)
        if idx % 11 == 0 and idx > 0:
            post_idx += 1
        post_html = ""
        if posts:
            p = posts[post_idx % len(posts)]
            sub   = p.get("subreddit", "")
            title = p.get("title", "").replace("<", "&lt;").replace(">", "&gt;")
            post_html = (
                f'<div class="post-slide" key="{post_idx}">'
                f'<div class="post-slide-sub">r/{sub}</div>'
                f'<div class="post-slide-title">{title}</div>'
                f'</div>'
            )

        at_cap = pct >= 93
        pct_display = "—%" if at_cap else f"{pct}%"
        running_label = (
            '<span class="prog-hint" style="display:inline;margin-left:8px;margin-bottom:0">'
            'still running</span>'
            if at_cap else ""
        )
        slot.markdown(
            f'<div class="prog-wrap">'
            f'<div class="prog-hint">{hint}</div>'
            f'<div class="prog-bg"><div class="prog-fill" style="width:{pct}%"></div></div>'
            f'<div class="prog-pct">{pct_display} &nbsp;·&nbsp; {elapsed_str}{running_label}</div>'
            f'{stuck_html}'
            f'{post_html}'
            f'</div>',
            unsafe_allow_html=True,
        )
        idx += 1
        time.sleep(0.35)

    t.join()
    slot.empty()

    if exc_holder:
        raise exc_holder["e"]
    return result_holder["r"]


def render_cost_bar(usage: dict) -> None:
    cost  = usage.get("cost_usd", 0)
    inp   = usage.get("input_tokens", 0)
    out   = usage.get("output_tokens", 0)
    # scale: 0–$0.50 maps to 0–100% of bar
    pct   = min(cost / 0.50 * 100, 100)
    if cost < 0.01:
        val_str = f"< $0.01"
    else:
        val_str = f"${cost:.3f}"
    tok_str = f"{inp:,} in · {out:,} out"
    st.markdown(
        f'<div class="cost-row">'
        f'<span class="cost-label">cost</span>'
        f'<div class="cost-track"><div class="cost-fill" style="width:{pct:.1f}%"></div></div>'
        f'<span class="cost-val">{val_str}</span>'
        f'<span class="cost-tok">{tok_str}</span>'
        f'</div>',
        unsafe_allow_html=True,
    )


def render_tool_calls(tool_calls: list[dict]) -> None:
    n = len(tool_calls)
    with st.expander(f"Method  {n} step{'s' if n != 1 else ''}", expanded=False):
        for i, tc in enumerate(tool_calls):
            st.markdown(
                f"<span class='step-title'>Step {i+1} -> {tc.get('tool','?')}</span>",
                unsafe_allow_html=True,
            )
            if tc.get("args"):
                st.json(tc["args"], expanded=False)
            res = tc.get("result") or {}
            if isinstance(res, dict):
                if res.get("table"):
                    try:
                        st.dataframe(pd.DataFrame(res["table"]), use_container_width=True, hide_index=True)
                    except Exception:
                        pass
                for key in ("saved_csv", "saved_png"):
                    if res.get(key):
                        st.markdown(f"<span class='spath'>-> {res[key]}</span>", unsafe_allow_html=True)
            if i < n - 1:
                st.markdown("---")


def render_export_buttons(answer: str, tool_calls: list[dict], turn_idx: int) -> None:
    artifacts = extract_artifacts(tool_calls)
    csvs = [a["path"] for a in artifacts if a["type"] == "csv"]
    pngs = [a["path"] for a in artifacts if a["type"] == "png"]

    items: list[tuple[str, bytes, str, str]] = []
    items.append(("answer.md", answer.encode("utf-8"), "text/markdown", f"answer_{turn_idx}.md"))
    for path in csvs:
        p = Path(path)
        if p.exists():
            items.append((p.name, p.read_bytes(), "text/csv", p.name))
    for path in pngs:
        p = Path(path)
        if p.exists():
            items.append((p.name, p.read_bytes(), "image/png", p.name))

    cols = st.columns(len(items))
    for col, (label, data, mime, fname) in zip(cols, items):
        with col:
            st.download_button(
                label=label, data=data, file_name=fname, mime=mime,
                key=f"dl_{turn_idx}_{fname}",
            )


# ── session state ──────────────────────────────────────────────────────────
for key, default in [("history", []), ("chat", []), ("turns", []), ("prefill", ""), ("authenticated", False), ("logged_out", False)]:
    if key not in st.session_state:
        st.session_state[key] = default

# seed from env if already set (e.g. from .env file) — but not if user explicitly logged out
if not st.session_state["authenticated"] and not st.session_state["logged_out"] and os.environ.get("ANTHROPIC_API_KEY"):
    st.session_state["authenticated"] = True

# ── dataset metadata ───────────────────────────────────────────────────────
meta          = dataset_snapshot()
posts_rows    = meta.get("posts", {}).get("rows")
comments_rows = meta.get("comments", {}).get("rows")
sub_count     = len(meta.get("posts", {}).get("subreddits") or [])
latest_date   = (meta.get("comments", {}).get("date_range") or {}).get("latest", "n/a")


# ── login gate ─────────────────────────────────────────────────────────────
if not st.session_state["authenticated"]:
    st.markdown("""
<style>
.login-wrap {
  display: flex; flex-direction: column; align-items: center;
  justify-content: center; min-height: 70vh; gap: 20px;
}
.login-title {
  font-size: 22px; font-weight: 700; letter-spacing: 0.04em; color: var(--ink);
}
.login-sub {
  font-size: 12px; color: var(--muted); margin-top: -12px;
}
</style>
<div class='login-wrap'>
  <div class='login-title'>Grasping Gooning</div>
  <div class='login-sub'>enter your Anthropic API key to continue</div>
</div>
""", unsafe_allow_html=True)

    col = st.columns([1, 2, 1])[1]
    with col:
        login_key = st.text_input(
            "API key", type="password", placeholder="sk-ant-…",
            label_visibility="collapsed",
        )
        if st.button("Enter ->", key="login_btn", use_container_width=True):
            if login_key.strip():
                ascii_key = login_key.encode("ascii", errors="ignore").decode("ascii")
                os.environ["ANTHROPIC_API_KEY"] = ascii_key
                st.session_state["authenticated"] = True
                st.session_state["logged_out"] = False
                st.rerun()
            else:
                st.error("Paste your API key above.")
    st.stop()

# ── sidebar ────────────────────────────────────────────────────────────────
with st.sidebar:
    st.markdown("""
<div class='sb-title'>Grasping Gooning</div>
<div class='sb-tagline'>reddit data analysis agent</div>
""", unsafe_allow_html=True)

    st.markdown("<div class='sidebar-box'><div class='sidebar-label'>Session</div></div>",
                unsafe_allow_html=True)
    if st.button("Clear conversation", key="clear"):
        st.session_state.update(history=[], chat=[], turns=[], prefill="")
        st.rerun()
    if st.button("Log out", key="logout"):
        os.environ.pop("ANTHROPIC_API_KEY", None)
        st.session_state.update(history=[], chat=[], turns=[], prefill="", authenticated=False, logged_out=True)
        st.rerun()

    # ── about (bottom of sidebar) ──────────────────────────────────────────
    st.markdown("<div class='sidebar-box'>", unsafe_allow_html=True)
    with st.expander("About"):
        earliest = (meta.get("posts", {}).get("date_range") or {}).get("earliest", "n/a")
        subs_list = meta.get("posts", {}).get("subreddits") or []
        st.markdown(f"""
<div class='sidebar-copy'>
A research tool for analysing the Reddit gooning corpus.<br>
Ask questions in plain English — the agent runs real code against the data and returns charts, tables, and findings.<br><br>
<span class='sidebar-stat'>{fmt(posts_rows)}</span> posts<br>
<span class='sidebar-stat'>{fmt(comments_rows)}</span> comments<br>
<span class='sidebar-stat'>{sub_count}</span> subreddits<br>
<span style='color:#444;font-size:9px'>{earliest} — {latest_date}</span><br><br>
<span style='color:#333;font-size:9px;letter-spacing:0.1em;text-transform:uppercase'>Subreddits</span><br>
<span style='color:#555;font-size:9px;line-height:1.9'>{" · ".join(subs_list[:15])}{"..." if len(subs_list) > 15 else ""}</span>
</div>
""", unsafe_allow_html=True)
    st.markdown("</div>", unsafe_allow_html=True)


# ── chat history ───────────────────────────────────────────────────────────
for i, msg in enumerate(st.session_state["chat"]):
    with st.chat_message(msg["role"]):
        role  = msg["role"]
        route = msg.get("route", "")
        label = "You" if role == "user" else "Answer"
        route_html = f"<span class='route-tag'>{route}</span>" if route and role == "assistant" else ""
        st.markdown(
            f"<div class='msg-meta'><span class='msg-label'>{label}</span>{route_html}</div>",
            unsafe_allow_html=True,
        )
        st.markdown("<div class='msg-body'>", unsafe_allow_html=True)
        st.markdown(msg["content"])
        st.markdown("</div>", unsafe_allow_html=True)
        for pj in (msg.get("plotly_jsons") or ([msg["plotly_json"]] if msg.get("plotly_json") else [])):
            render_plot(pj)
        if msg.get("usage"):
            render_cost_bar(msg["usage"])
        if msg.get("tool_calls"):
            render_tool_calls(msg["tool_calls"])
        if role == "assistant":
            render_export_buttons(msg["content"], msg.get("tool_calls") or [], i)


# ── chat input ─────────────────────────────────────────────────────────────
prefill  = st.session_state["prefill"]
question = st.chat_input("what do you want to know…")
if prefill:
    st.session_state["prefill"] = ""
effective_question = question or prefill

if effective_question:
    question = effective_question

    backend_history = build_backend_history(st.session_state["turns"])

    with st.chat_message("user"):
        st.markdown("<div class='msg-meta'><span class='msg-label'>You</span></div>",
                    unsafe_allow_html=True)
        st.markdown("<div class='msg-body'>", unsafe_allow_html=True)
        st.markdown(question)
        st.markdown("</div>", unsafe_allow_html=True)

    with st.chat_message("assistant"):
        progress_slot = st.empty()
        try:
            result = call_agent_with_progress(question, backend_history, list(st.session_state["turns"]), progress_slot)
        except Exception as exc:
            err_str = str(exc)
            is_auth_err = (
                type(exc).__name__ in ("AuthenticationError", "PermissionDeniedError")
                or "invalid x-api-key" in err_str.lower()
                or "401" in err_str
            )
            if is_auth_err:
                os.environ.pop("ANTHROPIC_API_KEY", None)
                st.session_state.update(authenticated=False, logged_out=True)
                st.error("API key rejected — please re-enter it.")
                st.rerun()
            elif "rate_limit" in err_str.lower():
                st.error("Rate limited. Wait a moment and try again.")
            elif "Unicode encoding error" in err_str or ("ascii" in err_str.lower() and "codec" in err_str.lower()):
                st.error("Encoding error — your API key may contain non-standard characters. Log out and re-enter it.")
            else:
                st.error(f"Something went wrong: {err_str[:300]}")
            st.stop()

        answer       = result.get("answer", "")
        tool_calls   = result.get("tool_calls", [])
        plotly_jsons = result.get("plotly_jsons") or ([result["plotly_json"]] if result.get("plotly_json") else [])
        route        = result.get("route", "")
        usage        = result.get("usage") or {}

        route_html = f"<span class='route-tag'>{route}</span>" if route else ""
        st.markdown(
            f"<div class='msg-meta'><span class='msg-label'>Answer</span>{route_html}</div>",
            unsafe_allow_html=True,
        )
        st.markdown("<div class='msg-body'>", unsafe_allow_html=True)
        st.markdown(answer)
        st.markdown("</div>", unsafe_allow_html=True)
        for pj in plotly_jsons:
            render_plot(pj)
        if usage:
            render_cost_bar(usage)
        if tool_calls:
            render_tool_calls(tool_calls)
        render_export_buttons(answer, tool_calls, len(st.session_state["turns"]))

    turn = {
        "question": question, "answer": answer,
        "tool_calls": tool_calls, "plotly_jsons": plotly_jsons,
        "artifacts": extract_artifacts(tool_calls), "route": route,
        "usage": usage,
    }
    st.session_state["turns"].append(turn)
    st.session_state["history"] = build_backend_history(st.session_state["turns"])
    st.session_state["chat"].append({"role": "user", "content": question})
    st.session_state["chat"].append({
        "role": "assistant", "content": answer,
        "tool_calls": tool_calls, "plotly_jsons": plotly_jsons,
        "route": route, "usage": usage,
    })