"""
BERTopic Topic Modelling Agent — Gradio Interface
"""
import os
os.environ.pop("GRADIO_HOT_RELOAD", None)
import gradio as gr

# ── Patch for Starlette 1.0.0 TemplateResponse signature change ──────────────
import starlette.templating

_original_template_response = starlette.templating.Jinja2Templates.TemplateResponse

def _patched_template_response(self, *args, **kwargs):
    # If Gradio calls it the old positional way: TemplateResponse("name.html", {"request": ...})
    if len(args) >= 2 and isinstance(args[0], str) and isinstance(args[1], dict):
        name, context = args[0], args[1]
        request = context.get("request")
        # Route it to the new signature explicitly
        return _original_template_response(self, request=request, name=name, context=context, **kwargs)
    
    # Otherwise, pass through normally
    return _original_template_response(self, *args, **kwargs)

starlette.templating.Jinja2Templates.TemplateResponse = _patched_template_response


# ── Patch the gradio_client JSON-schema parser crash safely ──────────────────
import gradio_client.utils as gc_utils
import copy

_original_parser = gc_utils._json_schema_to_python_type

# Updated to accept *args and **kwargs to handle newer Gradio version signatures safely
def _patched_parser(schema, *args, **kwargs):
    try:
        # Deepcopy to avoid mutating the original schema object
        safe_schema = copy.deepcopy(schema)
        
        # Recursively hunt down and replace boolean additionalProperties
        def fix_booleans(obj):
            if isinstance(obj, dict):
                if "additionalProperties" in obj and isinstance(obj["additionalProperties"], bool):
                    obj["additionalProperties"] = {}
                for v in obj.values():
                    fix_booleans(v)
            elif isinstance(obj, list):
                for item in obj:
                    fix_booleans(item)
                    
        fix_booleans(safe_schema)
        # Pass the safe schema along with any other arguments (like `defs`)
        return _original_parser(safe_schema, *args, **kwargs)
    except Exception:
        # Fallback to the original schema if our scrubber fails
        return _original_parser(schema, *args, **kwargs)

gc_utils._json_schema_to_python_type = _patched_parser


import pandas as pd
import json
from pathlib import Path

# ── Try to import the real agent; fall back to a stub so the UI still loads ──
try:
    from agent import BERTopicAgent
except ImportError:
    class BERTopicAgent:
        """Stub agent used when agent.py is not present."""
        def __init__(self):
            self.phase = "idle"
            self._charts: dict[str, str] = {}
            self._downloads: list[str] = []

        def handle_message(self, message: str, history: list, csv_path: str | None = None):
            yield (
                history + [[message, "⚠️ `agent.py` not found — running in UI-preview mode."]],
                self.phase,
                self._charts,
                self._downloads,
                [],
            )

        def handle_review(self, table_data: list, history: list):
            yield (
                history + [["[Review submitted]", "Review received. No real agent attached."]],
                self.phase,
                self._charts,
                self._downloads,
            )


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────

PHASES = ["idle", "loading", "embedding", "clustering", "labelling", "review", "done"]
PHASE_LABELS = {
    "idle":       "Idle",
    "loading":    "Loading data",
    "embedding":  "Embedding documents",
    "clustering": "Clustering topics",
    "labelling":  "Labelling topics",
    "review":     "Awaiting review",
    "done":       "Complete",
}

TABLE_COLS = ["#", "Topic Label", "Top Evidence", "Sentences", "Papers", "Approve", "Rename To", "Reasoning"]

EMPTY_TABLE = pd.DataFrame(columns=TABLE_COLS)


