#!/usr/bin/env python3
"""
generate_research_stack_usage_map.py — Build a searchable point graph of Research Stack.

Produces:
  - research_stack_usage_graph.json   (entities + edges)
  - research_stack_usage_graph.dot    (visual graph)
  - research_stack_usage_graph.md     (human-readable summary)

Entity kinds: module, theorem, def, script, doc, service, database, node, skill, goal
Edge kinds:  imports, contains, uses_service, touches_database, references_doc, implements, depends_on
"""
import json
import re
from collections import Counter, defaultdict
from pathlib import Path

ROOT = Path("/home/allaun/Research Stack")
OUT_DIR = Path("/tmp/SilverSight/docs")

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

def first_docstring(text: str) -> str:
    lines = text.splitlines()
    for i, line in enumerate(lines):
        s = line.strip()
        if s.startswith("/-"):
            buf = []
            for j in range(i + 1, len(lines)):
                l = lines[j].strip()
                if l.endswith("-/"):
                    buf.append(l[:-2].strip())
                    break
                buf.append(l.lstrip("- ").strip())
            return " ".join(buf).strip()
        if s.startswith("-- ") or s.startswith("# "):
            return s[2:].strip()
    return ""

def first_heading(text: str) -> str:
    for line in text.splitlines():
        s = line.strip()
        if s.startswith("# "):
            return s[2:].strip()
        if s.startswith("## "):
            return s[3:].strip()
    return ""

MATH_KINDS = {
    "braid":        ["braid","crossing","slug3","ternary","anyon","fusion","topo"],
    "number_theory":["sidon","erdos","prime","diophantine","goormaghtigh","zeckendorf"],
    "geometry":     ["manifold","geodesic","curvature","christoffel","riemannian","spherion"],
    "algebra":      ["quaternion","algebra","ring","field","group","monoid","category"],
    "analysis":     ["convergence","continuity","lipschitz","cauchy","measure","integral"],
    "quantum":      ["yangmills","lattice","hamiltonian","energy","entropy","quantum"],
    "fixedpoint":   ["fix16","q16","fixedpoint","saturate","phasemodulus"],
    "routing":      ["route","cfd","navier","burgers","canal","flow","pressure"],
    "avm":          ["avmisa","avm","instruction","receipt","emit"],
    "rrc":          ["rrc","corpus250","pist","classif","receipt"],
}

def dominant_kind(text: str) -> str:
    text_l = text.lower()
    scores = {k: sum(text_l.count(w) for w in ws) for k, ws in MATH_KINDS.items()}
    best = max(scores, key=scores.get)
    return best if scores[best] > 0 else "general"

# ── Entity containers ─────────────────────────────────────────────────────────

entities: dict[str, dict] = {}
edges: list[dict] = []

def add_entity(entity_kind: str, id: str, name: str, **props):
    if id in entities:
        entities[id].update(props)
        return
    entities[id] = {"kind": entity_kind, "id": id, "name": name, **props}

def add_edge(src: str, dst: str, kind: str, **props):
    edges.append({"src": src, "dst": dst, "kind": kind, **props})

# ── 1. Lean modules ───────────────────────────────────────────────────────────

LEAN_DIR = ROOT / "0-Core-Formalism/lean/Semantics"
SEMANTICS_DIR = LEAN_DIR / "Semantics"

RE_IMPORT = re.compile(r'^import\s+(\S+)')
RE_THEOREM = re.compile(r'^\s*(?:private\s+|protected\s+)?(theorem|lemma|def|structure|inductive)\s+(\w+)', re.MULTILINE)
RE_NAMESPACE = re.compile(r'^namespace\s+(\S+)', re.MULTILINE)

module_exports: dict[str, list[str]] = defaultdict(list)

