# mcp/knowledge_graph.py
from streamlit_agraph import Node, Edge, Config
import re

def build_agraph(papers, umls, drug_safety, umls_relations=None):
    nodes, edges = [], []

    # 1) concepts
    for c in umls:
        cui, name = c.get("cui"), c.get("name")
        types = c.get("types", [])
        if not cui or not name: 
            continue
        label = f"{name}\n({'/'.join(types)})"
        size  = 20 + int(c.get("score",0)*30)
        nodes.append(Node(
            id=f"cui_{cui}",
            label=label,
            size=size,
            color="#00b894",
            tooltip=c.get("definition","")
        ))

    # 2) concept→concept edges
    if umls_relations:
        for rels in umls_relations:
            for r in rels:
                s, t = f"cui_{r['cui']}", f"cui_{r['cui2']}"
                if s != t:
                    edges.append(Edge(source=s, target=t, label=r["label"]))

    # 3) papers
    for i,p in enumerate(papers):
        pid = f"paper_{i}"
        nodes.append(Node(
            id=pid, label=f"P{i+1}", size=15, color="#0984e3",
            tooltip=p.get("title","")
        ))
        txt = (p.get("title","") + " " + p.get("summary","")).lower()
        for c in umls:
            if c.get("name","").lower() in txt:
                edges.append(Edge(source=pid, target=f"cui_{c['cui']}", label="mentions"))

    # 4) drugs (as before)…
    # nodes.extend(...), edges.extend(...)

    cfg = Config(
        width="100%", height="600px", directed=True,
        nodeHighlightBehavior=True, highlightColor="#f1c40f",
        collapsible=True, node={"labelProperty":"label"}
    )
    return nodes, edges, cfg