import json
import os
import re
from pathlib import Path

import plotly.graph_objects as go
import gradio as gr
from huggingface_hub import InferenceClient

DATA_FILE = Path(__file__).parent / "data.json"

# Feature switch: when False, the login modal never shows and the Pro CTA never
# fires (regardless of the user's auth state). Flip to True to re-enable.
LOGIN_ENABLED = False

# Local dev: `export TRANSFORMERS_SERVE_URL=http://localhost:8000/v1`
# On HF Spaces: leave unset and set the HF_TOKEN secret instead.
_LOCAL_URL = os.environ.get("TRANSFORMERS_SERVE_URL")
MODEL = os.environ.get("CHAT_MODEL", "Qwen/Qwen2.5-Coder-7B-Instruct")

client = (
    InferenceClient(base_url=_LOCAL_URL, api_key="dummy")
    if _LOCAL_URL
    else InferenceClient()
)

SYSTEM = (
    "You control a 3D visualization of files in the huggingface/transformers "
    "repository. Each dot is one file. You do NOT have access to file "
    "contents — never paste code or pretend to show file contents.\n"
    "\n"
    "If the user is asking to focus the view on a specific file (verbs like "
    "zoom, focus, show, see, look, open, find, where), reply with ONE short "
    "sentence, then a final line containing exactly:\n"
    "[[ZOOM: <path>]]\n"
    "\n"
    "STRICT rules for <path>:\n"
    "- It MUST be one of the candidate paths provided for this turn, copied "
    "verbatim.\n"
    "- Never invent or modify a path. If no candidate fits, say so plainly "
    "and do NOT emit the directive.\n"
    "\n"
    "If the user is making small talk or asking a general question, answer "
    "briefly and do NOT emit the directive."
)

_ZOOM_RE = re.compile(r"\[\[ZOOM:\s*([^\]]+?)\s*\]\]")

# Tokens that appear in every (or nearly every) path or are conversational
# filler — useless for narrowing down candidates.
_STOP = {
    "the", "a", "an", "and", "or", "but", "of", "to", "in", "on", "for",
    "with", "this", "that", "is", "are", "be", "by", "at", "it", "its",
    "you", "your", "me", "my", "we", "our", "i", "can", "could", "would",
    "should", "please", "show", "find", "open", "view", "look", "see",
    "zoom", "focus", "navigate", "point", "file", "files", "code", "where",
    "what", "which", "how", "src", "transformers", "py",
}


def _candidate_paths(message, k=15):
    """Top-k file paths whose path tokens overlap the user's message."""
    tokens = [t for t in re.findall(r"[a-z0-9_]+", message.lower())
              if len(t) >= 3 and t not in _STOP]
    if not tokens:
        return []
    scored = []
    for path in _PATH_TO_IDX:
        p = path.lower()
        score = sum(1 for t in tokens if t in p)
        if score:
            scored.append((-score, len(path), path))
    scored.sort()
    return [path for _, _, path in scored[:k]]


def load_data():
    return json.loads(DATA_FILE.read_text())


_DATA = load_data()
# Map repo-relative file path -> point index. Hover strings look like
# "src/.../file.py<br>edits: N (last: YYYY-MM-DD)" — keep just the path.
_PATH_TO_IDX = {h.split("<br>", 1)[0]: i for i, h in enumerate(_DATA["hover"])}


def _find_point(target):
    """Resolve a path the model produced to an index into the point cloud."""
    target = target.strip().strip("`'\"")
    if target in _PATH_TO_IDX:
        return _PATH_TO_IDX[target]
    matches = [p for p in _PATH_TO_IDX if p.endswith(target) or target in p]
    if not matches:
        return None
    return _PATH_TO_IDX[min(matches, key=len)]


def _zoomed_fig(target, padding=0.4):
    idx = _find_point(target)
    if idx is None:
        return None
    x0, y0, z0 = _DATA["x"][idx], _DATA["y"][idx], _DATA["z"][idx]
    fig = make_point_cloud()
    fig.update_layout(scene=dict(
        xaxis=dict(range=[x0 - padding, x0 + padding]),
        yaxis=dict(range=[y0 - padding, y0 + padding]),
        zaxis=dict(range=[z0 - padding, z0 + padding]),
    ))
    return fig


def make_point_cloud():
    d = _DATA
    fig = go.Figure(
        data=[
            go.Scatter3d(
                x=d["x"],
                y=d["y"],
                z=d["z"],
                mode="markers",
                marker=dict(
                    size=3,
                    color=d["color"],
                    colorscale=[
                        (0.0, "#ff3b30"),
                        (0.5, "#ff9500"),
                        (1.0, "#34c759"),
                    ],
                    cmin=0.0,
                    cmax=1.0,
                    showscale=False,
                ),
                text=d["hover"],
                hovertemplate="%{text}<extra></extra>",
            )
        ]
    )
    fig.update_layout(
        template="plotly_dark",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        margin=dict(l=0, r=0, t=0, b=0),
        scene=dict(
            xaxis_title="",
            yaxis_title="",
            zaxis_title="",
            bgcolor="rgba(0,0,0,0)",
        ),
    )
    return fig