for lean in sorted(SEMANTICS_DIR.rglob("*.lean")):
    rel = lean.relative_to(LEAN_DIR)
    mod_id = ".".join(rel.with_suffix("").parts)
    text = lean.read_text(encoding="utf-8", errors="replace")

    imports = [m.group(1) for m in RE_IMPORT.finditer(text)]
    decls = RE_THEOREM.findall(text)
    ns_match = RE_NAMESPACE.search(text)
    namespace = ns_match.group(1) if ns_match else ""

    theorem_names = [n for k, n in decls if k in ("theorem", "lemma")]
    def_names = [n for k, n in decls if k == "def"]
    struct_names = [n for k, n in decls if k == "structure"]
    ind_names = [n for k, n in decls if k == "inductive"]

    kind = dominant_kind(text)
    doc = first_docstring(text)

    add_entity(
        "module", mod_id, mod_id.split(".")[-1],
        path=str(rel),
        namespace=namespace,
        doc=doc[:300],
        math_kind=kind,
        sorry_count=text.count("sorry"),
        theorem_count=len(theorem_names),
        def_count=len(def_names),
        struct_count=len(struct_names),
        inductive_count=len(ind_names),
        line_count=text.count("\n"),
    )

    # Containment edges to declarations
    for name in theorem_names[:30]:  # limit to keep graph manageable
        decl_id = f"{mod_id}.{name}"
        add_entity("theorem", decl_id, name, module=mod_id)
        add_edge(mod_id, decl_id, "contains")
    for name in def_names[:20]:
        decl_id = f"{mod_id}.{name}"
        add_entity("def", decl_id, name, module=mod_id)
        add_edge(mod_id, decl_id, "contains")

    # Import edges
    for imp in imports:
        if imp.startswith("Semantics."):
            add_edge(mod_id, imp, "imports")
        elif imp.startswith("Mathlib"):
            add_entity("external_lib", imp, imp, kind="mathlib")
            add_edge(mod_id, imp, "imports")

# ── 2. Scripts ────────────────────────────────────────────────────────────────

SCRIPT_DIRS = [
    ROOT / "scripts",
    ROOT / "4-Infrastructure" / "shim",
]

SERVICE_HINTS = {
    "neon-64gb": "neon_server",
    "100.92.88.64": "neon_server",
    "arxiv-pg": "neon_postgres",
    "appflowy": "appflowy_db",
    "vikunja": "vikunja_db",
    "ene": "ene_db",
    "ollama": "ollama_service",
    "gremlin": "gremlin_graph",
    "mathblob": "gremlin_graph",
    "headroom": "headroom_proxy",
    "100.88.57.96": "qfox_hermes",
    "hermes": "hermes_dashboard",
}

DB_HINTS = ["arxiv", "ene", "appflowy", "vikunja"]

for sdir in SCRIPT_DIRS:
    if not sdir.exists():
        continue
    for script in sorted(sdir.rglob("*.py")):
        rel = script.relative_to(ROOT)
        sid = f"script:{rel}"
        text = script.read_text(encoding="utf-8", errors="replace")
        doc = first_docstring(text)
        add_entity("script", sid, script.name, path=str(rel), doc=doc[:300])
        # Detect service/db usage
        text_l = text.lower()
        for hint, svc_id in SERVICE_HINTS.items():
            if hint.lower() in text_l:
                add_edge(sid, svc_id, "uses_service")
        for db in DB_HINTS:
            if db in text_l:
                add_edge(sid, f"db:{db}", "touches_database")

# ── 3. Docs / specs ───────────────────────────────────────────────────────────

DOCS_DIR = ROOT / "6-Documentation"
for doc_file in sorted(DOCS_DIR.rglob("*.md")):
    rel = doc_file.relative_to(ROOT)
    did = f"doc:{rel}"
    text = doc_file.read_text(encoding="utf-8", errors="replace")
    title = first_heading(text) or doc_file.name
    add_entity("doc", did, title, path=str(rel))

# ── 4. Skills ─────────────────────────────────────────────────────────────────

SKILL_DIRS = [
    ROOT / ".github" / "skills",
    ROOT / ".devin" / "skills",
]
for skdir in SKILL_DIRS:
    if not skdir.exists():
        continue
    for skill_dir in sorted(skdir.iterdir()):
        skill_md = skill_dir / "SKILL.md"
        if skill_md.exists():
            text = skill_md.read_text(encoding="utf-8", errors="replace")
            title = first_heading(text) or skill_dir.name
            add_entity("skill", f"skill:{skill_dir.name}", skill_dir.name, doc=title[:300])

# ── 5. Services / databases / nodes (manual inventory from AGENTS/contracts) ──

