"""Topology-only valve isolation search."""

from collections import defaultdict, deque
from typing import Any, Iterable


VALVE_LINK_TYPE = "valve"


def _parse_link_entry(link_entry: str) -> tuple[str, str, str, str]:
    parts = link_entry.split(":", 3)
    if len(parts) != 4:
        raise ValueError(f"Invalid link entry format: {link_entry}")
    return parts[0], parts[1], parts[2], parts[3]


def valve_isolation_analysis(
    link_entries: Iterable[str],
    accident_elements: str | list[str],
    disabled_valves: list[str] | None = None,
) -> dict[str, Any]:
    """Determine boundary valves and affected nodes from a topology snapshot."""

    disabled_valves_set = set(disabled_valves or [])
    target_elements = (
        [accident_elements]
        if isinstance(accident_elements, str)
        else accident_elements
    )

    pipe_adj: dict[str, set[str]] = defaultdict(set)
    all_valves: dict[str, tuple[str, str]] = {}
    link_lookup: dict[str, tuple[str, str, str]] = {}
    node_set: set[str] = set()
    for link_entry in link_entries:
        link_id, link_type, node1, node2 = _parse_link_entry(link_entry)
        link_type_name = str(link_type).lower()
        link_lookup[link_id] = (node1, node2, link_type_name)
        node_set.update((node1, node2))
        if link_type_name == VALVE_LINK_TYPE:
            all_valves[link_id] = (node1, node2)
        else:
            pipe_adj[node1].add(node2)
            pipe_adj[node2].add(node1)

    start_nodes: set[str] = set()
    for element in target_elements:
        if element in node_set:
            start_nodes.add(element)
        elif element in link_lookup:
            node1, node2, _ = link_lookup[element]
            start_nodes.update((node1, node2))
        else:
            raise ValueError(f"Accident element {element} was not found in topology")

    extra_adj: dict[str, list[str]] = defaultdict(list)
    boundary_valves: dict[str, tuple[str, str]] = {}
    for valve_id, (node1, node2) in all_valves.items():
        if valve_id in disabled_valves_set:
            extra_adj[node1].append(node2)
            extra_adj[node2].append(node1)
        else:
            boundary_valves[valve_id] = (node1, node2)

    affected_nodes: set[str] = set()
    queue = deque(start_nodes)
    while queue:
        node = queue.popleft()
        if node in affected_nodes:
            continue
        affected_nodes.add(node)
        queue.extend(pipe_adj.get(node, set()) - affected_nodes)
        queue.extend(
            neighbor
            for neighbor in extra_adj.get(node, ())
            if neighbor not in affected_nodes
        )

    must_close_valves: list[str] = []
    optional_valves: list[str] = []
    for valve_id, (node1, node2) in boundary_valves.items():
        node1_affected = node1 in affected_nodes
        node2_affected = node2 in affected_nodes
        if node1_affected and node2_affected:
            optional_valves.append(valve_id)
        elif node1_affected or node2_affected:
            must_close_valves.append(valve_id)

    must_close_valves.sort()
    optional_valves.sort()
    isolatable = bool(must_close_valves)
    result: dict[str, Any] = {
        "accident_elements": target_elements,
        "disabled_valves": disabled_valves,
        "affected_nodes": sorted(affected_nodes) if isolatable else [],
        "affected_node_count": len(affected_nodes),
        "must_close_valves": must_close_valves,
        "optional_valves": optional_valves,
        "isolatable": isolatable,
    }
    if len(target_elements) == 1:
        result["accident_element"] = target_elements[0]
    return result