"""
Deep Research in Supply Chain — Gradio UI (ChatGPT-inspired skeleton).

Multi-channel conversation interface that will wrap the Helicase
autonomous research agent. Research step is mocked for now so the
UX can be validated before Helicase + DashScope are wired in.

Runs on Hugging Face Spaces (free tier) behind an <iframe> on
supplychaindatahub.org/deep_research.
"""

from __future__ import annotations

import base64
import json
import os
import time
import uuid
import zlib
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Generator
from urllib.parse import quote

import gradio as gr
from fastapi import FastAPI
from fastapi.responses import HTMLResponse

from research_engine import run_research, _keys_configured, generate_kg_answer

# KG viewer is served by this same HF Space at /kg.
KG_VIEWER_BASE = os.getenv("KG_VIEWER_URL", "/kg")

_HERE = os.path.dirname(os.path.abspath(__file__))
with open(os.path.join(_HERE, "kg_viewer.html"), "r", encoding="utf-8") as _f:
    _KG_VIEWER_HTML = _f.read()


# ---------- session model ---------------------------------------------------

# Channels are stored as plain dicts so the whole state can be serialised
# into gr.BrowserState (localStorage) and survive page refreshes.
#   {
#     "id": str,
#     "title": str,
#     "created_at": ISO timestamp (str),
#     "messages": [{"role": "user"|"assistant", "content": str}],
#     "kg_nodes": [...],
#     "kg_edges": [...],
#   }


def _new_channel() -> dict:
    return {
        "id": str(uuid.uuid4())[:8],
        "title": "New chat",
        "created_at": datetime.now().isoformat(),
        "messages": [],
        "kg_nodes": [],
        "kg_edges": [],
        "kg_ref2url": {},
    }


def _empty_state() -> dict:
    ch = _new_channel()
    return {"channels": {ch["id"]: ch}, "active": ch["id"]}


def _hydrate(state) -> dict:
    """Restore state from BrowserState (may arrive as None or malformed)."""
    if not isinstance(state, dict) or not state.get("channels"):
        return _empty_state()
    # Make sure the active channel still exists.
    if state.get("active") not in state["channels"]:
        state["active"] = next(iter(state["channels"]))
    return state


def _get_channel(state: dict, channel_id: str) -> dict:
    return state["channels"][channel_id]


def _sidebar_choices(state: dict) -> list[tuple[str, str]]:
    items = []
    channels = sorted(
        state["channels"].values(), key=lambda c: c.get("created_at", ""), reverse=True
    )
    for ch in channels:
        label = ch["title"] if ch["title"] != "New chat" else "New chat"
        items.append((label, ch["id"]))
    return items


# ---------- streaming knowledge-graph renderer ------------------------------

NODE_COLORS = {
    "company":   {"bg": "#4f46e5", "border": "#3730a3"},   # indigo
    "material":  {"bg": "#10b981", "border": "#047857"},   # emerald
    "product":   {"bg": "#f59e0b", "border": "#b45309"},   # amber
    "process":   {"bg": "#ec4899", "border": "#9d174d"},   # pink
    "region":    {"bg": "#3b82f6", "border": "#1d4ed8"},   # blue
    "risk":      {"bg": "#ef4444", "border": "#991b1b"},   # red
}


def _render_kg_html(nodes: list[dict], edges: list[dict], kg_id: str, status: str = "", progress: float = 0.0) -> str:
    """Render an interactive vis-network KG as a self-contained HTML snippet."""
    vis_nodes = []
    for n in nodes:
        c = NODE_COLORS.get(n["type"], {"bg": "#64748b", "border": "#334155"})
        vis_nodes.append({
            "id": n["id"],
            "label": n["label"],
            "title": f"{n['type'].title()} · confidence {n.get('confidence', 0.8):.2f}",
            "group": n["type"],
            "color": {"background": c["bg"], "border": c["border"]},
            "font": {"color": "#fff", "size": 13},
            "shape": "dot",
            "size": 20,
        })
    vis_edges = [{
        "from": e["from"], "to": e["to"], "label": e["label"],
        "font": {"size": 10, "color": "#64748b", "strokeWidth": 0},
        "color": {"color": "#cbd5e1"},
        "arrows": "to",
    } for e in edges]

    nodes_json = json.dumps(vis_nodes)
    edges_json = json.dumps(vis_edges)

    progress_bar = (
        f'<div style="height:4px;background:#e5e7eb;border-radius:999px;overflow:hidden;margin:6px 0 10px 0;">'
        f'<div style="height:100%;width:{progress*100:.0f}%;background:linear-gradient(90deg,#4f46e5,#ec4899);transition:width 0.6s ease;"></div>'
        f'</div>'
    ) if progress < 1.0 else ""

    status_badge = f'<span style="font-size:12px;color:#4f46e5;font-weight:600;">{status}</span>' if status else ''

    return f"""
<div style="border:1px solid #e5e7eb;border-radius:12px;padding:14px;background:#ffffff;margin-top:10px;box-shadow:0 2px 8px rgba(15,23,42,0.04);">
  <div style="display:flex;align-items:center;justify-content:space-between;margin-bottom:6px;">
    <div style="font-weight:600;font-size:14px;color:#0f172a;">🕸️ Knowledge Graph · <span style="color:#64748b;font-weight:500;">{len(nodes)} nodes · {len(edges)} edges</span></div>
    {status_badge}
  </div>
  {progress_bar}
  <div id="{kg_id}" style="height:360px;border:1px solid #f1f5f9;border-radius:8px;background:#fafbfc;"></div>
  <div style="display:flex;gap:14px;flex-wrap:wrap;margin-top:8px;font-size:11px;color:#64748b;">
    <span><span style="display:inline-block;width:10px;height:10px;border-radius:50%;background:#4f46e5;margin-right:4px;"></span>Company</span>
    <span><span style="display:inline-block;width:10px;height:10px;border-radius:50%;background:#10b981;margin-right:4px;"></span>Material</span>
    <span><span style="display:inline-block;width:10px;height:10px;border-radius:50%;background:#f59e0b;margin-right:4px;"></span>Product</span>
    <span><span style="display:inline-block;width:10px;height:10px;border-radius:50%;background:#3b82f6;margin-right:4px;"></span>Region</span>
    <span><span style="display:inline-block;width:10px;height:10px;border-radius:50%;background:#ef4444;margin-right:4px;"></span>Risk</span>
  </div>
</div>
<script>
(function() {{
  function init() {{
    if (typeof vis === 'undefined') {{ setTimeout(init, 200); return; }}
    var el = document.getElementById("{kg_id}");
    if (!el) return;
    var nodes = new vis.DataSet({nodes_json});
    var edges = new vis.DataSet({edges_json});
    var net = new vis.Network(el, {{nodes:nodes, edges:edges}}, {{
      physics: {{ stabilization:{{iterations:80}}, barnesHut:{{gravitationalConstant:-2800, springLength:130}} }},
      interaction: {{ hover:true, tooltipDelay:120 }},
      nodes:{{ borderWidth: 2 }},
      edges:{{ smooth:{{type:'dynamic'}} }}
    }});
  }}
  if (typeof vis === 'undefined') {{
    var s = document.createElement('script');
    s.src = 'https://unpkg.com/vis-network@9.1.9/standalone/umd/vis-network.min.js';
    s.onload = init;
    document.head.appendChild(s);
  }} else {{
    init();
  }}
}})();
</script>
""".strip()


