# SPDX-FileCopyrightText: 2026 Jared Callaham <jared.callaham@gmail.com>
#
# SPDX-License-Identifier: GPL-3.0-or-later

"""Graph-based pruning for subsystem extraction.

This module provides graph analysis algorithms for identifying the minimal set of
blocks that influence a transfer function between two signals in a control diagram.
Uses bidirectional depth-first search (DFS) to find all blocks on paths from source
to destination, preserving internal feedback loops while excluding external couplings.

Complexity: O(V+E) where V = blocks, E = connections
Typical performance: <10ms for 50-block diagrams, <50ms for 100-block diagrams
"""

from typing import TYPE_CHECKING, Dict, List, Set

if TYPE_CHECKING:
    from lynx.blocks.base import Block
    from lynx.diagram import Diagram


def _build_connection_graph(
    diagram: "Diagram",
) -> tuple[Dict[str, List[str]], Dict[str, List[str]]]:
    """Build forward and backward adjacency lists from diagram connections.

    Args:
        diagram: Diagram to analyze

    Returns:
        Tuple of (forward_edges, backward_edges) where:
        - forward_edges[block_id] = list of target block IDs
        - backward_edges[block_id] = list of source block IDs
    """
    forward: Dict[str, List[str]] = {block.id: [] for block in diagram.blocks}
    backward: Dict[str, List[str]] = {block.id: [] for block in diagram.blocks}

    for conn in diagram.connections:
        forward[conn.source_block_id].append(conn.target_block_id)
        backward[conn.target_block_id].append(conn.source_block_id)

    return forward, backward


def _dfs_forward(
    forward_edges: Dict[str, List[str]], start_block_id: str, visited: Set[str]
) -> Set[str]:
    """Forward DFS: find all blocks reachable from start block.

    Args:
        forward_edges: Adjacency list for forward traversal
        start_block_id: Starting block ID
        visited: Set of already visited block IDs (for cycle detection)

    Returns:
        Set of block IDs reachable from start block
    """
    if start_block_id in visited:
        return set()  # Cycle detected, terminate

    visited.add(start_block_id)
    reachable = {start_block_id}

    for target_id in forward_edges.get(start_block_id, []):
        reachable |= _dfs_forward(forward_edges, target_id, visited)

    return reachable


def _dfs_backward(
    backward_edges: Dict[str, List[str]], start_block_id: str, visited: Set[str]
) -> Set[str]:
    """Backward DFS: find all blocks that can reach start block.

    Args:
        backward_edges: Adjacency list for backward traversal
        start_block_id: Starting block ID
        visited: Set of already visited block IDs (for cycle detection)

    Returns:
        Set of block IDs that can reach start block
    """
    if start_block_id in visited:
        return set()  # Cycle detected, terminate

    visited.add(start_block_id)
    reachable = {start_block_id}

    for source_id in backward_edges.get(start_block_id, []):
        reachable |= _dfs_backward(backward_edges, source_id, visited)

    return reachable


def _find_reachable_blocks(
    diagram: "Diagram", source_block: "Block", dest_block: "Block"
) -> Set[str]:
    """Find all blocks on any path from source to destination.

    Uses bidirectional reachability analysis: blocks must be both
    forward-reachable from source AND backward-reachable from destination.

    The key insight is that the intersection naturally excludes:
    - Blocks downstream of destination (forward-reachable but not backward-reachable)
    - Blocks upstream of source (backward-reachable but not forward-reachable)
    - Unrelated blocks (neither forward nor backward reachable)

    Args:
        diagram: Diagram to search
        source_block: Starting block
        dest_block: Ending block

    Returns:
        Set of block IDs that are on at least one path

    Raises:
        SignalNotFoundError: If no path exists from source to destination
    """
    forward_edges, backward_edges = _build_connection_graph(diagram)

    # Forward pass: blocks reachable from source
    forward_reachable = _dfs_forward(forward_edges, source_block.id, set())

    # Backward pass: blocks that can reach destination
    backward_reachable = _dfs_backward(backward_edges, dest_block.id, set())

    # Check if destination is reachable from source
    if dest_block.id not in forward_reachable:
        from ..diagram import SignalNotFoundError

        raise SignalNotFoundError(
            signal_name=f"{source_block.id} → {dest_block.id}",
            searched_locations=["forward reachability"],
            custom_message="No path exists from source to destination",
        )

    # Intersection: blocks on any path from source to destination
    # This automatically excludes:
    # - Blocks downstream of dest (in forward but not backward)
    # - Blocks upstream of source (in backward but not forward)
    path_blocks = forward_reachable & backward_reachable

    # Always include source and destination (even for same-block extraction)
    path_blocks.add(source_block.id)
    path_blocks.add(dest_block.id)

    return path_blocks


def prune_diagram(diagram: "Diagram", keep_blocks: Set[str]) -> "Diagram":
    """Remove all blocks not in keep_blocks set.

    Creates a pruned copy of the diagram with only the blocks in keep_blocks.
    Connections to removed blocks are automatically cleaned up.

    Args:
        diagram: Diagram to prune (will be cloned, not modified)
        keep_blocks: Set of block IDs to preserve

    Returns:
        Pruned diagram (clone with removed blocks)
    """
    # Clone to avoid modifying original
    pruned = diagram._clone()

    # Identify blocks to remove
    blocks_to_remove = [
        block.id for block in pruned.blocks if block.id not in keep_blocks
    ]

    # Remove blocks (connections are auto-cleaned by remove_block)
    for block_id in blocks_to_remove:
        pruned.remove_block(block_id)

    return pruned