def phase_html(phase: str) -> str:
    """Render a step-progress bar for the current pipeline phase."""
    steps = [p for p in PHASES if p != "idle"]
    try:
        current_idx = steps.index(phase) if phase in steps else -1
    except ValueError:
        current_idx = -1

    dots = []
    for i, step in enumerate(steps):
        if i < current_idx:
            cls, icon = "step done", "✓"
        elif i == current_idx:
            cls, icon = "step active", "●"
        else:
            cls, icon = "step", "○"
        dots.append(
            f'<div class="{cls}"><span class="icon">{icon}</span>'
            f'<span class="label">{PHASE_LABELS[step]}</span></div>'
        )

    bar_pct = 0 if current_idx < 0 else int((current_idx + 1) / len(steps) * 100)

    return f"""
<div class="phase-wrap">
  <div class="phase-steps">{''.join(dots)}</div>
  <div class="phase-bar-bg">
    <div class="phase-bar-fill" style="width:{bar_pct}%"></div>
  </div>
</div>
"""


def make_chart_html(iframe_src: str) -> str:
    if not iframe_src:
        return "<p style='color:#888;padding:1rem'>No chart available yet.</p>"
    return (
        f'<iframe src="{iframe_src}" width="100%" height="520" '
        f'frameborder="0" scrolling="no" style="border-radius:8px;"></iframe>'
    )


# ──────────────────────────────────────────────────────────────────────────────
# CSS
# ──────────────────────────────────────────────────────────────────────────────