# Scripted "growing" KG for the mock — incrementally builds Tesla-4680-ish supply chain.
_MOCK_STEPS = [
    {
        "status": "🔍 Iteration 1 · searching the web…",
        "nodes": [
            {"id": "tesla", "label": "Tesla", "type": "company", "confidence": 0.98},
            {"id": "cell_4680", "label": "4680 Cell", "type": "product", "confidence": 0.95},
        ],
        "edges": [{"from": "tesla", "to": "cell_4680", "label": "manufactures"}],
    },
    {
        "status": "🧠 Iteration 2 · extracting materials…",
        "nodes": [
            {"id": "nmc", "label": "NMC Cathode", "type": "material", "confidence": 0.92},
            {"id": "anode", "label": "Silicon-Graphite Anode", "type": "material", "confidence": 0.88},
        ],
        "edges": [
            {"from": "cell_4680", "to": "nmc", "label": "contains"},
            {"from": "cell_4680", "to": "anode", "label": "contains"},
        ],
    },
    {
        "status": "🔎 Iteration 3 · tracing cathode metals…",
        "nodes": [
            {"id": "cobalt", "label": "Cobalt", "type": "material", "confidence": 0.90},
            {"id": "nickel", "label": "Nickel", "type": "material", "confidence": 0.91},
            {"id": "lithium", "label": "Lithium", "type": "material", "confidence": 0.93},
        ],
        "edges": [
            {"from": "nmc", "to": "cobalt", "label": "contains"},
            {"from": "nmc", "to": "nickel", "label": "contains"},
            {"from": "nmc", "to": "lithium", "label": "contains"},
        ],
    },
    {
        "status": "🧭 Iteration 4 · identifying suppliers…",
        "nodes": [
            {"id": "glencore", "label": "Glencore", "type": "company", "confidence": 0.87},
            {"id": "bhp", "label": "BHP", "type": "company", "confidence": 0.86},
            {"id": "albemarle", "label": "Albemarle", "type": "company", "confidence": 0.84},
        ],
        "edges": [
            {"from": "glencore", "to": "cobalt", "label": "supplies"},
            {"from": "bhp", "to": "nickel", "label": "supplies"},
            {"from": "albemarle", "to": "lithium", "label": "supplies"},
        ],
    },
    {
        "status": "🌍 Iteration 5 · mapping geographies & risks…",
        "nodes": [
            {"id": "drc", "label": "DR Congo", "type": "region", "confidence": 0.95},
            {"id": "australia", "label": "Australia", "type": "region", "confidence": 0.94},
            {"id": "risk_esg", "label": "ESG / child-labour risk", "type": "risk", "confidence": 0.80},
        ],
        "edges": [
            {"from": "cobalt", "to": "drc", "label": "mined in"},
            {"from": "lithium", "to": "australia", "label": "mined in"},
            {"from": "cobalt", "to": "risk_esg", "label": "associated with"},
        ],
    },
]


def _encode_kg_url(question: str, nodes: list[dict], edges: list[dict], ref2url: dict | None = None) -> str:
    """Compress the KG into a URL-safe base64 payload for the viewer."""
    payload = json.dumps(
        {"prompt": question, "nodes": nodes, "edges": edges, "ref2url": ref2url or {}},
        ensure_ascii=False, separators=(",", ":"),
    )
    compressed = zlib.compress(payload.encode("utf-8"), level=9)
    b64 = base64.urlsafe_b64encode(compressed).decode("ascii").rstrip("=")
    return f"{KG_VIEWER_BASE}#data={b64}"


def _resolve_inline_citations(text: str, ref2url: dict) -> str:
    """Turn inline [n] citations in the NL answer into clickable links."""
    import re
    def replace(match):
        idx = match.group(1)
        ref = (ref2url or {}).get(str(idx)) or {}
        url = ref.get("url", "")
        title = (ref.get("title") or "").replace('"', "'")
        if url:
            return f"<a href=\"{url}\" target=\"_blank\" rel=\"noopener\" title=\"{title}\" style=\"color:#4f46e5;font-weight:600;text-decoration:none;\">[{idx}]</a>"
        return f"[{idx}]"
    return re.sub(r"\[(\d+)\]", replace, text)


