Graph visualization — Cytoscape network with click handlers

from __future__ import annotations

from imperal_sdk import Extension, ChatExtension, ActionResult, ui

ext = Extension(
    "my-ext",
    display_name="My Extension",
    description="An extension demonstrating graph visualization with node drill-down.",
    actions_explicit=True,
)

chat = ChatExtension(
    ext=ext,
    tool_name="tool_my_ext_chat",
    description="My Extension — browse entity relationships in a graph.",
)

# ── Stub graph data — replace with ctx.store or API fetch ─────────────────
#
# In production, fetch entities + relationships from your own backend
# (or ctx.store / an external API) and map them to this flat dict format.

_NODES = [
    {"id": "person-1",  "label": "Alice",       "type": "person"},
    {"id": "person-2",  "label": "Bob",         "type": "person"},
    {"id": "company-1", "label": "Acme Corp",   "type": "company"},
    {"id": "company-2", "label": "Globex Inc",  "type": "company"},
    {"id": "event-1",   "label": "Contract A",  "type": "event"},
]

_EDGES = [
    {"id": "e1", "source": "person-1",  "target": "company-1", "label": "works_at"},
    {"id": "e2", "source": "person-2",  "target": "company-1", "label": "works_at"},
    {"id": "e3", "source": "person-1",  "target": "company-2", "label": "director"},
    {"id": "e4", "source": "company-1", "target": "event-1",   "label": "party_to"},
    {"id": "e5", "source": "company-2", "target": "event-1",   "label": "party_to"},
]

_ENTITY_DETAIL: dict[str, dict] = {
    "person-1":  {"name": "Alice",      "type": "person",  "role": "Director"},
    "person-2":  {"name": "Bob",        "type": "person",  "role": "Employee"},
    "company-1": {"name": "Acme Corp",  "type": "company", "reg": "GB12345678"},
    "company-2": {"name": "Globex Inc", "type": "company", "reg": "GB87654321"},
    "event-1":   {"name": "Contract A", "type": "event",   "date": "2026-01-15"},
}


# ── Left sidebar — graph overview ─────────────────────────────────────────

@ext.panel(
    "graph",
    slot="left",
    title="Relationships",
    icon="🌐",
    default_width=400,
    min_width=320,
    max_width=720,
)
async def graph_panel(ctx: object, **kwargs: object) -> object:
    # ui.Graph accepts flat dicts: {"id", "label", "type"} for nodes,
    # {"id", "source", "target", "label"} for edges.
    # The SDK also accepts Cytoscape {"data": {...}} wrappers — both are normalised.
    #
    # color_by="type" assigns distinct colors per unique value of the "type" field.
    # on_node_click fires when the user selects a node; the node's "id" is
    # injected as the "node_id" kwarg in the target function call.
    return ui.Stack(
        gap=2,
        children=[
            ui.Header(
                "Entity network",
                level=2,
                subtitle=f"{len(_NODES)} entities · {len(_EDGES)} connections",
            ),
            ui.Graph(
                nodes=_NODES,
                edges=_EDGES,
                layout="cose-bilkent",
                height=520,
                color_by="type",
                edge_label_visible=True,
                on_node_click=ui.Call("__panel__entity_detail"),
            ),
        ],
    )


# ── Center panel — entity drill-down ──────────────────────────────────────

@ext.panel(
    "entity_detail",
    slot="center",
    center_overlay=True,
    title="Entity detail",
    icon="ℹ️",
)
async def entity_detail(ctx: object, node_id: str = "", **kwargs: object) -> object:
    # on_node_click injects the selected node's "id" as "node_id".
    # Return ui.Empty() at discovery (node_id is "").
    if not node_id:
        return ui.Empty(
            message="Select a node in the graph to view details",
            icon="🖱️",
        )

    entity = _ENTITY_DETAIL.get(node_id)
    if entity is None:
        return ui.Error(message=f"Entity '{node_id}' not found.")

    # Build a KeyValue card with all fields from the entity dict.
    kv_items = [{"key": k.capitalize(), "value": str(v)} for k, v in entity.items()]

    return ui.Stack(
        gap=4,
        children=[
            ui.Header(entity["name"], level=2, subtitle=entity["type"].capitalize()),
            ui.KeyValue(items=kv_items, columns=2),
        ],
    )


# ── chat.function — fetch the graph via chat ──────────────────────────────

@chat.function(
    "show_entity_graph",
    description="Display the entity relationship graph in the panel.",
    action_type="read",
)
async def show_entity_graph(ctx: object) -> ActionResult:
    return ActionResult(
        status="ok",
        summary=f"Graph loaded — {len(_NODES)} entities, {len(_EDGES)} connections.",
        refresh_panels=["graph"],
    )