CUSTOM_CSS = """
/* ── Google Fonts ── */
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700;800&family=JetBrains+Mono:wght@400;500&family=Lato:wght@300;400;700&display=swap');

/* ── Root palette ── */
:root {
  --bg:        #0d0f14;
  --surface:   #151820;
  --border:    #252a38;
  --accent:    #4f8ef7;
  --accent2:   #a259ff;
  --success:   #2dce89;
  --warn:      #f7c94f;
  --danger:    #f76f6f;
  --text:      #e8ecf4;
  --muted:     #7a8299;
  --radius:    10px;
  --mono:      'JetBrains Mono', monospace;
  --sans:      'Lato', sans-serif;
  --display:   'Syne', sans-serif;
}

/* ── Global resets ── */
body, .gradio-container {
  background: var(--bg) !important;
  color: var(--text) !important;
  font-family: var(--sans) !important;
}

/* ── App header ── */
.app-header {
  padding: 2rem 0 1.2rem;
  border-bottom: 1px solid var(--border);
  margin-bottom: 1.5rem;
}
.app-header h1 {
  font-family: var(--display);
  font-size: 2rem;
  font-weight: 800;
  letter-spacing: -0.03em;
  background: linear-gradient(135deg, var(--accent) 0%, var(--accent2) 100%);
  -webkit-background-clip: text;
  -webkit-text-fill-color: transparent;
  margin: 0 0 .25rem;
}
.app-header p {
  color: var(--muted);
  font-size: .9rem;
  margin: 0;
}

/* ── Section cards ── */
.section-card {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 1.25rem 1.5rem;
  margin-bottom: 1.2rem;
}
.section-title {
  font-family: var(--display);
  font-size: .75rem;
  font-weight: 700;
  letter-spacing: .12em;
  text-transform: uppercase;
  color: var(--accent);
  margin: 0 0 1rem;
}

/* ── Phase progress bar ── */
.phase-wrap {
  padding: .5rem 0 .8rem;
}
.phase-steps {
  display: flex;
  gap: 0;
  margin-bottom: .6rem;
  flex-wrap: wrap;
}
.step {
  display: flex;
  flex-direction: column;
  align-items: center;
  flex: 1;
  min-width: 70px;
  position: relative;
  opacity: .38;
}
.step::after {
  content: '';
  position: absolute;
  top: 10px;
  left: 50%;
  width: 100%;
  height: 2px;
  background: var(--border);
  z-index: 0;
}
.step:last-child::after { display: none; }
.step .icon {
  font-size: .85rem;
  z-index: 1;
  background: var(--surface);
  padding: 0 4px;
  color: var(--muted);
}
.step .label {
  font-size: .62rem;
  font-family: var(--mono);
  color: var(--muted);
  margin-top: 4px;
  text-align: center;
}
.step.done { opacity: .7; }
.step.done .icon { color: var(--success); }
.step.active { opacity: 1; }
.step.active .icon {
  color: var(--accent);
  animation: pulse 1.2s ease-in-out infinite;
}
.step.active .label { color: var(--accent); font-weight: 600; }
@keyframes pulse { 0%,100%{opacity:1} 50%{opacity:.4} }

.phase-bar-bg {
  height: 4px;
  background: var(--border);
  border-radius: 4px;
  overflow: hidden;
}
.phase-bar-fill {
  height: 100%;
  background: linear-gradient(90deg, var(--accent), var(--accent2));
  border-radius: 4px;
  transition: width .5s ease;
}

/* ── Chatbot ── */
.chatbot-wrap .message.user { background: #1a2035 !important; }
.chatbot-wrap .message.bot  { background: #151c2e !important; }

/* ── Buttons ── */
.btn-primary {
  background: linear-gradient(135deg, var(--accent), var(--accent2)) !important;
  color: #fff !important;
  border: none !important;
  font-family: var(--display) !important;
  font-weight: 700 !important;
  letter-spacing: .04em !important;
  border-radius: var(--radius) !important;
  padding: .55rem 1.4rem !important;
  transition: opacity .2s, transform .1s !important;
}
.btn-primary:hover { opacity: .88 !important; transform: translateY(-1px) !important; }
.btn-primary:active { transform: translateY(0) !important; }

.btn-secondary {
  background: var(--surface) !important;
  border: 1px solid var(--border) !important;
  color: var(--text) !important;
  font-family: var(--display) !important;
  font-weight: 600 !important;
  border-radius: var(--radius) !important;
}
.btn-secondary:hover { border-color: var(--accent) !important; color: var(--accent) !important; }

/* ── Tabs ── */
.tab-nav button {
  font-family: var(--display) !important;
  font-weight: 600 !important;
  font-size: .82rem !important;
  letter-spacing: .06em !important;
  text-transform: uppercase !important;
  color: var(--muted) !important;
  border-bottom: 2px solid transparent !important;
  padding: .5rem 1rem !important;
  background: transparent !important;
}
.tab-nav button.selected {
  color: var(--accent) !important;
  border-bottom-color: var(--accent) !important;
}

/* ── Review Table (custom HTML — replaces gr.Dataframe) ── */
#review-table-wrap {
  overflow-x: auto;
  width: 100%;
  border-radius: var(--radius);
  border: 1px solid var(--border);
}
.rt-table {
  width: 100%;
  min-width: 960px;
  border-collapse: collapse;
  font-family: var(--mono);
  font-size: .78rem;
}
.rt-table th {
  background: #1a1f2e;
  color: var(--accent);
  font-family: var(--display);
  font-size: .68rem;
  letter-spacing: .1em;
  text-transform: uppercase;
  padding: .65rem .8rem;
  text-align: left;
  border-bottom: 2px solid var(--border);
  white-space: nowrap;
  position: sticky;
  top: 0;
  z-index: 2;
}
.rt-table td {
  padding: .55rem .8rem;
  border-bottom: 1px solid var(--border);
  vertical-align: top;
  color: var(--text);
}
.rt-table tr:nth-child(even) td { background: rgba(255,255,255,.02); }
.rt-table tr:hover td { background: rgba(79,142,247,.06); transition: background .15s; }
.rt-col-num    { width: 45px;  text-align: center; color: var(--muted); }
.rt-col-label  { min-width: 160px; font-weight: 600; }
.rt-col-evid   { min-width: 240px; line-height: 1.6; color: var(--muted); }
.rt-col-sent   { width: 80px;  text-align: center; }
.rt-col-paper  { width: 70px;  text-align: center; }
.rt-col-appr   { width: 72px;  text-align: center; }
.rt-col-rename { min-width: 140px; }
.rt-col-reason { min-width: 200px; }
.rt-check {
  width: 18px; height: 18px; cursor: pointer;
  accent-color: var(--success);
}
.rt-input, .rt-textarea {
  width: 100%; box-sizing: border-box;
  padding: .32rem .55rem;
  background: #0d0f14;
  border: 1px solid var(--border);
  color: var(--text);
  border-radius: 5px;
  font-family: var(--sans);
  font-size: .78rem;
  transition: border-color .15s;
}
.rt-input:focus, .rt-textarea:focus {
  border-color: var(--accent);
  outline: none;
  box-shadow: 0 0 0 2px rgba(79,142,247,.18);
}
.rt-textarea { resize: vertical; min-height: 46px; }
.rt-empty { color: var(--muted); padding: 1.5rem; font-style: italic; text-align: center; }

/* ── File upload ── */
.file-drop {
  border: 2px dashed var(--border) !important;
  border-radius: var(--radius) !important;
  background: var(--surface) !important;
  color: var(--muted) !important;
  transition: border-color .2s, background .2s !important;
}
.file-drop:hover {
  border-color: var(--accent) !important;
  background: #151c2e !important;
}

/* ── Textbox ── */
textarea, input[type=text] {
  background: var(--surface) !important;
  border: 1px solid var(--border) !important;
  color: var(--text) !important;
  font-family: var(--sans) !important;
  border-radius: var(--radius) !important;
}
textarea:focus, input[type=text]:focus {
  border-color: var(--accent) !important;
  box-shadow: 0 0 0 3px rgba(79,142,247,.15) !important;
}

/* ── Dropdown ── */
select, .gr-dropdown {
  background: var(--surface) !important;
  border: 1px solid var(--border) !important;
  color: var(--text) !important;
  border-radius: var(--radius) !important;
}

/* ── Misc ── */
label { color: var(--muted) !important; font-size: .8rem !important; font-family: var(--mono) !important; }
"""