def respond(message, history):
    sys_prompt = SYSTEM
    candidates = _candidate_paths(message)
    if candidates:
        sys_prompt += (
            "\n\nCandidate files for this turn (pick one of these exact paths "
            "if you emit a ZOOM directive):\n"
            + "\n".join(f"- {p}" for p in candidates)
        )
    messages = [{"role": "system", "content": sys_prompt}, *history,
                {"role": "user", "content": message}]
    new_history = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": ""},
    ]
    full = ""
    stream = client.chat_completion(
        messages, model=MODEL, max_tokens=512, stream=True,
    )
    for chunk in stream:
        delta = chunk.choices[0].delta.content or ""
        full += delta
        new_history[-1]["content"] = _ZOOM_RE.sub("", full).rstrip()
        yield "", new_history, gr.skip()

    m = _ZOOM_RE.search(full)
    if m:
        fig = _zoomed_fig(m.group(1))
        if fig is not None:
            yield "", new_history, fig


CSS = """
html, body { margin: 0 !important; padding: 0 !important; height: 100vh !important; max-height: 100vh !important; overflow: hidden !important; }
gradio-app { display: block !important; height: 100% !important; }
footer { display: none !important; }

/* Header overlay — sits above the main split, doesn't take vertical space. */
#header { position: absolute !important; top: 0 !important; left: 0 !important; right: 0 !important;
          z-index: 10 !important; display: flex !important; align-items: center !important;
          justify-content: space-between !important; padding: 8px 12px !important;
          gap: 8px !important; pointer-events: none !important; }
#header > * { pointer-events: auto !important; }
#logo, #logo img { height: 40px !important; width: auto !important; max-width: 40px !important;
                   object-fit: contain !important; }
#logo { background: transparent !important; border: 0 !important; padding: 0 !important; flex: 0 0 auto !important; }
#pro-cta a { display: inline-block; padding: 6px 12px; border-radius: 999px;
             background: linear-gradient(90deg,#ff9500,#34c759); color: #111 !important;
             font-weight: 600; font-size: 13px; text-decoration: none; }
#pro-cta { flex: 0 0 auto !important; }

/* Login modal — full-screen overlay until dismissed or auth completes. */
#login-modal { position: fixed !important; inset: 0 !important; z-index: 100 !important;
               display: flex !important; align-items: center !important; justify-content: center !important;
               background: rgba(0,0,0,0.6) !important; backdrop-filter: blur(4px) !important;
               padding: 24px !important; margin: 0 !important; }
#login-modal .gr-group, #login-modal > div { background: #1a1a1a !important; border-radius: 12px !important;
               padding: 24px !important; max-width: 360px !important; text-align: center !important; }
#login-modal h2 { margin: 0 0 8px 0; font-size: 18px; }
#login-modal p { margin: 0 0 16px 0; opacity: 0.8; font-size: 14px; }

.gradio-container { height: 100% !important; max-width: 100% !important; width: 100% !important; margin: 0 !important; padding: 0 !important; overflow: hidden !important; min-height: 0 !important; }
.gradio-container .main, .gradio-container .wrap, .gradio-container .contain { height: 100% !important; padding: 0 !important; margin: 0 !important; max-width: 100% !important; min-height: 0 !important; }
#main-row { height: 100% !important; gap: 0 !important; margin: 0 !important; padding: 0 !important; flex-wrap: nowrap !important; min-height: 0 !important; }
#left-col, #right-col { height: 100% !important; padding: 0 !important; margin: 0 !important; min-width: 0 !important; min-height: 0 !important; }
#left-col > *, #right-col > * { border-radius: 0 !important; }
#point-plot, #point-plot > div, #point-plot .js-plotly-plot, #point-plot .plot-container { height: 100% !important; width: 100% !important; }
#right-col { display: flex !important; flex-direction: column !important; }
#chatbot { flex: 1 1 auto !important; height: auto !important; min-height: 0 !important; border-radius: 0 !important; }
#msg-input { flex: 0 0 auto !important; margin: 0 !important; border-radius: 0 !important; }
#msg-input textarea { min-height: 0 !important; }
#style-shim { display: none !important; }
"""