def _summary(question: str, nodes: list[dict], edges: list[dict], iterations: int, converged: bool, reason: str, final_u: float, ref2url: dict | None = None) -> str:
    type_counts: dict[str, int] = {}
    for n in nodes:
        t = n.get("type", "entity")
        type_counts[t] = type_counts.get(t, 0) + 1
    bits = ", ".join(f"**{c}** {t}{'s' if c != 1 else ''}" for t, c in sorted(type_counts.items(), key=lambda x: -x[1]))
    viewer = _encode_kg_url(question, nodes, edges, ref2url)
    status = (
        f"🧬 **Converged** after {iterations} helical iteration{'s' if iterations != 1 else ''} — {reason}."
        if converged
        else f"⏹️ Reached max iterations ({iterations}) without full convergence — {reason}."
    )

    # GraphRAG-style natural language answer grounded in the KG.
    nl_answer = ""
    try:
        raw = generate_kg_answer(question, nodes, edges, ref2url or {})
        if raw:
            nl_answer = _resolve_inline_citations(raw, ref2url or {})
    except Exception:
        nl_answer = ""

    answer_block = f"{nl_answer}\n\n---\n\n" if nl_answer else ""

    return (
        f"{answer_block}"
        f"**Knowledge graph built**: {len(nodes)} nodes · {len(edges)} edges — {bits}. "
        f"{status}\n\n"
        f"<a href=\"{viewer}\" target=\"_blank\" rel=\"noopener\" "
        "style=\"display:inline-block;background:#4f46e5;color:#fff;padding:10px 18px;"
        "border-radius:10px;font-weight:600;text-decoration:none;margin-top:4px;\">"
        "📊 Open Knowledge Graph →</a>\n\n"
        "<sub>Ask a follow-up to extend the graph — it uses the same KG as context.</sub>"
    )


def _run_research_streaming(question: str, prior_kg_nodes: list[dict] | None = None, prior_kg_edges: list[dict] | None = None, prior_ref2url: dict | None = None
                            ) -> Generator[tuple[str, list[dict], list[dict], dict], None, None]:
    """Yield ``(chat_text, nodes, edges, ref2url)`` — chat text updates as
    research proceeds; final yield carries the summary with a link to the
    /kg/ viewer page plus the full citation map."""
    nodes: list[dict] = list(prior_kg_nodes or [])
    edges: list[dict] = list(prior_kg_edges or [])
    ref2url: dict = dict(prior_ref2url or {})
    seen_ids = {n["id"] for n in nodes}

    if _keys_configured():
        # ----- Helicase-style convergence (mirrors helicase_orchestrator.py) -----
        # Three criteria — any one triggers convergence:
        #   1. CONFIDENT — U_memory < τ_U for `patience` consecutive iterations
        #   2. STAGNANT  — |ΔU_memory| < τ_Δ for `patience` iterations AND iteration > 2
        #   3. MAX       — iteration reached max_iterations
        max_iterations        = int(os.getenv("MAX_HELIX_ITERATIONS", "10"))
        uncertainty_threshold = float(os.getenv("UNCERTAINTY_THRESHOLD", "0.3"))   # τ_U
        convergence_delta     = float(os.getenv("CONVERGENCE_DELTA",   "0.05"))    # τ_Δ (5%)
        convergence_patience  = int(os.getenv("CONVERGENCE_PATIENCE",  "3"))

        yield "🔬 Starting helical research loop…", nodes, edges, ref2url

        u_memory_history: list[float] = []
        converge_count = 0
        stagnant_count = 0
        converged = False
        reason = "max iterations reached"
        iteration = 0

        for iteration in range(1, max_iterations + 1):
            # Run one helical round (search → reason → extract → add to KG).
            for status, n, e in run_research(question, nodes, edges, ref2url=ref2url):
                if status == "keys_missing":
                    break
                nodes, edges = n, e
                yield f"{status} · Helix {iteration}/{max_iterations}", nodes, edges, ref2url

            # Compute U_memory proxy: 1 − average node confidence.
            if nodes:
                avg_conf = sum(n.get("confidence", 0.75) for n in nodes) / len(nodes)
                u_memory = max(0.0, min(1.0, 1.0 - avg_conf))
            else:
                u_memory = 1.0

            # --- Stagnant check ---
            if u_memory_history:
                prev_u = u_memory_history[-1]
                relative_change = abs(prev_u - u_memory) / prev_u if prev_u > 0 else 0.0
                if relative_change < convergence_delta:
                    stagnant_count += 1
                else:
                    stagnant_count = 0
            u_memory_history.append(u_memory)

            # --- Criterion 1: CONFIDENT (U_memory below τ_U for `patience` rounds) ---
            if u_memory < uncertainty_threshold:
                converge_count += 1
                if converge_count >= convergence_patience:
                    converged = True
                    reason = f"confident (U={u_memory:.3f} < τ_U={uncertainty_threshold:.2f} for {converge_count} rounds)"
            else:
                converge_count = 0

            # --- Criterion 2: STAGNANT (ΔU below τ_Δ for `patience` rounds, after round 2) ---
            if not converged and stagnant_count >= convergence_patience and iteration > 2:
                converged = True
                reason = f"stagnant (|ΔU|<{convergence_delta*100:.0f}% for {stagnant_count} rounds)"

            yield (
                f"🌀 Helix {iteration}: U_mem={u_memory:.3f} · "
                f"stagnant={stagnant_count}/{convergence_patience} · "
                f"confident={converge_count}/{convergence_patience}",
                nodes, edges, ref2url,
            )

            if converged:
                yield f"🧬 CONVERGED — {reason}", nodes, edges, ref2url
                break

        final_u = u_memory_history[-1] if u_memory_history else 1.0
        yield _summary(question, nodes, edges, iterations=iteration, converged=converged, reason=reason, final_u=final_u, ref2url=ref2url), nodes, edges, ref2url
        return

    # ---- MOCK fallback ----
    for i, step in enumerate(_MOCK_STEPS):
        for n in step["nodes"]:
            if n["id"] not in seen_ids:
                nodes.append(n)
                seen_ids.add(n["id"])
        edges.extend(step["edges"])
        yield step["status"], nodes, edges, ref2url
        time.sleep(1.2)
    avg_conf = sum(n.get("confidence", 0.75) for n in nodes) / max(1, len(nodes))
    yield _summary(
        question, nodes, edges,
        iterations=len(_MOCK_STEPS), converged=True,
        reason="mock run complete", final_u=max(0.0, 1.0 - avg_conf),
        ref2url=ref2url,
    ), nodes, edges, ref2url


