"""Mustalih Living — Gradio Space app (Gradio 6.13+)"""

import os, json, random, tempfile
from collections import defaultdict

import gradio as gr
import networkx as nx
from huggingface_hub import hf_hub_download
from pyvis.network import Network

DATASET_REPO = "FatimahEmadEldin/icaire-ai-glossary-enriched"

# -----------------------------------------------------------------------------
# Load data from the paired dataset (or local fallback for dev)
# -----------------------------------------------------------------------------
def load_data():
    try:
        path = hf_hub_download(
            repo_id=DATASET_REPO,
            filename="glossary_enriched.json",
            repo_type="dataset",
        )
    except Exception:
        # Fallback to local file if the dataset repo isn't reachable
        path = "glossary_enriched.json"
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


DATASET = load_data()
TERMS = DATASET["terms"]
TERMS_BY_NAME = {t["english_term"]: t for t in TERMS if t.get("english_term")}
ALL_TERMS = sorted(TERMS_BY_NAME.keys())


# -----------------------------------------------------------------------------
# Build the knowledge graph
# -----------------------------------------------------------------------------
def build_graph():
    G = nx.Graph()
    for t in TERMS:
        name = t.get("english_term")
        if not name:
            continue
        G.add_node(
            name,
            cluster=t.get("primary_cluster", "UNKNOWN"),
            arabic=t.get("arabic_term", ""),
            difficulty=t.get("difficulty", "intermediate"),
        )
    for t in TERMS:
        source = t.get("english_term")
        graph_raw = t.get("graph_raw", {}) or {}
        for edge_type in ["prerequisites", "unlocks", "related_concepts",
                          "used_with", "part_of"]:
            for target in graph_raw.get(edge_type, []):
                if target in TERMS_BY_NAME and target != source:
                    G.add_edge(source, target, type=edge_type)
    return G


G = build_graph()

# Clusters index
CLUSTERS = defaultdict(list)
for t in TERMS:
    c = t.get("primary_cluster", "UNKNOWN")
    if t.get("english_term"):
        CLUSTERS[c].append(t["english_term"])
CLUSTER_NAMES = sorted(CLUSTERS.keys())
CLUSTER_CHOICES = ["All"] + CLUSTER_NAMES

# Story tracks index
TRACK_TERMS = defaultdict(list)
for t in TERMS:
    story = t.get("story_assignments_v2", {}) or {}
    for assignment in story.get("story_assignments", []):
        track_id = assignment.get("track")
        if track_id:
            TRACK_TERMS[track_id].append({
                "term": t["english_term"],
                "arabic": t.get("arabic_term", ""),
                "position": assignment.get("position_in_track", 5),
                "role": assignment.get("role", "supporting"),
                "hook_ar": assignment.get("one_line_hook_ar", ""),
                "hook_en": assignment.get("one_line_hook_en", ""),
                "chapter": assignment.get("chapter_hint", ""),
            })
for track in TRACK_TERMS:
    TRACK_TERMS[track].sort(key=lambda x: x["position"])

TRACK_DISPLAY = {
    "DATA_FOUNDATIONS": "Data Foundations / أساسيات البيانات",
    "HOW_A_MODEL_LEARNS": "How a Model Learns / كيف يتعلم النموذج",
    "NEURAL_NETWORKS_AND_DEEP": "Neural Networks / الشبكات العصبية",
    "CLASSICAL_ML_AND_STATS": "Classical ML / تعلم الآلة التقليدي",
    "APPLIED_AI": "Applied AI / الذكاء التطبيقي",
    "TRUSTWORTHY_AI": "Trustworthy AI / الذكاء الموثوق",
    "AI_INFRASTRUCTURE": "AI Infrastructure / البنية التحتية",
}
TRACK_CHOICES = [k for k in TRACK_DISPLAY.keys() if k in TRACK_TERMS]


# -----------------------------------------------------------------------------
# Mermaid rendering — inject the Mermaid JS library
# (Gradio 6's Markdown supports ```mermaid``` blocks but requires the JS)
# -----------------------------------------------------------------------------
MERMAID_JS = """
<script type="module">
  import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.esm.min.mjs';
  mermaid.initialize({ startOnLoad: false, theme: 'neutral' });
  function renderAll() {
    document.querySelectorAll('pre code.language-mermaid, code.language-mermaid').forEach((el, i) => {
      if (el.dataset.processed) return;
      const container = document.createElement('div');
      container.className = 'mermaid';
      container.innerHTML = el.textContent;
      el.parentElement.replaceWith(container);
      el.dataset.processed = '1';
    });
    mermaid.run();
  }
  new MutationObserver(() => setTimeout(renderAll, 100))
    .observe(document.body, { childList: true, subtree: true });
  setTimeout(renderAll, 500);
</script>
"""


