from __future__ import annotations
"""
aml_engine/dependency_graph.py - Rule dependency DAG.

Builds a directed acyclic graph (DAG) of rule dependencies using NetworkX.
Provides topological ordering so dependent rules run after their prerequisites.
"""


import logging
from typing import Dict, List, Optional

try:
    import networkx as nx
    NETWORKX_AVAILABLE = True
except ImportError:
    NETWORKX_AVAILABLE = False

from RULE.aml_engine.rule_models import AMLRule

logger = logging.getLogger(__name__)


class RuleDependencyGraph:
    """
    Directed Acyclic Graph of AML rule dependencies.

    Nodes: rule_id strings
    Edges: (dependency_id → rule_id)  →  dependency must run first

    Usage:
        graph = RuleDependencyGraph(rules)
        ordered_ids = graph.topological_order()
        deps = graph.dependencies_of("R-011")
    """

    def __init__(self, rules: List[AMLRule]):
        self._rules: Dict[str, AMLRule] = {r.rule_id: r for r in rules}
        self._graph: Optional[object] = None

        if NETWORKX_AVAILABLE:
            self._graph = self._build_graph(rules)
        else:
            logger.warning(
                "NetworkX not available. Topological ordering will fall back to "
                "file order. Install with: pip install networkx"
            )

    @staticmethod
    def _build_graph(rules: List[AMLRule]):
        """Build a DiGraph with edges dep_id → rule_id."""
        graph = nx.DiGraph()
        for rule in rules:
            graph.add_node(rule.rule_id, rule=rule)
        for rule in rules:
            for dep_id in rule.interdependent_rules:
                if dep_id in {r.rule_id for r in rules}:
                    graph.add_edge(dep_id, rule.rule_id)
                else:
                    logger.warning(
                        "Rule '%s' references unknown dependency '%s' — edge skipped.",
                        rule.rule_id, dep_id,
                    )
        return graph

    def topological_order(self) -> List[str]:
        """Return rule IDs in topological execution order (dependencies first)."""
        if not NETWORKX_AVAILABLE or self._graph is None:
            logger.warning("Falling back to original file order for rule execution.")
            return list(self._rules.keys())
        try:
            ordered = list(nx.topological_sort(self._graph))
            return ordered
        except nx.NetworkXUnfeasible:
            raise ValueError(
                "Circular rule dependency detected — cannot compute topological order."
            )

    def dependencies_of(self, rule_id: str) -> List[str]:
        """Return direct dependency rule IDs for a given rule."""
        if self._graph is None:
            return []
        if rule_id not in self._graph:
            return []
        return list(self._graph.predecessors(rule_id))

    def dependents_of(self, rule_id: str) -> List[str]:
        """Return rule IDs that depend on the given rule."""
        if self._graph is None:
            return []
        if rule_id not in self._graph:
            return []
        return list(self._graph.successors(rule_id))

    def independent_rules(self) -> List[str]:
        """Return rule IDs that have no dependencies (in-degree == 0)."""
        if self._graph is None:
            return [rid for rid, rule in self._rules.items() if not rule.interdependent_rules]
        return [n for n, deg in self._graph.in_degree() if deg == 0]

    def rule_by_id(self, rule_id: str) -> Optional[AMLRule]:
        """Look up an AMLRule by its ID."""
        return self._rules.get(rule_id)

    def summary(self) -> Dict:
        """Return a dict summary of the graph."""
        total = len(self._rules)
        independent = len(self.independent_rules())
        return {
            "total_rules": total,
            "independent_rules": independent,
            "interdependent_rules": total - independent,
            "has_networkx": NETWORKX_AVAILABLE,
        }


def build_rule_graph(rules: List[AMLRule]) -> RuleDependencyGraph:
    """Build a RuleDependencyGraph from a list of AMLRule objects."""
    return RuleDependencyGraph(rules)