add_entity("service", "neon_server", "neon-64gb", host="100.92.88.64", role="ARM64 VPS / netcup")
add_entity("service", "neon_postgres", "PostgreSQL on neon", container="arxiv-pg", host="neon-64gb")
add_entity("service", "neon_ollama", "Ollama on neon", url="http://100.92.88.64:11434", models="DeepSeek-Prover,Goedel-Prover,Qwen,hermes3")
add_entity("service", "vikunja_service", "Vikunja on neon", url="http://100.92.88.64:3456")
add_entity("service", "headroom_proxy_neon", "Headroom proxy on neon", url="http://100.92.88.64:8787")
add_entity("service", "gremlin_graph", "Azure Cosmos DB Gremlin / mathblob", account="mathblob", graph="concepts")
add_entity("service", "qfox_hermes", "Hermes dashboard on qfox-1", url="http://100.88.57.96:9119")
add_entity("service", "qfox_ollama", "Ollama on qfox-1", url="http://127.0.0.1:11434")

add_entity("database", "db:arxiv", "arxiv", host="neon_postgres", tables="arxiv_papers,arxiv_paper_codes8,citations")
add_entity("database", "db:ene", "ene", host="neon_postgres", schema="ENE graph schema")
add_entity("database", "db:appflowy", "appflowy", host="neon_postgres", app="AppFlowy-Cloud")
add_entity("database", "db:vikunja", "vikunja", host="neon_postgres", app="Vikunja")

add_entity("node", "node:qfox-1", "qfox-1", ip="100.88.57.96", role="workstation / GPU / Garage")
add_entity("node", "node:cupfox", "cupfox", ip="100.72.130.76", role="k3s control-plane / Garage")
add_entity("node", "node:nixos-laptop", "nixos-laptop", ip="100.102.173.61", role="k3s worker / Garage")
add_entity("node", "node:racknerd", "racknerd", ip="100.80.39.40", role="k3s worker / VPS")
add_entity("node", "node:neon-64gb", "neon-64gb", ip="100.92.88.64", role="ARM64 VPS / Postgres / Ollama")
add_entity("node", "node:steamdeck", "steamdeck", ip="100.85.244.73", role="k3s worker / GPU")

# Link services to nodes
add_edge("neon_postgres", "node:neon-64gb", "runs_on")
add_edge("neon_ollama", "node:neon-64gb", "runs_on")
add_edge("vikunja_service", "node:neon-64gb", "runs_on")
add_edge("headroom_proxy_neon", "node:neon-64gb", "runs_on")
add_edge("qfox_hermes", "node:qfox-1", "runs_on")
add_edge("qfox_ollama", "node:qfox-1", "runs_on")

# Link databases to service
for db in ["db:arxiv", "db:ene", "db:appflowy", "db:vikunja"]:
    add_edge(db, "neon_postgres", "hosted_by")

# ── 6. Goals / roadmaps ───────────────────────────────────────────────────────

GOAL_FILES = [
    ROOT / "6-Documentation" / "docs" / "roadmaps" / "ROADMAP.md",
    ROOT / "TODO_MAP.md",
]
for gf in GOAL_FILES:
    if gf.exists():
        text = gf.read_text(encoding="utf-8", errors="replace")
        title = first_heading(text) or gf.name
        gid = f"goal:{gf.relative_to(ROOT)}"
        add_entity("goal", gid, title, path=str(gf.relative_to(ROOT)))

# ── 7. Cross-link scripts to docs by name overlap ─────────────────────────────

for eid, ent in list(entities.items()):
    if ent["kind"] == "script":
        name = ent["name"].lower().replace("_", " ")
        for did, dent in entities.items():
            if dent["kind"] == "doc":
                title_l = dent.get("name", "").lower()
                if any(word in title_l for word in name.replace(".py", "").split() if len(word) > 4):
                    add_edge(eid, did, "references_doc")

# ── 8. Write JSON ─────────────────────────────────────────────────────────────

OUT_DIR.mkdir(parents=True, exist_ok=True)
json_path = OUT_DIR / "research_stack_usage_graph.json"
json_path.write_text(json.dumps({
    "entities": list(entities.values()),
    "edges": edges,
    "stats": {
        "entity_count": len(entities),
        "edge_count": len(edges),
        "by_kind": dict(Counter(e["kind"] for e in entities.values())),
    }
}, indent=2), encoding="utf-8")