# ---------- handlers --------------------------------------------------------

_IN_PROGRESS_MARKERS = ("🔬 Starting", "🔬 Planning", "🔍 Running", "🧠 Extracting",
                        "✨ Added", "🌀 Helix", "🔁 Round", "💤 Round", "🔬 Starting helical")


def on_load(state):
    """Restore from BrowserState. Replace any stale in-progress assistant
    message (left over from a refresh mid-stream) with a clean note."""
    state = _hydrate(state)
    for ch in state["channels"].values():
        msgs = ch.get("messages") or []
        if msgs and msgs[-1].get("role") == "assistant":
            content = msgs[-1].get("content", "")
            if any(m in content for m in _IN_PROGRESS_MARKERS) and "📊" not in content and "CONVERGED" not in content:
                msgs[-1] = {"role": "assistant",
                            "content": "⏸️ *Research was interrupted (page refreshed). Ask again to continue, or start a new chat.*"}
    ch = _get_channel(state, state["active"])
    welcome_visible = len(ch["messages"]) == 0
    return (
        state,
        gr.update(choices=_sidebar_choices(state), value=state["active"]),
        list(ch["messages"]),
        gr.update(visible=welcome_visible),
    )


def on_new_channel(state):
    state = _hydrate(state)
    ch = _new_channel()
    state["channels"][ch["id"]] = ch
    state["active"] = ch["id"]
    return (
        state,
        gr.update(choices=_sidebar_choices(state), value=ch["id"]),
        [],
        "",
        gr.update(visible=True),
    )


def on_switch_channel(state, channel_id):
    state = _hydrate(state)
    if channel_id and channel_id in state["channels"]:
        state["active"] = channel_id
    ch = _get_channel(state, state["active"])
    welcome_visible = len(ch["messages"]) == 0
    return state, list(ch["messages"]), gr.update(visible=welcome_visible)


def on_ask(state, question):
    """Streaming handler — chat text only; KG opens in /kg viewer via link."""
    state = _hydrate(state)
    question = (question or "").strip()
    if not question:
        ch = _get_channel(state, state["active"])
        yield state, list(ch["messages"]), "", gr.update(), gr.update()
        return

    ch = _get_channel(state, state["active"])
    ch["messages"].append({"role": "user", "content": question})
    if ch["title"] == "New chat":
        ch["title"] = question[:42] + ("…" if len(question) > 42 else "")
    ch["messages"].append({"role": "assistant", "content": "🔬 Starting research…"})

    sidebar = gr.update(choices=_sidebar_choices(state), value=state["active"])
    hide_welcome = gr.update(visible=False)

    yield state, list(ch["messages"]), "", sidebar, hide_welcome

    for chat_text, nodes, edges, ref2url in _run_research_streaming(
        question, ch["kg_nodes"], ch["kg_edges"], ch.get("kg_ref2url", {})
    ):
        ch["messages"][-1] = {"role": "assistant", "content": chat_text}
        ch["kg_nodes"] = list(nodes)
        ch["kg_edges"] = list(edges)
        ch["kg_ref2url"] = dict(ref2url)
        yield state, list(ch["messages"]), "", sidebar, hide_welcome


# ---------- UI --------------------------------------------------------------

