# mcp/knowledge_graph.py

from streamlit_agraph import Node, Edge, Config
import re

def build_agraph(papers, umls, drug_safety):
    """
    Build interactive agraph nodes and edges.
    Handles drug_safety entries that may be dict or list.
    """

    nodes, edges = [], []

    # Add UMLS concept nodes
    for c in umls:
        cui = c.get("cui")
        name = c.get("name", "")
        if cui and name:
            nid = f"concept_{cui}"
            nodes.append(Node(id=nid, label=name, size=25, color="#00b894"))

    # Add drug nodes, handling list or dict
    drug_names = []
    for i, dr in enumerate(drug_safety):
        if not dr:
            continue

        # Normalize to single dict
        recs = dr if isinstance(dr, list) else [dr]
        for j, rec in enumerate(recs):
            # Attempt to extract a drug name
            dn = rec.get("drug_name") or rec.get("patient", {}).get("drug", "") or rec.get("medicinalproduct", "")
            dn = dn or f"drug_{i}_{j}"
            did = f"drug_{i}_{j}"
            drug_names.append((did, dn))
            nodes.append(Node(id=did, label=dn, size=25, color="#d35400"))

    # Add paper nodes and link to concepts & drugs
    for pi, p in enumerate(papers):
        pid = f"paper_{pi}"
        nodes.append(Node(id=pid, label=f"P{pi+1}", tooltip=p["title"], size=15, color="#0984e3"))

        text = f"{p.get('title','')} {p.get('summary','')}".lower()
        # Link to concepts
        for c in umls:
            cname = c.get("name", "")
            cui = c.get("cui")
            if cname and cui and cname.lower() in text:
                edges.append(Edge(source=pid, target=f"concept_{cui}", label="mentions"))
        # Link to drugs
        for did, dn in drug_names:
            if dn.lower() in text:
                edges.append(Edge(source=pid, target=did, label="mentions"))

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