# ──────────────────────────────────────────────────────────────────────────────
# Build UI
# ──────────────────────────────────────────────────────────────────────────────

def build_app():
    agent = BERTopicAgent()

    with gr.Blocks(
        title="BERTopic Agent",
        theme=gr.themes.Base(
            primary_hue="blue",
            neutral_hue="slate",
            font=[gr.themes.GoogleFont("Lato"), "sans-serif"],
        ),
        css=CUSTOM_CSS
    ) as demo:

        # ── State ──────────────────────────────────────────────────────────
        state_csv_path   = gr.State(None)          # path to uploaded CSV
        state_charts     = gr.State({})             # {label: iframe_src}
        state_downloads  = gr.State([])             # list of file paths

        # ── Header ────────────────────────────────────────────────────────
        gr.HTML("""
        <div class="app-header">
          <h1>⬡ BERTopic Agent</h1>
          <p>Conversational topic modelling powered by BERTopic + LLM review</p>
        </div>
        """)

        # ── Phase progress ─────────────────────────────────────────────────
        phase_display = gr.HTML(value=phase_html("idle"), elem_id="phase-display")

        # ══════════════════════════════════════════════════════════════════
        # SECTION 1 — Data Input
        # ══════════════════════════════════════════════════════════════════
        gr.HTML('<div class="section-card"><p class="section-title">① Data Input</p>')
        with gr.Row():
            csv_upload = gr.File(
                label="Upload CSV (one document per row)",
                file_types=[".csv"],
                elem_classes=["file-drop"],
            )
            with gr.Column(scale=2):
                gr.Markdown(
                    "**Expected format:** CSV with at least a `text` column. "
                    "Optional columns: `id`, `title`, `source`.\n\n"
                    "After uploading, tell the agent which column holds your documents.",
                    elem_id="upload-hint",
                )
        gr.HTML('</div>')

        # ══════════════════════════════════════════════════════════════════
        # SECTION 2 — Agent Conversation
        # ══════════════════════════════════════════════════════════════════
        gr.HTML('<div class="section-card"><p class="section-title">② Agent Conversation</p>')
        chatbot = gr.Chatbot(
            label="",
            height=420,
            show_label=False,
            # type="tuples",
            elem_classes=["chatbot-wrap"],
            avatar_images=(
                None,                          # user — use default
                "https://em-content.zobj.net/source/apple/391/robot_1f916.png",
            ),
        )
        with gr.Row():
            user_input = gr.Textbox(
                placeholder="Ask the agent to run topic modelling, explain results, adjust parameters…",
                lines=1,
                show_label=False,
                scale=8,
            )
            send_btn = gr.Button("Send ⟶", elem_classes=["btn-primary"], scale=1)
        gr.HTML('</div>')

        # ══════════════════════════════════════════════════════════════════
        # SECTION 3 — Results
        # ══════════════════════════════════════════════════════════════════
        gr.HTML('<div class="section-card"><p class="section-title">③ Results</p>')
        with gr.Tabs(elem_classes=["tab-nav"]):

            # ── Tab A: Review Table ──────────────────────────────────────
            with gr.Tab("📋  Review Table"):
                gr.HTML(
                    "<p style='color:var(--muted,#7a8299);font-size:.8rem;margin:0 0 .6rem'>"
                    "Check ✓ Approve, optionally fill Rename To / Reasoning, then click Submit.</p>"
                )
                review_table_html = gr.HTML(
                    value="<p class='rt-empty' style='color:#7a8299;padding:1rem;font-style:italic'>"
                          "No topics yet — run topic modelling first.</p>",
                    label="",
                )
                # Hidden textbox holds JSON-serialised table rows, kept in sync by JS
                review_json_state = gr.Textbox(
                    value="[]",
                    visible=False,
                    elem_id="review-json-state",
                )
                submit_review_btn = gr.Button(
                    "Submit Review  ✓",
                    elem_classes=["btn-primary"],
                    size="sm",
                )

            # ── Tab B: Charts ────────────────────────────────────────────
            with gr.Tab("📊  Charts"):
                chart_selector = gr.Dropdown(
                    label="Select chart",
                    choices=[],
                    interactive=True,
                )
                chart_display = gr.HTML(
                    value="<p style='color:#7a8299;padding:1rem'>Charts will appear here after topic modelling completes.</p>"
                )

            # ── Tab C: Download ──────────────────────────────────────────
            with gr.Tab("⬇  Download"):
                download_files = gr.File(
                    label="Output files",
                    file_count="multiple",
                    interactive=False,
                )

        gr.HTML('</div>')

        # ══════════════════════════════════════════════════════════════════
        # Event handlers
        # ══════════════════════════════════════════════════════════════════

        # ── Store CSV path when file is uploaded ─────────────────────────
        def on_csv_upload(file):
            if file is None:
                return None
            return file.name

        csv_upload.change(
            fn=on_csv_upload,
            inputs=[csv_upload],
            outputs=[state_csv_path],
        )

        # ── Send message to agent ─────────────────────────────────────────
        def on_send(message, history, csv_path, charts_state, downloads_state):
            """Stream agent responses and update all UI elements."""
            if not message.strip():
                yield history, "", phase_html("idle"), gr.update(), gr.update(), None, charts_state, downloads_state
                return

            # Accumulate streamed updates
            last_history    = history
            last_phase      = "idle"
            last_charts     = charts_state
            last_downloads  = downloads_state
            new_topics      = []

            for result in agent.handle_message(message, history, csv_path=csv_path):
                # Agent yields: (history, phase, charts_dict, downloads_list, topic_rows)
                (last_history, last_phase,
                 last_charts, last_downloads, new_topics) = result

                _html, _json = _rows_to_html(new_topics) if new_topics else (gr.update(), gr.update())
                yield (
                    last_history,
                    "",                          # clear input
                    phase_html(last_phase),
                    _html,
                    _json,
                    gr.update(),                 # downloads updated below
                    last_charts,
                    last_downloads,
                )

            # Final yield with downloads
            _html, _json = _rows_to_html(new_topics) if new_topics else (gr.update(), gr.update())
            yield (
                last_history,
                "",
                phase_html(last_phase),
                _html,
                _json,
                last_downloads if last_downloads else None,
                last_charts,
                last_downloads,
            )

        def _rows_to_html(rows):
            """Convert agent row data → (html_str, json_str).
            Accepts list-of-lists, list-of-dicts, list-of-tuples, or pd.DataFrame.
            """
            if rows is None:
                rows = []

            # ── Normalise to list-of-lists in TABLE_COLS order ──────────────
            # Alias map: agent may use slightly different key names
            _ALIASES = {
                "top evidence (truncated)": "Top Evidence",
                "top evidence":             "Top Evidence",
                "topic label":              "Topic Label",
                "rename to":                "Rename To",
                "approve":                  "Approve",
                "reasoning":                "Reasoning",
                "sentences":                "Sentences",
                "papers":                   "Papers",
                "#":                        "#",
            }

            if isinstance(rows, pd.DataFrame):
                # Normalise column names via alias map
                rows.columns = [_ALIASES.get(c.lower(), c) for c in rows.columns]
                for col in TABLE_COLS:
                    if col not in rows.columns:
                        rows[col] = ""
                padded = rows[TABLE_COLS].fillna("").values.tolist()

            elif rows and isinstance(rows[0], dict):
                # Remap keys through alias table then extract in TABLE_COLS order
                normalised = []
                for r in rows:
                    remapped = {_ALIASES.get(k.lower(), k): v for k, v in r.items()}
                    normalised.append([remapped.get(col, "") for col in TABLE_COLS])
                padded = normalised

            else:
                # List of lists / tuples — pad / truncate to exactly 8 columns
                padded = []
                for r in rows:
                    row = list(r) + [""] * max(0, 8 - len(list(r)))
                    padded.append(row[:8])

            if not padded:
                html = "<p class='rt-empty'>No topics yet — run topic modelling first.</p>"
                return html, "[]"

            # ── Serialise to JSON (drives JS state) ─────────────────────────
            # Convert non-serialisable types (numpy bools, NaN, etc.) safely
            def _safe(v):
                if v is None or (isinstance(v, float) and v != v):   # NaN check
                    return ""
                try:
                    import numpy as np
                    if isinstance(v, (np.bool_, np.integer, np.floating)):
                        return v.item()
                except ImportError:
                    pass
                return v

            safe_padded = [[_safe(c) for c in row] for row in padded]
            json_data = json.dumps(safe_padded, ensure_ascii=False)

            # ── Build HTML rows ──────────────────────────────────────────────
            def _esc(v, attr=False):
                s = str(v) if v is not None else ""
                s = s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
                if attr:
                    s = s.replace('"', "&quot;").replace("'", "&#39;")
                return s

            rows_html = ""
            for i, row in enumerate(safe_padded):
                num, label, evidence, sentences, papers, approve, rename, reasoning = row
                checked = "checked" if str(approve).lower() in ("true", "1", "yes", "✓") else ""
                rows_html += f"""
        <tr data-row="{i}">
          <td class="rt-col-num">{_esc(num)}</td>
          <td class="rt-col-label">{_esc(label)}</td>
          <td class="rt-col-evid">{_esc(evidence)}</td>
          <td class="rt-col-sent">{_esc(sentences)}</td>
          <td class="rt-col-paper">{_esc(papers)}</td>
          <td class="rt-col-appr">
            <input type="checkbox" class="rt-check" {checked}
              onchange="rtUpdate({i},5,this.checked)">
          </td>
          <td class="rt-col-rename">
            <input type="text" class="rt-input"
              value="{_esc(rename, attr=True)}"
              oninput="rtUpdate({i},6,this.value)"
              placeholder="new label…">
          </td>
          <td class="rt-col-reason">
            <textarea class="rt-textarea" rows="2"
              oninput="rtUpdate({i},7,this.value)"
              placeholder="notes…">{_esc(reasoning)}</textarea>
          </td>
        </tr>"""

            # ── Assemble full HTML block ─────────────────────────────────────
            # Use a unique suffix so re-renders don't collide with stale globals
            uid = id(safe_padded)
            html = f"""<div id="review-table-wrap">
<script>
(function(){{
  // Initialise or refresh the shared row-data store
  window._rtData = {json_data};

  window.rtUpdate = function(idx, col, val) {{
    if (!window._rtData || !window._rtData[idx]) return;
    window._rtData[idx][col] = val;
    // Push updated JSON into the hidden Gradio textbox
    // Try both possible DOM locations Gradio uses across versions
    var el = (
      document.querySelector('#review-json-state textarea') ||
      document.querySelector('[id*="review-json-state"] textarea') ||
      document.querySelector('textarea[data-testid="textbox"]')
    );
    if (!el) {{ console.warn("rtUpdate: hidden textbox not found"); return; }}
    var setter = Object.getOwnPropertyDescriptor(
      window.HTMLTextAreaElement.prototype, 'value'
    ).set;
    setter.call(el, JSON.stringify(window._rtData));
    el.dispatchEvent(new Event('input', {{bubbles: true}}));
  }};
}})();
</script>
<table class="rt-table">
  <thead><tr>
    <th class="rt-col-num">#</th>
    <th class="rt-col-label">Topic Label</th>
    <th class="rt-col-evid">Top Evidence</th>
    <th class="rt-col-sent">Sentences</th>
    <th class="rt-col-paper">Papers</th>
    <th class="rt-col-appr">Approve</th>
    <th class="rt-col-rename">Rename To</th>
    <th class="rt-col-reason">Reasoning</th>
  </tr></thead>
  <tbody>{rows_html}</tbody>
</table></div>"""
            return html, json_data

        send_outputs = [
            chatbot,
            user_input,
            phase_display,
            review_table_html,
            review_json_state,
            download_files,
            state_charts,
            state_downloads,
        ]

        send_btn.click(
            fn=on_send,
            inputs=[user_input, chatbot, state_csv_path, state_charts, state_downloads],
            outputs=send_outputs,
        )
        user_input.submit(
            fn=on_send,
            inputs=[user_input, chatbot, state_csv_path, state_charts, state_downloads],
            outputs=send_outputs,
        )

        # ── Submit review ─────────────────────────────────────────────────
        def on_submit_review(json_str, history, charts_state, downloads_state):
            """Pass edited table back to the agent as a review submission."""
            try:
                rows = json.loads(json_str) if json_str and json_str.strip() not in ("", "[]") else []
            except Exception:
                rows = []

            last_history   = history
            last_phase     = "review"
            last_charts    = charts_state
            last_downloads = downloads_state

            for result in agent.handle_review(rows, history):
                (last_history, last_phase,
                 last_charts, last_downloads) = result

                yield (
                    last_history,
                    phase_html(last_phase),
                    last_downloads if last_downloads else None,
                    last_charts,
                    last_downloads,
                )

            yield (
                last_history,
                phase_html(last_phase),
                last_downloads if last_downloads else None,
                last_charts,
                last_downloads,
            )

        submit_review_btn.click(
            fn=on_submit_review,
            inputs=[review_json_state, chatbot, state_charts, state_downloads],
            outputs=[chatbot, phase_display, download_files, state_charts, state_downloads],
        )

        # ── Chart selector ────────────────────────────────────────────────
        def on_chart_change(chart_name, charts_state):
            src = (charts_state or {}).get(chart_name, "")
            return make_chart_html(src)

        chart_selector.change(
            fn=on_chart_change,
            inputs=[chart_selector, state_charts],
            outputs=[chart_display],
        )

        # ── Update chart dropdown when charts_state changes ──────────────
        def on_charts_state_change(charts_state):
            choices = list((charts_state or {}).keys())
            if choices:
                return gr.update(choices=choices, value=choices[0])
            return gr.update(choices=[], value=None)

        state_charts.change(
            fn=on_charts_state_change,
            inputs=[state_charts],
            outputs=[chart_selector],
        )

    return demo


# ──────────────────────────────────────────────────────────────────────────────
# Entry point
# ──────────────────────────────────────────────────────────────────────────────

app = build_app()

app.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    show_error=True
)