# ChatGPT-inspired CSS — polished / professional.
CSS = """
/* Hide Gradio chrome. */
footer { display: none !important; }
#sdr-root .gradio-container { max-width: 100% !important; padding: 0 !important; }

/* Google-font fallback — use system first, then Inter. */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

/* Global frame — dark blue ambient background + auto-height window. */
html, body, gradio-app {
    background: radial-gradient(1200px 600px at 20% 10%, #1e3a8a 0%, transparent 55%),
                radial-gradient(1000px 700px at 85% 90%, #312e81 0%, transparent 55%),
                linear-gradient(135deg, #0b1020 0%, #0f172a 55%, #0b1a34 100%) !important;
    min-height: 100vh;
    margin: 0;
    padding: 0;
}

#sdr-root {
    background: transparent;
    color: #111827;
    font-family: "Inter", ui-sans-serif, system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
    min-height: 100vh;
    padding: 20px !important;
    box-sizing: border-box !important;
    letter-spacing: -0.005em;
}

/* Floating "window" — at least a full viewport tall, scroll-friendly. */
#sdr-root > div:first-child {
    max-width: 1280px;
    min-height: calc(100vh - 40px);
    margin: 0 auto !important;
    background: #ffffff;
    border-radius: 16px;
    overflow: hidden;
    box-shadow:
        0 30px 80px rgba(2, 6, 23, 0.45),
        0 10px 30px rgba(2, 6, 23, 0.30),
        0 0 0 1px rgba(255,255,255,0.06);
    display: flex;
}
#sdr-root > div:first-child > div[class*="row"] { flex: 1; display: flex; width: 100%; min-height: 0; }

/* ============ Sidebar — full height flex ============ */
#sdr-sidebar {
    background: linear-gradient(180deg, #0f172a 0%, #1e293b 100%);
    color: #e5e7eb;
    border-right: 1px solid #1f2937;
    padding: 18px 14px !important;
    display: flex !important;
    flex-direction: column !important;
    gap: 8px;
    height: 100%;
    min-height: 0;
    overflow-y: auto;
}

/* Brand block at top. */
#sdr-brand {
    display: flex; align-items: center; gap: 10px;
    padding: 4px 6px 14px 6px;
    border-bottom: 1px solid rgba(255,255,255,0.08);
    margin-bottom: 10px;
}
#sdr-brand .sdr-logo {
    width: 32px; height: 32px; border-radius: 8px;
    background: linear-gradient(135deg,#6366f1,#ec4899);
    display: grid; place-items: center; font-size: 16px;
    box-shadow: 0 2px 8px rgba(99,102,241,0.35);
}
#sdr-brand .sdr-title { font-weight: 700; font-size: 14px; color:#fff; line-height:1.1; }
#sdr-brand .sdr-sub { font-size: 11px; color:#94a3b8; margin-top:2px; }

/* Section heading. */
#sdr-sidebar .prose h3, #sdr-section-head {
    margin: 10px 6px 6px 6px; font-size: 11px; color: #94a3b8;
    text-transform: uppercase; letter-spacing: 0.6px; font-weight: 600;
}

/* New chat button — crisp white-on-dark pill, ChatGPT style. */
#sdr-new-btn { margin-bottom: 2px; }
#sdr-new-btn button {
    background: #ffffff !important;
    color: #0f172a !important;
    border: 1px solid rgba(255,255,255,0.9) !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
    font-size: 14px !important;
    padding: 10px 14px !important;
    width: 100% !important;
    box-shadow: 0 2px 10px rgba(0,0,0,0.18);
    transition: all 0.15s ease;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    gap: 8px !important;
}
#sdr-new-btn button:hover {
    background: #f1f5f9 !important;
    box-shadow: 0 6px 18px rgba(0,0,0,0.28);
    transform: translateY(-1px);
}

/* Channel list — pills with brighter colour for inactive items. */
#sdr-channels label {
    border-radius: 10px !important;
    padding: 9px 12px !important;
    font-size: 13.5px !important;
    color: #e2e8f0 !important;
    background: rgba(255,255,255,0.04) !important;
    border: 1px solid rgba(255,255,255,0.06) !important;
    transition: all 0.15s ease;
    margin: 3px 0 !important;
    cursor: pointer;
    display: block;
}
#sdr-channels label:hover {
    background: rgba(255,255,255,0.10) !important;
    border-color: rgba(255,255,255,0.15) !important;
    color: #f8fafc !important;
    transform: translateX(2px);
}
#sdr-channels label:has(input:checked) {
    background: linear-gradient(135deg, rgba(99,102,241,0.28), rgba(167,81,247,0.22)) !important;
    border-color: #818cf8 !important;
    color: #ffffff !important;
    box-shadow: 0 2px 10px rgba(99,102,241,0.18);
    font-weight: 600 !important;
}
#sdr-channels input[type="radio"] { display:none; }
#sdr-channels span { color: inherit !important; }

/* Sidebar footer. */
#sdr-footer {
    margin-top: auto;
    padding: 12px 8px 0 8px;
    border-top: 1px solid rgba(255,255,255,0.08);
    font-size: 11px;
    color: #94a3b8;
    line-height: 1.5;
}
#sdr-footer .sdr-user {
    display:flex; align-items:center; gap: 10px; padding:8px 4px;
    border-radius: 8px;
}
#sdr-footer .sdr-avatar {
    width: 28px; height: 28px; border-radius: 50%;
    background: linear-gradient(135deg,#6366f1,#ec4899);
    display:grid; place-items:center; font-size: 12px; color:#fff; font-weight:600;
}
#sdr-footer .sdr-org { color:#cbd5e1; font-weight:500; font-size:12px; }
#sdr-signout-btn { margin-top: 6px; }
#sdr-signout-btn button {
    background: rgba(239,68,68,0.10) !important;
    color: #fca5a5 !important;
    border: 1px solid rgba(239,68,68,0.30) !important;
    border-radius: 8px !important;
    font-size: 12px !important;
    padding: 7px 10px !important;
    width: 100% !important;
}
#sdr-signout-btn button:hover {
    background: rgba(239,68,68,0.18) !important;
    color: #fecaca !important;
}

/* ============ Main area — full height flex ============ */
#sdr-main {
    padding: 0 !important;
    background: #ffffff;
    display: flex !important;
    flex-direction: column !important;
    min-height: calc(100vh - 40px);
    overflow-y: auto;
}
#sdr-header {
    padding: 14px 28px 12px 28px;
    border-bottom: 1px solid #f1f5f9;
    background: rgba(255,255,255,0.9);
    backdrop-filter: saturate(180%) blur(10px);
    display:flex; align-items:center; justify-content:space-between; gap:14px;
    position: sticky; top: 0; z-index: 2;
}
#sdr-header .sdr-badge {
    background: linear-gradient(135deg,#eef2ff,#fdf4ff);
    color: #4338ca; font-size: 11px; font-weight: 600;
    padding: 3px 10px; border-radius: 999px;
    border: 1px solid #e0e7ff;
}
#sdr-header h2 { margin: 0; font-size: 17px; font-weight: 700; color:#0f172a; }
#sdr-header p { margin: 2px 0 0 0; font-size: 12.5px; color: #64748b; }

/* Welcome hero + 2×2 sample prompt grid. */
#sdr-welcome { padding: 18px 32px 4px 32px; flex: 0 0 auto; }
#sdr-welcome-hero { text-align:center; padding: 4px 0 14px 0; }
#sdr-welcome-hero .sdr-hero-icon { font-size: 36px; margin-bottom: 4px; }
#sdr-welcome-hero .sdr-hero-title { font-size: 20px; font-weight: 700; color:#0f172a; margin-bottom: 4px; }
#sdr-welcome-hero .sdr-hero-sub { color:#64748b; font-size: 13px; max-width: 580px; margin: 0 auto; line-height:1.5; }

#sdr-card-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 10px;
    max-width: 720px;
    margin: 6px auto 0 auto;
}
.sdr-card {
    background: #ffffff;
    border: 1px solid #e5e7eb;
    border-radius: 14px;
    padding: 14px 16px;
    cursor: pointer;
    box-shadow: 0 1px 2px rgba(15,23,42,0.03);
    transition: all 0.15s ease;
    display: flex;
    flex-direction: column;
    gap: 4px;
    user-select: none;
}
.sdr-card:hover {
    border-color: #a5b4fc;
    box-shadow: 0 6px 18px rgba(99,102,241,0.14), 0 2px 6px rgba(15,23,42,0.06);
    transform: translateY(-2px);
    background: linear-gradient(135deg,#fdfdff,#faf5ff);
}
.sdr-card-emoji { font-size: 22px; line-height: 1; }
.sdr-card-title { font-weight: 700; color:#0f172a; font-size: 14.5px; }
.sdr-card-desc { color:#64748b; font-size: 12.5px; line-height: 1.45; }

/* Chatbot — compact status area above the KG panel. */
#sdr-chat {
    border: none !important;
    background: transparent !important;
    padding: 8px 16px 0 16px !important;
    flex: 0 0 auto !important;
    overflow-y: auto !important;
}

/* KG panel — the main artifact. */
#sdr-kg-panel {
    flex: 1 1 auto !important;
    padding: 8px 20px 4px 20px !important;
    min-height: 340px !important;
    overflow: hidden;
}
#sdr-kg-panel > div { width: 100%; height: 100%; }
#sdr-kg-panel > div > div:first-child,
#sdr-kg-panel > div > div:last-child {
    width: 100%;
    height: 100%;
}
#sdr-chat .message { font-size: 15px !important; line-height: 1.68 !important; color:#1f2937 !important; }
#sdr-chat .bubble, #sdr-chat .user-row, #sdr-chat .bot-row { background: transparent !important; }
#sdr-chat .message-wrap { max-width: 820px !important; margin: 0 auto !important; }
/* User messages — subtle indigo-tinted bubble on the right. */
#sdr-chat .user > div:first-child {
    background: #f0f2ff !important;
    border-radius: 18px 18px 4px 18px !important;
    padding: 10px 14px !important;
    color: #1e1b4b !important;
}
/* Assistant messages — no bubble, pure text. */
#sdr-chat .bot > div:first-child { background: transparent !important; padding: 4px 2px !important; }

/* Input area — pinned near bottom. */
#sdr-input-row {
    max-width: 820px;
    margin: 8px auto 4px auto;
    background: #fff;
    border: 1px solid #e5e7eb;
    border-radius: 18px;
    padding: 6px 8px;
    box-shadow: 0 4px 16px rgba(15, 23, 42, 0.06), 0 1px 2px rgba(15, 23, 42, 0.04);
    transition: all 0.2s ease;
    flex: 0 0 auto !important;
}
#sdr-input-row:focus-within {
    border-color: #a5b4fc;
    box-shadow: 0 4px 20px rgba(99, 102, 241, 0.12), 0 0 0 3px rgba(99, 102, 241, 0.08);
}
#sdr-input-row textarea {
    border: none !important;
    background: transparent !important;
    font-size: 15px !important;
    padding: 10px 12px !important;
    resize: none !important;
    color: #111827 !important;
}
#sdr-input-row textarea::placeholder { color: #9ca3af !important; }

#sdr-send-btn button {
    background: linear-gradient(135deg,#4f46e5,#7c3aed) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 12px !important;
    min-width: 60px !important;
    height: 44px !important;
    padding: 0 14px !important;
    font-weight: 600 !important;
    box-shadow: 0 2px 8px rgba(79, 70, 229, 0.28);
    transition: all 0.15s ease;
}
#sdr-send-btn button:hover { transform: translateY(-1px); box-shadow: 0 4px 14px rgba(79, 70, 229, 0.4); }

/* Login screen. */
#sdr-login { max-width: 420px; margin: 0 auto; padding: 40px 24px; }
#sdr-login-card { background: #ffffff; padding: 24px; border-radius: 14px; box-shadow: 0 20px 50px rgba(2,6,23,0.35); border: 1px solid #e5e7eb; }
#sdr-login-card input { border-radius: 10px !important; }
#sdr-login-btn button {
    width: 100% !important; background: linear-gradient(135deg,#4f46e5,#7c3aed) !important;
    color: #fff !important; border: none !important; border-radius: 12px !important;
    padding: 12px 16px !important; font-weight: 600 !important; margin-top: 8px !important;
    box-shadow: 0 2px 10px rgba(79,70,229,0.28);
}
#sdr-login-btn button:hover { transform: translateY(-1px); box-shadow: 0 4px 16px rgba(79,70,229,0.4); }
#sdr-login-msg { text-align: center; }

#sdr-note {
    text-align: center;
    font-size: 10.5px;
    color: #b3b7c2;
    letter-spacing: 0.1px;
    padding: 4px 16px 16px 16px;
    max-width: 760px;
    margin: 0 auto;
    line-height: 1.45;
}
#sdr-note::before {
    content: "ⓘ  ";
    color: #cbd5e1;
    font-size: 11px;
}
"""

