"""
BERTopic Topic Modelling Agent — Gradio Interface
"""
import os
os.environ.pop("GRADIO_HOT_RELOAD", None)
import gradio as gr
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;
}

/* ── Dataframe ── */
.dataframe-wrap table {
  font-family: var(--mono) !important;
  font-size: .78rem !important;
}
.dataframe-wrap th {
  background: #1a1f2e !important;
  color: var(--accent) !important;
  font-family: var(--display) !important;
  font-size: .7rem !important;
  letter-spacing: .08em !important;
  text-transform: uppercase !important;
}
.dataframe-wrap td { border-color: var(--border) !important; }

/* ── 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") 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"):
                review_table = gr.Dataframe(
                    value=EMPTY_TABLE,
                    headers=TABLE_COLS,
                    datatype=["number", "str", "str", "number", "number", "bool", "str", "str"],
                    interactive=True,
                    wrap=True,
                    elem_classes=["dataframe-wrap"],
                    label="Topic review — edit Approve / Rename To / Reasoning, then click Submit",
                )
                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: str, history: list, csv_path, charts_state, downloads_state):
            """Stream agent responses and update all UI elements."""
            if not message.strip():
                yield history, "", phase_html("idle"), [], 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

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

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

        def _rows_to_df(rows):
            if not rows:
                return EMPTY_TABLE
            df = pd.DataFrame(rows, columns=TABLE_COLS)
            return df

        send_outputs = [
            chatbot,
            user_input,
            phase_display,
            review_table,
            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(table_df, history, charts_state, downloads_state):
            """Pass edited table back to the agent as a review submission."""
            rows = table_df.values.tolist() if isinstance(table_df, pd.DataFrame) else table_df or []

            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_table, chatbot, state_charts, state_downloads],
            outputs=[chatbot, phase_display, download_files, state_charts, state_downloads],
        )

        # ── Chart selector ────────────────────────────────────────────────
        def on_chart_change(chart_name: str, charts_state: dict):
            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: dict):
            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
# ──────────────────────────────────────────────────────────────────────────────

# Build the app at module level so HuggingFace Spaces can import it correctly.
# The `if __name__ == "__main__"` guard prevented launch() from running on Spaces
# because HF imports app.py rather than executing it as a script.
app = build_app()
app.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=False,
    show_error=True,
    ssr_mode=False,
    theme=gr.themes.Base(
        primary_hue="blue",
        neutral_hue="slate",
        font=[gr.themes.GoogleFont("Lato"), "sans-serif"],
    ),
    css=CUSTOM_CSS,
)