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

def build_agraph(res: Dict) -> (list, list, Config):
    nodes, edges = [], []
    # add each paper as a node
    for i,p in enumerate(res["papers"]):
        nid = f"paper_{i}"
        nodes.append(Node(id=nid, label=p["title"], size=20, color="#0984e3"))
        # connect to AI summary?
    # add UMLS concepts
    for u in res["umls"]:
        cid = f"cui_{u['cui']}"
        label = f"{u['name']} ({u['cui']})"
        nodes.append(Node(id=cid, label=label, size=25, color="#00b894"))
        # connect concept → first paper
        edges.append(Edge(source=cid, target="paper_0", label="mentioned_in"))
    # genes
    g = res.get("gene",{})
    if g:
        gid = "gene_node"
        nodes.append(Node(id=gid, label=g.get("symbol",g.get("name","gene")), color="#d63031"))
        edges.append(Edge(source=gid, target="cui_"+res["umls"][0]["cui"], label="related"))
    # variants
    for v in res["variants"]:
        vid = f"var_{v['mutationId']}"
        nodes.append(Node(id=vid, label=v["mutationId"], color="#fdcb6e", size=15))
        edges.append(Edge(source=vid, target=gid, label="affects"))
    # trials
    for t in res["trials"]:
        tid = t["NCTId"][0]
        nodes.append(Node(id=tid, label=tid, color="#6c5ce7"))
        edges.append(Edge(source=tid, target=gid, label="studies"))

    cfg = Config(width="100%", height="600", directed=True,
                 nodeHighlightBehavior=True, highlightColor="#fdcb6e")
    return nodes, edges, cfg