"""
mermaid_graph.converter
====================
Convert between the extended node_link_data dict and NetworkX graphs.

Subgraph membership is stored as a ``_subgraph_path`` node attribute
(list of subgraph IDs from outermost to innermost) so that the tree
structure can be losslessly restored when converting back to JSON.

Subgraph labels are stored separately as a ``_subgraph_labels`` graph
attribute (mapping from subgraph ID to label) so that labels survive
the round-trip even when subgraphs are rebuilt from node attributes.
"""

from __future__ import annotations

from copy import deepcopy
from typing import Any

import networkx as nx


# ════════════════════════════════════════════════════════════════════════
#  JSON → NetworkX
# ════════════════════════════════════════════════════════════════════════

def to_networkx(data: dict[str, Any]) -> nx.DiGraph | nx.MultiDiGraph | nx.Graph | nx.MultiGraph:
    """
    Convert an extended node_link_data dict to a NetworkX graph.

    * Mermaid metadata is preserved in ``G.graph`` attributes.
    * Each node gets a ``_subgraph_path`` attribute listing its
      subgraph ancestors (outermost first).
    * If the data contains parallel edges (``multigraph: true``),
      a ``MultiDiGraph`` / ``MultiGraph`` is used so that no edges
      are silently dropped.
    """
    # Deep-copy to avoid mutating the caller's data
    enriched = deepcopy(data)

    # 1) Build subgraph membership map and label map
    membership: dict[str, list[str]] = {}
    label_map: dict[str, str] = {}
    _build_membership(
        enriched.get("graph", {}).get("subgraphs", []),
        parent_path=[],
        result=membership,
        label_map=label_map,
    )

    # 2) Attach membership to nodes
    for node in enriched["nodes"]:
        node["_subgraph_path"] = membership.get(node["id"], [])

    # 3) Store label map in graph metadata
    enriched.setdefault("graph", {})["_subgraph_labels"] = label_map

    # 4) Normalise key: our format uses "links", nx 3.x expects "edges"
    enriched = _links_to_edges(enriched)

    # 5) Delegate to NetworkX
    G = nx.node_link_graph(enriched)
    return G


def _build_membership(
    subgraphs: list[dict],
    parent_path: list[str],
    result: dict[str, list[str]],
    label_map: dict[str, str],
) -> None:
    for sg in subgraphs:
        current_path = parent_path + [sg["id"]]
        sg_label = sg.get("label", sg["id"])
        label_map[sg["id"]] = sg_label
        for node_id in sg.get("nodes", []):
            result[node_id] = current_path
        _build_membership(
            sg.get("subgraphs", []),
            current_path,
            result,
            label_map,
        )


# ════════════════════════════════════════════════════════════════════════
#  NetworkX → JSON
# ════════════════════════════════════════════════════════════════════════

def from_networkx(G: nx.DiGraph | nx.MultiDiGraph | nx.Graph | nx.MultiGraph) -> dict[str, Any]:
    """
    Convert a NetworkX graph back to the extended node_link_data dict.

    If nodes carry ``_subgraph_path`` attributes, the subgraph tree
    in ``graph.subgraphs`` is automatically reconstructed.
    """
    data = nx.node_link_data(G)

    # Normalise key: nx 3.x outputs "edges", our format uses "links"
    data = _edges_to_links(data)

    # Rebuild subgraphs from node attributes if not already present
    subgraphs_in_graph = data.get("graph", {}).get("subgraphs")
    if subgraphs_in_graph is None:
        label_map = data.get("graph", {}).pop("_subgraph_labels", {})
        data.setdefault("graph", {})["subgraphs"] = _rebuild_subgraphs(
            data["nodes"], label_map
        )
    else:
        # Clean up the internal label map if subgraphs are already present
        data.get("graph", {}).pop("_subgraph_labels", None)

    # Clean up temporary attributes
    for node in data["nodes"]:
        node.pop("_subgraph_path", None)

    return data


# ── Key normalisation helpers (nx 2.x uses "links", nx 3.x uses "edges") ─

def _links_to_edges(data: dict[str, Any]) -> dict[str, Any]:
    """Rename 'links' → 'edges' if needed (for nx.node_link_graph)."""
    if "links" in data and "edges" not in data:
        result = {k: v for k, v in data.items() if k != "links"}
        result["edges"] = data["links"]
        return result
    return data


def _edges_to_links(data: dict[str, Any]) -> dict[str, Any]:
    """Rename 'edges' → 'links' if needed (for our JSON format)."""
    if "edges" in data and "links" not in data:
        result = {k: v for k, v in data.items() if k != "edges"}
        result["links"] = data["edges"]
        return result
    return data


def _rebuild_subgraphs(
    nodes: list[dict],
    label_map: dict[str, str] | None = None,
) -> list[dict]:
    """Reconstruct subgraph tree from node ``_subgraph_path`` attributes."""
    if label_map is None:
        label_map = {}

    tree: dict[str, dict[str, Any]] = {}

    for node in nodes:
        path: list[str] = node.get("_subgraph_path", [])
        if not path:
            continue

        # Ensure every level of the path exists in the tree
        current_level = tree
        for sg_id in path:
            if sg_id not in current_level:
                current_level[sg_id] = {
                    "id": sg_id,
                    "label": label_map.get(sg_id, sg_id),
                    "nodes": [],
                    "children": {},
                }
            current_level = current_level[sg_id]["children"]

        # Place node in deepest subgraph
        target = _find_at_path(tree, path)
        target["nodes"].append(node["id"])

    return _tree_to_list(tree)


def _find_at_path(tree: dict, path: list[str]) -> dict:
    current = tree[path[0]]
    for sg_id in path[1:]:
        current = current["children"][sg_id]
    return current


def _tree_to_list(tree: dict[str, dict]) -> list[dict]:
    result = []
    for sg in tree.values():
        entry: dict[str, Any] = {
            "id": sg["id"],
            "label": sg.get("label", sg["id"]),
            "nodes": sg["nodes"],
            "subgraphs": _tree_to_list(sg.get("children", {})),
        }
        if sg.get("style") is not None:
            entry["style"] = sg["style"]
        result.append(entry)
    return result