# Client-side double-click-to-focus. Plotly's own `plotly_doubleclick` fires
# only for the empty canvas (and is bound to "reset view"), so detect two
# clicks on the same point inside a short window instead. Pure JS — no server
# round-trip, no rebuild of the figure.
FOCUS_JS = """
<script>
(function() {
  var MARK = '__dblclick_focus_attached__';
  var GAP_MS = 400;
  var ZOOM_RADIUS = 0.4;
  function attach(el) {
    if (el[MARK]) return;
    el[MARK] = true;
    var last = {x: null, y: null, z: null, t: 0};
    el.on('plotly_click', function(ev) {
      if (!ev.points || !ev.points.length) return;
      var pt = ev.points[0];
      var now = Date.now();
      var same = last.t && (now - last.t) < GAP_MS
              && last.x === pt.x && last.y === pt.y && last.z === pt.z;
      if (same) {
        window.Plotly.relayout(el, {
          'scene.xaxis.range': [pt.x - ZOOM_RADIUS, pt.x + ZOOM_RADIUS],
          'scene.yaxis.range': [pt.y - ZOOM_RADIUS, pt.y + ZOOM_RADIUS],
          'scene.zaxis.range': [pt.z - ZOOM_RADIUS, pt.z + ZOOM_RADIUS],
        });
        last = {t: 0};
      } else {
        last = {x: pt.x, y: pt.y, z: pt.z, t: now};
      }
    });
  }
  function poll() {
    var el = document.querySelector('#point-plot .js-plotly-plot');
    if (el && el.on && window.Plotly) attach(el);
  }
  // Poll instead of one-shot: Gradio re-creates the plot element when the
  // figure is replaced (e.g. the LLM zoom path), and we need to re-bind then.
  setInterval(poll, 500);
  poll();
})();
</script>
"""

_THEME_CSS = gr.themes.Citrus()._get_theme_css()

with gr.Blocks(title="File Point Cloud", fill_height=True) as demo:
    # HF Spaces serves `demo` directly without calling launch(), so CSS/theme
    # passed to launch() never runs there. Injecting them via gr.HTML works in both.
    gr.HTML(
        f"<style>{_THEME_CSS}{CSS}</style>{FOCUS_JS}",
        elem_id="style-shim", container=False, padding=False,
    )

    # Login modal — only mounts/shows when LOGIN_ENABLED. on_load() further hides
    # it if the user is already authenticated. The LoginButton is created
    # conditionally because instantiating any OAuth component triggers Gradio
    # to attach OAuth routes at startup, which crashes on a Space that hasn't
    # opted into OAuth via the README `hf_oauth: true` flag.
    with gr.Group(elem_id="login-modal", visible=LOGIN_ENABLED) as login_modal:
        gr.HTML(
            "<h2>Welcome</h2>"
            "<p>Log in with your Hugging Face account to continue.</p>"
        )
        if LOGIN_ENABLED:
            gr.LoginButton()
        skip_btn = gr.Button("Continue without logging in", variant="secondary", size="sm")
    skip_btn.click(lambda: gr.update(visible=False), outputs=[login_modal])

    # Header overlay: logo (top-left) + Pro CTA (top-right). Absolutely
    # positioned via CSS so it doesn't interfere with the 50/50 split below.
    with gr.Row(elem_id="header"):
        gr.Image(
            value="logo.png", elem_id="logo", show_label=False, container=False,
            interactive=False, buttons=[],
        )
        # Pro CTA — hidden by default, shown on .load() if the user is logged in
        # but lacks Pro. For the placeholder it just links out to the upgrade page.
        pro_cta = gr.HTML(
            '<div><a href="https://huggingface.co/subscribe/pro" target="_blank" '
            'rel="noopener">Upgrade to Pro →</a></div>',
            elem_id="pro-cta", visible=False,
        )

    # Two signatures because the OAuthProfile annotation also pulls in OAuth
    # wiring; only use it when LOGIN_ENABLED.
    if LOGIN_ENABLED:
        def on_load(profile: gr.OAuthProfile | None):
            logged_in = profile is not None
            # TODO(pro): hit the HF API with the OAuth token to read the user's
            # plan; for now we always offer the CTA to logged-in users.
            is_pro = False
            return (
                gr.update(visible=not logged_in),
                gr.update(visible=logged_in and not is_pro),
            )
    else:
        def on_load():
            return gr.update(visible=False), gr.update(visible=False)

    demo.load(on_load, inputs=None, outputs=[login_modal, pro_cta])

    with gr.Row(elem_id="main-row", equal_height=True):
        with gr.Column(scale=1, elem_id="left-col", min_width=0):
            plot = gr.Plot(value=make_point_cloud(), show_label=False, elem_id="point-plot")
        with gr.Column(scale=1, elem_id="right-col", min_width=0):
            chatbot = gr.Chatbot(elem_id="chatbot", show_label=False)
            msg = gr.Textbox(
                placeholder="Ask about the files…",
                show_label=False,
                elem_id="msg-input",
                lines=1,
                max_lines=8,
                container=False,
            )
            msg.submit(respond, [msg, chatbot], [msg, chatbot, plot])

if __name__ == "__main__":
    demo.launch()