"""Step 7 — 경계 관로 추출 및 밸브/유량계 배정, Step 8 — 모델에 반영."""

from __future__ import annotations

import copy
from dataclasses import dataclass

import networkx as nx
import wntr

from .validate import BaselineResult


@dataclass
class BoundaryLink:
    pipe_id: str
    from_node: str
    to_node: str
    block_a: int
    block_b: int
    action: str  # 'close' | 'meter'
    baseline_flow: float
    diameter: float


def find_boundary_links(G: nx.Graph, labels: dict[str, int]) -> list[BoundaryLink]:
    """경계 엣지 = 양 끝 노드의 블록 라벨이 다른 엣지. 병렬관은 개별 pipe_id로 전개."""
    boundary: list[BoundaryLink] = []
    for u, v, data in G.edges(data=True):
        if u not in labels or v not in labels:
            continue
        block_a, block_b = labels[u], labels[v]
        if block_a == block_b:
            continue
        for pipe_id in data.get("pipe_ids", []):
            boundary.append(
                BoundaryLink(
                    pipe_id=pipe_id,
                    from_node=u,
                    to_node=v,
                    block_a=block_a,
                    block_b=block_b,
                    action="close",
                    baseline_flow=0.0,
                    diameter=data.get("diameter", 0.0),
                )
            )
    return boundary


def assign_valve_or_meter(
    boundary: list[BoundaryLink],
    baseline: BaselineResult,
    inlets_per_block: int = 1,
) -> list[BoundaryLink]:
    """블록별 유입 경로: baseline 유량 상위 inlets_per_block개는 유량계, 나머지는 차단."""
    for link in boundary:
        link.baseline_flow = baseline.link_flow_mean_abs.get(link.pipe_id, 0.0)

    by_block: dict[int, list[BoundaryLink]] = {}
    for link in boundary:
        by_block.setdefault(link.block_a, []).append(link)
        by_block.setdefault(link.block_b, []).append(link)

    meter_pipe_ids: set[str] = set()
    for links in by_block.values():
        ranked = sorted(links, key=lambda l: l.baseline_flow, reverse=True)
        meter_pipe_ids.update(l.pipe_id for l in ranked[:inlets_per_block])

    for link in boundary:
        link.action = "meter" if link.pipe_id in meter_pipe_ids else "close"

    return boundary


def apply_blocks_to_model(
    wn: "wntr.network.WaterNetworkModel", boundary: list[BoundaryLink]
) -> "wntr.network.WaterNetworkModel":
    wn_blocked = copy.deepcopy(wn)
    for link in boundary:
        if link.action == "close":
            wn_blocked.get_link(link.pipe_id).initial_status = "Closed"
    return wn_blocked