# ── 9. Write DOT ──────────────────────────────────────────────────────────────

dot = ["digraph ResearchStackUsage {", "  rankdir=LR;", "  node [fontname=\"monospace\", fontsize=10];"]
kind_colors = {
    "module": "lightblue",
    "theorem": "lightcyan",
    "def": "lightcyan2",
    "script": "lightgoldenrod1",
    "doc": "palegreen",
    "service": "lightpink",
    "database": "lightyellow",
    "node": "lightgrey",
    "skill": "plum",
    "goal": "lightsalmon",
    "external_lib": "white",
}
for eid, ent in entities.items():
    color = kind_colors.get(ent["kind"], "white")
    label = f"{ent['kind']}\\n{ent['name'][:40]}"
    dot.append(f'  "{eid}" [style=filled, fillcolor={color}, label="{label}"];')

for edge in edges:
    dot.append(f'  "{edge["src"]}" -> "{edge["dst"]}" [label="{edge["kind"]}"];')
dot.append("}")
(OUT_DIR / "research_stack_usage_graph.dot").write_text("\n".join(dot), encoding="utf-8")

# ── 10. Write Markdown summary ────────────────────────────────────────────────

md = ["# Research Stack Searchable Usage Map", ""]
stats = Counter(e["kind"] for e in entities.values())
md.append(f"- **Entities:** {len(entities)}")
md.append(f"- **Edges:** {len(edges)}")
md.append("")
md.append("## Entity counts")
for kind, n in sorted(stats.items(), key=lambda x: -x[1]):
    md.append(f"- `{kind}`: {n}")
md.append("")

md.append("## Top modules by in-degree")
indeg = Counter(e["dst"] for e in edges if e["kind"] == "imports")
for mod, n in indeg.most_common(20):
    ent = entities.get(mod)
    name = ent["name"] if ent else mod
    md.append(f"- `{name}` — imported by {n} modules")
md.append("")

md.append("## Top theorem-producing modules")
mods = [e for e in entities.values() if e["kind"] == "module"]
for m in sorted(mods, key=lambda x: -x.get("theorem_count", 0))[:20]:
    md.append(f"- `{m['name']}` — theorems={m.get('theorem_count',0)} defs={m.get('def_count',0)} sorries={m.get('sorry_count',0)} kind={m.get('math_kind','?')}")
md.append("")

md.append("## Math-kind distribution")
kind_counts = Counter(m.get("math_kind", "?") for m in mods)
for k, n in kind_counts.most_common():
    md.append(f"- `{k}`: {n}")
md.append("")

md.append("## Neon-hosted services")
for svc in ["neon_postgres", "neon_ollama", "vikunja_service", "headroom_proxy_neon"]:
    ent = entities.get(svc)
    if ent:
        md.append(f"- **{ent['name']}** — {ent.get('url','')} {ent.get('container','')} {ent.get('models','')}")
md.append("")

md.append("## Databases on neon")
for db in ["db:arxiv", "db:ene", "db:appflowy", "db:vikunja"]:
    ent = entities.get(db)
    if ent:
        md.append(f"- **{ent['name']}** — {ent.get('tables','')} {ent.get('schema','')}")
md.append("")

md.append("## How to use this graph")
md.append("- Load `research_stack_usage_graph.json` into any search/index tool.")
md.append("- Render the DOT with: `dot -Tsvg research_stack_usage_graph.dot -o graph.svg`")
md.append("- Query examples: find modules with `math_kind=braid` and `sorry_count>0`;")
md.append("  find scripts that touch `db:arxiv`; find docs referenced by a script.")
md.append("")

(OUT_DIR / "research_stack_usage_graph.md").write_text("\n".join(md), encoding="utf-8")

print(f"Wrote {json_path}")
print(f"Wrote {OUT_DIR / 'research_stack_usage_graph.dot'}")
print(f"Wrote {OUT_DIR / 'research_stack_usage_graph.md'}")
print(f"Entities: {len(entities)}, edges: {len(edges)}")