from collections import defaultdict, deque
from functools import lru_cache
import time
from typing import Any

from app.services.tjnetwork import (
    get_network_link_nodes,
    is_node,
    get_link_properties,
)


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]


@lru_cache(maxsize=16)
def _get_network_topology(network: str):
    """
    解析并缓存网络拓扑，大幅减少重复的 API 调用和字符串解析开销。
    返回:
    - pipe_adj: 永久连通的管道/泵邻接表 (dict[str, set])
    - all_valves: 所有阀门字典 {id: (n1, n2)}
    - link_lookup: 链路快速查表 {id: (n1, n2, type)} 用于快速定位事故点
    - node_set: 所有已知节点集合
    """
    pipe_adj = defaultdict(set)
    all_valves = {}
    link_lookup = {}
    node_set = set()

    # 此处假设 get_network_link_nodes 获取全网数据
    for link_entry in get_network_link_nodes(network):
        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.add(node1)
        node_set.add(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)

    return pipe_adj, all_valves, link_lookup, node_set


def valve_isolation_analysis(
    network: str, accident_elements: str | list[str], disabled_valves: list[str] = None
) -> dict[str, Any]:
    """
    关阀搜索/分析：基于拓扑结构确定事故隔离所需关阀。
    :param network: 模型名称
    :param accident_elements: 事故点（节点或管道/泵/阀门ID），可以是单个ID字符串或ID列表
    :param disabled_valves: 故障/无法关闭的阀门ID列表
    :return: dict，包含受影响节点、必须关闭阀门、可选阀门等信息
    """
    if disabled_valves is None:
        disabled_valves_set = set()
    else:
        disabled_valves_set = set(disabled_valves)

    if isinstance(accident_elements, str):
        target_elements = [accident_elements]
    else:
        target_elements = accident_elements

    # 1. 获取缓存拓扑 (极快，无 IO)
    pipe_adj, all_valves, link_lookup, node_set = _get_network_topology(network)

    # 2. 确定起点，优先查表避免 API 调用
    start_nodes = set()
    for element in target_elements:
        if element in node_set:
            start_nodes.add(element)
        elif element in link_lookup:
            n1, n2, _ = link_lookup[element]
            start_nodes.add(n1)
            start_nodes.add(n2)
        else:
            # 仅当缓存中没找到时（极少见），才回退到慢速 API
            if is_node(network, element):
                start_nodes.add(element)
            else:
                props = get_link_properties(network, element)
                n1, n2 = props.get("node1"), props.get("node2")
                if n1 and n2:
                    start_nodes.add(n1)
                    start_nodes.add(n2)
                else:
                    raise ValueError(
                        f"Accident element {element} invalid or missing endpoints"
                    )

    # 3. 处理故障阀门 (构建临时增量图)
    # 我们不修改 cached pipe_adj，而是建立一个 extra_adj
    extra_adj = defaultdict(list)
    boundary_valves = {}  # 当前有效的边界阀门

    for vid, (n1, n2) in all_valves.items():
        if vid in disabled_valves_set:
            # 故障阀门：视为连通管道
            extra_adj[n1].append(n2)
            extra_adj[n2].append(n1)
        else:
            # 正常阀门：视为潜在边界
            boundary_valves[vid] = (n1, n2)

    # 4. BFS 搜索 (叠加 pipe_adj 和 extra_adj)
    affected_nodes: set[str] = set()
    queue = deque(start_nodes)
    while queue:
        node = queue.popleft()
        if node in affected_nodes:
            continue
        affected_nodes.add(node)

        # 遍历永久管道邻居
        if node in pipe_adj:
            for neighbor in pipe_adj[node]:
                if neighbor not in affected_nodes:
                    queue.append(neighbor)

        # 遍历故障阀门带来的额外邻居
        if node in extra_adj:
            for neighbor in extra_adj[node]:
                if neighbor not in affected_nodes:
                    queue.append(neighbor)

    # 5. 结果聚合
    must_close_valves: list[str] = []
    optional_valves: list[str] = []

    for valve_id, (n1, n2) in boundary_valves.items():
        in_n1 = n1 in affected_nodes
        in_n2 = n2 in affected_nodes
        if in_n1 and in_n2:
            optional_valves.append(valve_id)
        elif in_n1 or in_n2:
            must_close_valves.append(valve_id)

    must_close_valves.sort()
    optional_valves.sort()

    result = {
        "accident_elements": target_elements,
        "disabled_valves": disabled_valves,
        "affected_nodes": sorted(affected_nodes),
        "must_close_valves": must_close_valves,
        "optional_valves": optional_valves,
        "isolatable": len(must_close_valves) > 0,
    }

    if len(target_elements) == 1:
        result["accident_element"] = target_elements[0]

    return result