"""NetworkX integration — convert AGE graphs to NetworkX graph objects."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any

from piggie.exceptions import IntegrationError

if TYPE_CHECKING:
    import networkx as nx

    from piggie.driver import PiggieDriver


def to_networkx(
    driver: PiggieDriver,
    graph: str,
    *,
    labels: list[str] | None = None,
    rel_types: list[str] | None = None,
    directed: bool = True,
    node_properties: bool = True,
    edge_properties: bool = True,
) -> nx.DiGraph | nx.Graph:
    """Convert an AGE graph (or subgraph) to a NetworkX graph.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        labels: Vertex labels to include. None means all labels.
        rel_types: Edge types to include. None means all types.
        directed: If True, return a DiGraph; otherwise Graph.
        node_properties: Include vertex properties as node attributes.
        edge_properties: Include edge properties as edge attributes.

    Returns:
        A NetworkX DiGraph or Graph.
    """
    try:
        import networkx as nx
    except ImportError:
        raise ImportError(
            "networkx is required. Install it with: pip install piggie[networkx]"
        ) from None

    G: nx.DiGraph | nx.Graph = nx.DiGraph() if directed else nx.Graph()

    # Discover labels if not specified
    if labels is None:
        schema = driver.get_schema(graph=graph)
        labels = [lbl["name"] for lbl in schema["vertex_labels"]]

    if rel_types is None:
        schema = driver.get_schema(graph=graph)
        rel_types = [lbl["name"] for lbl in schema["edge_labels"]]

    # Load nodes per label
    for label in labels:
        try:
            result = driver.cypher(f"MATCH (n:{label}) RETURN n", graph=graph)
        except Exception as exc:
            raise IntegrationError(f"Failed to query nodes for label {label}: {exc}") from exc

        for record in result:
            v = record["n"]
            attrs: dict[str, Any] = {"label": v.label}
            if node_properties:
                attrs.update(v.properties)
            G.add_node(v.id, **attrs)

    # Load edges per rel type
    for rel_type in rel_types:
        try:
            result = driver.cypher(
                f"MATCH ()-[r:{rel_type}]->() RETURN r", graph=graph
            )
        except Exception as exc:
            raise IntegrationError(
                f"Failed to query edges for type {rel_type}: {exc}"
            ) from exc

        for record in result:
            e = record["r"]
            attrs = {"label": e.label}
            if edge_properties:
                attrs.update(e.properties)
            G.add_edge(e.start_id, e.end_id, **attrs)

    return G


def to_networkx_sql(
    driver: PiggieDriver,
    graph: str,
    *,
    directed: bool = True,
    node_properties: bool = True,
    edge_properties: bool = True,
) -> nx.DiGraph | nx.Graph:
    """Convert an AGE graph to NetworkX via direct SQL reads.

    Bypasses Cypher planning and agtype ANTLR4 parsing by reading AGE's
    internal label tables directly. Typically 10-20x faster than
    :func:`to_networkx` for large graphs.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        directed: If True, return a DiGraph; otherwise Graph.
        node_properties: Include vertex properties as node attributes.
        edge_properties: Include edge properties as edge attributes.

    Returns:
        A NetworkX DiGraph or Graph.
    """
    try:
        import networkx as nx
    except ImportError:
        raise ImportError(
            "networkx is required. Install it with: pip install piggie[networkx]"
        ) from None

    conn = driver.connection
    schema = driver.get_schema(graph=graph)
    G: nx.DiGraph | nx.Graph = nx.DiGraph() if directed else nx.Graph()

    # Load nodes from each vertex label table.
    # Cast properties to text to bypass the agtype psycopg adapter
    # (agtype deserialization is ~250x slower than json.loads on text).
    for label_info in schema["vertex_labels"]:
        label = label_info["name"]
        try:
            with conn.cursor() as cur:
                if node_properties:
                    cur.execute(
                        f'SELECT id::text::bigint, properties::text '
                        f'FROM {graph}."{label}"'
                    )
                    for row in cur:
                        props = json.loads(row[1]) if row[1] else {}
                        props["label"] = label
                        G.add_node(row[0], **props)
                else:
                    cur.execute(
                        f'SELECT id::text::bigint FROM {graph}."{label}"'
                    )
                    for row in cur:
                        G.add_node(row[0], label=label)
        except Exception as exc:
            raise IntegrationError(
                f"Failed SQL read for vertex label {label}: {exc}"
            ) from exc

    # Load edges from each edge label table
    for label_info in schema["edge_labels"]:
        label = label_info["name"]
        try:
            with conn.cursor() as cur:
                if edge_properties:
                    cur.execute(
                        f'SELECT start_id::text::bigint, end_id::text::bigint, '
                        f'properties::text FROM {graph}."{label}"'
                    )
                    for row in cur:
                        props = json.loads(row[2]) if row[2] else {}
                        props["label"] = label
                        G.add_edge(row[0], row[1], **props)
                else:
                    cur.execute(
                        f'SELECT start_id::text::bigint, end_id::text::bigint '
                        f'FROM {graph}."{label}"'
                    )
                    for row in cur:
                        G.add_edge(row[0], row[1], label=label)
        except Exception as exc:
            raise IntegrationError(
                f"Failed SQL read for edge label {label}: {exc}"
            ) from exc

    return G
