"""DAG construction and traversal using networkx.

Replaces rTisane's dagitty-based graph operations.
"""

from __future__ import annotations

import networkx as nx

from .relationships import Causes, Relates


def build_graph(relationships) -> nx.DiGraph:
    """Build a directed graph from a list of Assumption/Hypothesis wrappers.

    - Causes  → single directed edge (cause → effect)
    - Relates → two directed edges tagged ``ambiguous=True``
    """
    g = nx.DiGraph()
    for entry in relationships:
        rel = entry.relationship
        if isinstance(rel, Causes):
            g.add_edge(rel.cause.name, rel.effect.name, ambiguous=False)
        elif isinstance(rel, Relates):
            g.add_edge(rel.lhs.name, rel.rhs.name, ambiguous=True)
            g.add_edge(rel.rhs.name, rel.lhs.name, ambiguous=True)
    return g


def parents(graph: nx.DiGraph, node: str) -> list[str]:
    """Direct predecessors of *node*."""
    if node not in graph:
        return []
    return list(graph.predecessors(node))


def children(graph: nx.DiGraph, node: str) -> list[str]:
    """Direct successors of *node*."""
    if node not in graph:
        return []
    return list(graph.successors(node))


def ancestors(graph: nx.DiGraph, node: str) -> set[str]:
    """All ancestors of *node* (including *node* itself, matching dagitty)."""
    if node not in graph:
        return {node}
    return nx.ancestors(graph, node) | {node}


def descendants(graph: nx.DiGraph, node: str) -> set[str]:
    """All descendants of *node* (including *node* itself)."""
    if node not in graph:
        return {node}
    return nx.descendants(graph, node) | {node}


def find_paths(graph: nx.DiGraph, source: str, target: str) -> list[list[str]]:
    """All simple paths from *source* to *target*."""
    if source not in graph or target not in graph:
        return []
    return list(nx.all_simple_paths(graph, source, target))