# -----------------------------------------------------------------------------
# View helpers
# -----------------------------------------------------------------------------
def render_term_detail(term_name):
    if not term_name or term_name not in TERMS_BY_NAME:
        return "Select a term…", ""
    t = TERMS_BY_NAME[term_name]
    feel = t.get("one_sentence_feel", {}) or {}
    detailed = t.get("detailed_explanation", {}) or {}
    graph_raw = t.get("graph_raw", {}) or {}

    md = f"""## {t.get('arabic_term', '')}
### {t.get('english_term', '')}

**Cluster:** `{t.get('primary_cluster', '')}` · **Difficulty:** {t.get('difficulty', '')}

---

### Feel / الإحساس
> {feel.get('ar', '')}
>
> *{feel.get('en', '')}*

---

### Definition / التعريف
{t.get('arabic_def', '')}

{t.get('english_def', '')}

---

### Detailed explanation
{detailed.get('ar', '')}

{detailed.get('en', '')}

---

### Relationships
"""
    for edge_type in ["prerequisites", "unlocks", "part_of",
                      "alternative_to", "contrasts_with", "related_concepts"]:
        items = graph_raw.get(edge_type, [])
        if items:
            md += f"- **{edge_type.replace('_', ' ').title()}:** "
            md += ", ".join(f"`{i}`" for i in items) + "\n"

    if t.get("code_example_python"):
        md += f"\n---\n\n### Code example\n```python\n{t['code_example_python']}\n```\n"

    mermaid = t.get("ai_mermaid", "")
    mermaid_md = f"```mermaid\n{mermaid}\n```" if mermaid and mermaid.strip() else ""
    return md, mermaid_md


def list_cluster_terms(cluster_id):
    if not cluster_id:
        return ""
    terms = CLUSTERS.get(cluster_id, [])
    if not terms:
        return "No terms in this cluster."
    md = f"### {cluster_id} — {len(terms)} terms\n\n"
    for name in sorted(terms):
        t = TERMS_BY_NAME.get(name, {})
        ar = t.get("arabic_term", "")
        diff = t.get("difficulty", "")
        md += f"- **{name}** — {ar} · *{diff}*\n"
    return md


def render_track(track_id):
    if not track_id:
        return ""
    items = TRACK_TERMS.get(track_id, [])
    title = TRACK_DISPLAY.get(track_id, track_id)
    md = f"## {title}\n\n{len(items)} terms in this track.\n\n---\n\n"
    prev_chapter = None
    for i, item in enumerate(items, 1):
        chapter = item["chapter"] or "general"
        if chapter != prev_chapter:
            md += f"\n### Chapter: {chapter}\n"
            prev_chapter = chapter
        md += f"\n**{i}. {item['term']}** — {item['arabic']} · _{item['role']}_\n\n"
        if item["hook_ar"]:
            md += f"> {item['hook_ar']}\n\n"
        if item["hook_en"]:
            md += f"> *{item['hook_en']}*\n\n"
    return md


def render_graph_html(cluster_filter):
    net = Network(height="600px", width="100%", bgcolor="#fafaf7",
                  font_color="#1a1a19", notebook=False, cdn_resources="in_line")
    net.barnes_hut(spring_length=200)

    if cluster_filter and cluster_filter != "All":
        node_set = set(CLUSTERS.get(cluster_filter, []))
    else:
        degrees = dict(G.degree())
        node_set = set(sorted(degrees, key=lambda n: -degrees[n])[:60])

    color_map = {
        "TRANSFORMER_ARCHITECTURE": "#7F77DD",
        "NEURAL_NETWORK_BASICS": "#7F77DD",
        "LAYERS_AND_ACTIVATIONS": "#7F77DD",
        "OPTIMIZATION_ALGORITHMS": "#BA7517",
        "LOSS_FUNCTIONS": "#BA7517",
        "DATA_COLLECTION_AND_LABELING": "#1D9E75",
        "DATA_QUALITY_AND_CLEANING": "#1D9E75",
        "AI_ETHICS_PRINCIPLES": "#D85A30",
        "BIAS_AND_FAIRNESS_TYPES": "#D85A30",
        "FAIRNESS_METRICS": "#D85A30",
        "PROMPTING_TECHNIQUES": "#378ADD",
        "LARGE_LANGUAGE_MODELS": "#378ADD",
    }

    for n in node_set:
        attrs = G.nodes[n]
        color = color_map.get(attrs.get("cluster"), "#888780")
        label = f"{attrs.get('arabic', '')}\n{n}"
        net.add_node(n, label=label, color=color,
                     title=f"{n}\n{attrs.get('cluster', '')}", size=15)

    for src, tgt, data in G.edges(data=True):
        if src in node_set and tgt in node_set:
            net.add_edge(src, tgt, color="#cfccc1", width=0.6)

    with tempfile.NamedTemporaryFile("w", suffix=".html", delete=False) as f:
        net.save_graph(f.name)
        with open(f.name, "r", encoding="utf-8") as rf:
            return rf.read()


