"""Fast graph compilation for dbt build.

Monkey-patches Linker.add_test_edges with dbt's own add_test_edges_2
algorithm (which is gated behind a flag we cannot control), and replaces
find_cycles with a Kahn's-algorithm implementation that is O(V+E) instead
of networkx's much slower edge-DFS.
"""

from __future__ import annotations

import dataclasses
import typing as t
from collections import defaultdict, deque

if t.TYPE_CHECKING:
    from dbt.compilation import Linker
    from dbt.contracts.graph.manifest import Manifest

UniqueID = str


def _get_tests_for_node(manifest: Manifest, unique_id: UniqueID) -> t.List[UniqueID]:
    """Get a list of tests that depend on the node with the provided unique id."""
    tests = []
    if unique_id in manifest.child_map:
        for child_unique_id in manifest.child_map[unique_id]:
            if child_unique_id.startswith("test."):
                tests.append(child_unique_id)
    return tests


@dataclasses.dataclass
class _SeenDetails:
    node_id: UniqueID
    visits: int = 0
    ancestors: t.Set[UniqueID] = dataclasses.field(default_factory=set)
    awaits_tests: t.Set[t.Tuple[UniqueID, t.Tuple[UniqueID, ...]]] = dataclasses.field(
        default_factory=set
    )


def _fast_add_test_edges(linker: Linker, manifest: Manifest) -> None:
    """Direct port of dbt's add_test_edges_2 / _get_test_edges_2 / _get_multi_test_edges."""

    graph = linker.graph
    new_edges = _get_test_edges_2(graph, manifest)
    for e in new_edges:
        graph.add_edge(e[0], e[1], edge_type="parent_test")


def _get_test_edges_2(
    graph: "nx.DiGraph", manifest: Manifest
) -> t.List[t.Tuple[UniqueID, UniqueID]]:
    new_edges: t.List[t.Tuple[UniqueID, UniqueID]] = []

    source_nodes: t.List[UniqueID] = []
    executable_nodes: t.Set[UniqueID] = set()
    multi_tested_nodes: t.Set[UniqueID] = set()
    single_tested_nodes: t.Dict[UniqueID, t.List[UniqueID]] = defaultdict(list)

    for node_id in graph.nodes:
        manifest_node = manifest.nodes.get(node_id, None)
        if manifest_node is None:
            continue

        if next(graph.predecessors(node_id), None) is None:
            source_nodes.append(node_id)

        if manifest_node.resource_type.value != "test":
            executable_nodes.add(node_id)
        else:
            test_deps = manifest_node.depends_on_nodes
            if len(test_deps) == 1:
                single_tested_nodes[test_deps[0]].append(node_id)
            elif len(test_deps) > 1:
                multi_tested_nodes.update(manifest_node.depends_on_nodes)

    # Add new edges for single-tested nodes.
    for node_id, test_ids in single_tested_nodes.items():
        succs = [s for s in graph.successors(node_id) if s in executable_nodes]
        for succ_id in succs:
            for test_id in test_ids:
                new_edges.append((test_id, succ_id))

    # Get the edges for multi-tested nodes separately, if needed.
    if len(multi_tested_nodes) > 0:
        multi_test_edges = _get_multi_test_edges(
            graph, manifest, source_nodes, executable_nodes, multi_tested_nodes
        )
        new_edges += multi_test_edges

    return new_edges


def _get_multi_test_edges(
    graph: "nx.DiGraph",
    manifest: Manifest,
    source_nodes: t.List[UniqueID],
    executable_nodes: t.Set[UniqueID],
    multi_tested_nodes: t.Set[UniqueID],
) -> t.List[t.Tuple[UniqueID, UniqueID]]:
    new_edges: t.List[t.Tuple[UniqueID, UniqueID]] = []
    ready: deque[UniqueID] = deque(source_nodes)
    details = {node_id: _SeenDetails(node_id) for node_id in source_nodes}

    while len(ready) > 0:
        curr_details = details[ready.pop()]
        test_ids = _get_tests_for_node(manifest, curr_details.node_id)
        new_awaits_for_succs = curr_details.awaits_tests.copy()
        for test_id in test_ids:
            deps: t.List[UniqueID] = sorted(manifest.nodes[test_id].depends_on_nodes)
            if len(deps) > 1:
                new_awaits_for_succs.add((test_id, tuple(deps)))

        for succ_id in [s for s in graph.successors(curr_details.node_id) if s in executable_nodes]:
            suc_details = details.get(succ_id, None)
            if suc_details is None:
                suc_details = _SeenDetails(succ_id)
                details[succ_id] = suc_details
            suc_details.visits += 1
            suc_details.awaits_tests.update(new_awaits_for_succs)
            suc_details.ancestors.update(curr_details.ancestors)
            if curr_details.node_id in multi_tested_nodes:
                suc_details.ancestors.add(curr_details.node_id)

            if suc_details.visits == graph.in_degree(succ_id):
                if len(suc_details.awaits_tests) > 0:
                    removes = set()
                    for awt in suc_details.awaits_tests:
                        if not any(True for a in awt[1] if a not in suc_details.ancestors):
                            removes.add(awt)
                            new_edges.append((awt[0], succ_id))
                    suc_details.awaits_tests.difference_update(removes)
                ready.appendleft(succ_id)

        del details[curr_details.node_id]

    return new_edges


def _fast_find_cycles(linker: Linker) -> t.Optional[str]:
    """Detect cycles using Kahn's algorithm (topological sort).

    O(V+E) vs networkx's edge DFS which is much slower on large graphs.
    Returns None if no cycle, or a cycle description string if found.
    """
    graph = linker.graph
    in_degree: t.Dict[str, int] = {}
    children: t.Dict[str, t.List[str]] = defaultdict(list)

    for node_id in graph.nodes:
        in_degree[node_id] = 0
    for u, v in graph.edges():
        children[u].append(v)
        in_degree[v] += 1

    queue = deque(n for n, d in in_degree.items() if d == 0)
    visited = 0

    while queue:
        node_id = queue.popleft()
        visited += 1
        for child in children[node_id]:
            in_degree[child] -= 1
            if in_degree[child] == 0:
                queue.append(child)

    if visited == len(in_degree):
        return None

    # Cycle detected — find nodes involved for the error message
    cycle_nodes = sorted(n for n, d in in_degree.items() if d > 0)
    return " --> ".join(cycle_nodes[:10])


def install_fast_graph() -> None:
    """Monkey-patch Linker.add_test_edges and find_cycles with fast implementations."""
    import logging

    from dbt.compilation import Linker

    logger = logging.getLogger(__name__)

    Linker.add_test_edges = _fast_add_test_edges
    Linker.find_cycles = _fast_find_cycles
    logger.debug("Installed fast graph compilation (add_test_edges + find_cycles)")