EMPTY_PLACEHOLDER = ""

SAMPLE_PROMPTS = [
    ("🔋", "Lithium-ion batteries",
     "Map tier-1 and tier-2 suppliers of cobalt, lithium and nickel used in the cathode of Tesla's 4680 battery cell, including country of origin, refining facilities, alternative materials and known ESG risks."),
    ("🧲", "EV rare-earth magnets",
     "Identify the mining locations, refining and separation facilities, and top manufacturers of NdFeB rare-earth magnets used in permanent-magnet synchronous motors for electric vehicles, including geopolitical risks and substitution efforts."),
    ("💊", "Pharmaceutical APIs",
     "Trace the tier-1 and tier-2 active pharmaceutical ingredient (API) suppliers for a representative generic cardiovascular drug, focusing on India and China manufacturers and regulatory/FDA inspection history."),
    ("🌾", "Cocoa supply chain",
     "Map the supply chain of cocoa for the global chocolate industry — origin countries, major cooperatives, traders, processors, and known deforestation or child-labour risks."),
]


with gr.Blocks(title="Deep Research — Supply Chain AI", elem_id="sdr-root", css=CSS) as demo:
    # BrowserState persists across page refreshes via browser localStorage.
    state = gr.BrowserState(default_value={})
    auth_state = gr.BrowserState(default_value=False)

    # ---------- LOGIN SCREEN (hidden until _restore_auth decides) ----------
    with gr.Column(visible=False, elem_id="sdr-login") as login_col:
        gr.HTML(
            "<div style='text-align:center;padding:60px 20px 24px;'>"
            "<div style='font-size:48px;text-shadow:0 2px 10px rgba(99,102,241,0.4);'>🔬</div>"
            "<h2 style='color:#f8fafc;margin:10px 0 6px;font-weight:700;letter-spacing:-0.3px;text-shadow:0 2px 12px rgba(0,0,0,0.3);'>Deep Research — Supply Chain AI</h2>"
            "<p style='color:#cbd5e1;margin:0 auto 28px;max-width:480px;line-height:1.55;'>"
            "Public beta from the <b style='color:#e0e7ff;'>Supply Chain AI Lab (SCAIL)</b>, University of Cambridge. "
            "Please sign in to continue."
            "</p>"
            "</div>"
        )
        with gr.Column(elem_id="sdr-login-card"):
            login_user = gr.Textbox(label="Username", placeholder="username", max_lines=1, elem_id="sdr-login-user")
            login_pass = gr.Textbox(label="Password", placeholder="password", type="password", max_lines=1, elem_id="sdr-login-pass")
            login_btn = gr.Button("Sign in", variant="primary", elem_id="sdr-login-btn")
            login_msg = gr.HTML("", elem_id="sdr-login-msg")

    # ---------- MAIN APP (hidden until login) ----------
    with gr.Row(equal_height=False, visible=False) as main_row:
        # ---------- Sidebar ----------
        with gr.Column(scale=1, min_width=260, elem_id="sdr-sidebar"):
            gr.HTML(
                "<div id='sdr-brand'>"
                "<div class='sdr-logo'>🔬</div>"
                "<div>"
                "<div class='sdr-title'>Deep Research</div>"
                "<div class='sdr-sub'>Supply Chain AI · SCAIL</div>"
                "</div>"
                "</div>"
            )
            with gr.Group(elem_id="sdr-new-btn"):
                new_btn = gr.Button("＋ New chat", variant="primary", size="md")
            gr.HTML("<div id='sdr-section-head'>Recent</div>")
            channel_list = gr.Radio(
                choices=[],
                value=None,
                label="",
                interactive=True,
                container=False,
                elem_id="sdr-channels",
            )
            gr.HTML(
                "<div id='sdr-footer'>"
                "<div class='sdr-user'>"
                "<div class='sdr-avatar'>S</div>"
                "<div>"
                "<div class='sdr-org'>SCAIL · Cambridge</div>"
                "<div>Public beta</div>"
                "</div>"
                "</div>"
                "</div>"
            )
            with gr.Group(elem_id="sdr-signout-btn"):
                signout_btn = gr.Button("↩ Sign out", size="sm")

        # ---------- Main chat area ----------
        with gr.Column(scale=4, elem_id="sdr-main"):
            gr.HTML(
                "<div id='sdr-header'>"
                "<div>"
                "<h2>Deep Research — Supply Chain AI</h2>"
                "<p>A public beta from the Supply Chain AI Lab (SCAIL), University of Cambridge.</p>"
                "</div>"
                "<span class='sdr-badge'>BETA</span>"
                "</div>"
            )

            # Welcome hero + 2×2 sample-prompt grid (hidden after first question).
            cards_html = "<div id='sdr-welcome-hero'>" \
                         "<div class='sdr-hero-icon'>🔬</div>" \
                         "<div class='sdr-hero-title'>Deep Research — Supply Chain AI</div>" \
                         "<div class='sdr-hero-sub'>Ask a detailed supply chain question. The agent will search the web, reason across evidence, and build an interactive knowledge graph.</div>" \
                         "</div><div id='sdr-card-grid'>"
            for i, (emoji, title, _) in enumerate(SAMPLE_PROMPTS):
                desc_short = SAMPLE_PROMPTS[i][2]
                desc_short = desc_short[:95] + "…" if len(desc_short) > 95 else desc_short
                cards_html += (
                    f"<div class='sdr-card' onclick=\"document.querySelector('#sdr-hidden-{i} button').click()\">"
                    f"<div class='sdr-card-emoji'>{emoji}</div>"
                    f"<div class='sdr-card-title'>{title}</div>"
                    f"<div class='sdr-card-desc'>{desc_short}</div>"
                    f"</div>"
                )
            cards_html += "</div>"
            welcome_block = gr.HTML(cards_html, elem_id="sdr-welcome")

            # Hidden buttons — JS triggers these to fill the textarea.
            sample_btns = []
            with gr.Row(visible=False):
                for i in range(len(SAMPLE_PROMPTS)):
                    sample_btns.append(gr.Button("", elem_id=f"sdr-hidden-{i}"))

            # Chat: text-only Q&A. KGs open in the /kg/ viewer via link.
            chat = gr.Chatbot(
                label="",
                elem_id="sdr-chat",
                height=520,
                show_label=False,
                render_markdown=True,
                sanitize_html=False,
                type="messages",
            )

            with gr.Row(elem_id="sdr-input-row"):
                question_box = gr.Textbox(
                    label="",
                    placeholder="Message Deep Research…  (e.g. map tier-1 & tier-2 suppliers of cobalt in Tesla 4680 cells)",
                    lines=2,
                    max_lines=8,
                    scale=9,
                    container=False,
                    show_label=False,
                )
                with gr.Column(scale=1, min_width=70, elem_id="sdr-send-btn"):
                    ask_btn = gr.Button("↑", variant="primary")

            footer_note = gr.HTML(
                "<div id='sdr-note'>"
                "Deep Research can make mistakes. Verify important findings against primary sources."
                " · <span style='color:#818cf8;font-weight:500;'>Public beta</span>"
                "</div>",
                visible=False,
            )

    # ---------- auth handlers ----------
    def do_login(u, p, auth_st):
        if _auth((u or "").strip(), (p or "").strip()):
            return (
                True,
                gr.update(visible=False),                                     # login_col
                gr.update(visible=True),                                      # main_row
                gr.update(visible=True),                                      # footer_note
                "",                                                           # login_msg
            )
        return (
            False,
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(visible=False),
            "<div style='color:#dc2626;margin-top:10px;font-size:13px;'>Invalid credentials. Please try again.</div>",
        )

    def do_signout(state):
        return (
            False,                              # auth_state cleared
            {},                                 # state wiped (channels + KGs cleared)
            gr.update(visible=True),            # login_col
            gr.update(visible=False),           # main_row
            gr.update(visible=False),           # footer_note
            "",                                 # login_msg
            "",                                 # username field cleared
            "",                                 # password field cleared
        )

    signout_btn.click(
        do_signout,
        inputs=[state],
        outputs=[auth_state, state, login_col, main_row, footer_note, login_msg, login_user, login_pass],
        api_name=False,
    )

    login_btn.click(
        do_login,
        inputs=[login_user, login_pass, auth_state],
        outputs=[auth_state, login_col, main_row, footer_note, login_msg],
        api_name=False,
    )
    login_pass.submit(
        do_login,
        inputs=[login_user, login_pass, auth_state],
        outputs=[auth_state, login_col, main_row, footer_note, login_msg],
        api_name=False,
    )

    # Single consolidated load handler — prevents flash by deciding both
    # (a) which panel (login vs chat) is visible and (b) initial chat state
    # in one round-trip.
    def _initial_load(state, auth_st):
        state = _hydrate(state)
        ch = _get_channel(state, state["active"])
        welcome_visible = bool(auth_st) and len(ch["messages"]) == 0
        return (
            state,
            gr.update(choices=_sidebar_choices(state), value=state["active"]),
            list(ch["messages"]),
            gr.update(visible=welcome_visible),
            gr.update(visible=not bool(auth_st)),  # login_col
            gr.update(visible=bool(auth_st)),      # main_row
            gr.update(visible=bool(auth_st)),      # footer_note
        )

    demo.load(
        _initial_load,
        inputs=[state, auth_state],
        outputs=[state, channel_list, chat, welcome_block, login_col, main_row, footer_note],
        api_name=False,
    )

    new_btn.click(
        on_new_channel,
        inputs=state,
        outputs=[state, channel_list, chat, question_box, welcome_block],
        api_name=False,
    )
    channel_list.change(
        on_switch_channel,
        inputs=[state, channel_list],
        outputs=[state, chat, welcome_block],
        api_name=False,
    )
    ask_btn.click(
        on_ask,
        inputs=[state, question_box],
        outputs=[state, chat, question_box, channel_list, welcome_block],
        api_name=False,
    )
    question_box.submit(
        on_ask,
        inputs=[state, question_box],
        outputs=[state, chat, question_box, channel_list, welcome_block],
        api_name=False,
    )

    # Hidden buttons triggered by the HTML card onClick — populate the textarea.
    for btn, (_emoji, _title, full_prompt) in zip(sample_btns, SAMPLE_PROMPTS):
        btn.click(lambda p=full_prompt: p, inputs=None, outputs=question_box, api_name=False)


demo.queue()


# ---------- login gate ------------------------------------------------------
# APP_USERS env var: "user1:pass1,user2:pass2" — configure via HF Space Secret.
def _parse_users(s: str) -> list[tuple[str, str]]:
    out: list[tuple[str, str]] = []
    for pair in (s or "").split(","):
        pair = pair.strip()
        if ":" in pair:
            u, p = pair.split(":", 1)
            if u and p:
                out.append((u.strip(), p.strip()))
    return out


USERS = _parse_users(os.getenv("APP_USERS", "scail:deepresearch2026"))


def _auth(username: str, password: str) -> bool:
    return any(u == username and p == password for u, p in USERS)


# FastAPI app mounts Gradio + serves the KG viewer at /kg on the same origin.
app = FastAPI()


@app.get("/kg", response_class=HTMLResponse)
def kg_viewer():
    return HTMLResponse(content=_KG_VIEWER_HTML)


app = gr.mount_gradio_app(app, demo, path="/", ssr_mode=False)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)