def search_terms(query):
    if not query or len(query.strip()) < 2:
        return "Type at least 2 characters."
    q = query.strip().lower()
    hits = []
    for t in TERMS:
        if (q in t.get("english_term", "").lower()
            or q in t.get("arabic_term", "")
            or q in t.get("arabic_def", "").lower()
            or q in t.get("english_def", "").lower()):
            hits.append(t)
            if len(hits) >= 20:
                break
    if not hits:
        return f"No matches for **{query}**."
    md = f"### {len(hits)} matches for *{query}*\n\n"
    for t in hits:
        md += f"- **{t['english_term']}** — {t.get('arabic_term', '')}\n"
    return md


# -----------------------------------------------------------------------------
# Build the Gradio interface (Gradio 6 idioms)
# -----------------------------------------------------------------------------
default_term = "Gradient Descent" if "Gradient Descent" in TERMS_BY_NAME else ALL_TERMS[0]

with gr.Blocks(
    title="Mustalih Living — مُصطلِح الحيّ",
    theme=gr.themes.Soft(primary_hue="slate"),
    head=MERMAID_JS,
) as demo:
    gr.Markdown(
        f"# Mustalih Living · مُصطلِح الحيّ\n"
        f"**Arabic-first interactive AI glossary** · {len(TERMS)} terms · "
        f"7 story tracks · {G.number_of_edges()} typed relationships"
    )

    with gr.Tabs():

        with gr.Tab("📖 Term detail"):
            with gr.Row():
                term_picker = gr.Dropdown(
                    choices=ALL_TERMS, value=default_term,
                    label="Pick a term", allow_custom_value=False,
                )
                random_btn = gr.Button("🎲 Random term")
            term_md = gr.Markdown()
            term_mermaid = gr.Markdown()

            term_picker.change(
                render_term_detail, inputs=term_picker,
                outputs=[term_md, term_mermaid],
            )
            random_btn.click(
                lambda: random.choice(ALL_TERMS),
                outputs=term_picker,
            )

        with gr.Tab("🕸 Knowledge graph"):
            cluster_picker = gr.Dropdown(
                choices=CLUSTER_CHOICES, value="All",
                label="Filter by cluster",
            )
            graph_html = gr.HTML(render_graph_html("All"))
            cluster_picker.change(
                render_graph_html, inputs=cluster_picker,
                outputs=graph_html,
            )

        with gr.Tab("📚 Story tracks"):
            track_picker = gr.Dropdown(
                choices=TRACK_CHOICES,
                value=TRACK_CHOICES[0] if TRACK_CHOICES else None,
                label="Pick a story track",
            )
            track_md = gr.Markdown(
                render_track(TRACK_CHOICES[0]) if TRACK_CHOICES else ""
            )
            track_picker.change(
                render_track, inputs=track_picker,
                outputs=track_md,
            )

        with gr.Tab("🔖 Clusters"):
            cluster_browse = gr.Dropdown(
                choices=CLUSTER_NAMES,
                value=CLUSTER_NAMES[0] if CLUSTER_NAMES else None,
                label="Explore a cluster",
            )
            cluster_md = gr.Markdown(
                list_cluster_terms(CLUSTER_NAMES[0]) if CLUSTER_NAMES else ""
            )
            cluster_browse.change(
                list_cluster_terms, inputs=cluster_browse,
                outputs=cluster_md,
            )

        with gr.Tab("🔍 Search"):
            query = gr.Textbox(
                label="Search in Arabic or English",
                placeholder="e.g., Attention, الانتباه, gradient…",
            )
            search_btn = gr.Button("Search")
            search_md = gr.Markdown()
            search_btn.click(search_terms, inputs=query, outputs=search_md)
            query.submit(search_terms, inputs=query, outputs=search_md)

        with gr.Tab("📊 Stats"):
            has_feel = sum(1 for t in TERMS if t.get("one_sentence_feel"))
            has_mermaid = sum(1 for t in TERMS if t.get("ai_mermaid"))
            has_story = sum(1 for t in TERMS if t.get("story_assignments_v2"))
            has_graph = sum(1 for t in TERMS if t.get("graph_raw"))
            gr.Markdown(f"""
### Dataset coverage

| Field | Coverage |
|---|---|
| Total terms | {len(TERMS)} |
| With feel metaphor | {has_feel} ({100*has_feel/len(TERMS):.1f}%) |
| With Mermaid UML | {has_mermaid} ({100*has_mermaid/len(TERMS):.1f}%) |
| With story assignment | {has_story} ({100*has_story/len(TERMS):.1f}%) |
| With graph edges | {has_graph} ({100*has_graph/len(TERMS):.1f}%) |

### Clusters
Found {len(CLUSTER_NAMES)} clusters across the corpus.
""")

    # Initial render on page load (Gradio 6 idiom: demo.load inside blocks)
    demo.load(
        fn=lambda: render_term_detail(default_term),
        outputs=[term_md, term_mermaid],
    )


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