"""Step 1 — EPANET 모델 로드 및 단순 그래프 변환."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import networkx as nx
import pandas as pd
import wntr

logger = logging.getLogger(__name__)


@dataclass
class NetworkBundle:
    wn: "wntr.network.WaterNetworkModel"
    graph: nx.Graph
    isolated_nodes: list[str]
    connectivity_warning: Optional[str] = None


def load_network(inp_path: str | Path, connections_csv: Optional[str | Path] = None) -> NetworkBundle:
    wn = wntr.network.WaterNetworkModel(str(inp_path))

    G = _build_simple_graph(wn)
    _inject_node_attributes(wn, G)
    _load_connections(G, connections_csv)

    isolated = [n for n in G.nodes if G.degree(n) == 0]
    if isolated:
        logger.warning("%d개 노드가 어떤 관로에도 연결되어 있지 않습니다: %s", len(isolated), isolated[:10])

    connectivity_warning = None
    if G.number_of_nodes() > 0 and not nx.is_connected(G):
        n_components = nx.number_connected_components(G)
        connectivity_warning = f"네트워크가 {n_components}개의 연결되지 않은 구성요소로 나뉘어 있습니다"
        logger.warning(connectivity_warning)

    return NetworkBundle(wn=wn, graph=G, isolated_nodes=isolated, connectivity_warning=connectivity_warning)


def _build_simple_graph(wn: "wntr.network.WaterNetworkModel") -> nx.Graph:
    multi = wn.to_graph()  # MultiDiGraph, 노드 좌표 등 기본 속성 포함

    G = nx.Graph()
    G.add_nodes_from(multi.nodes)

    grouped: dict[tuple[str, str], list[str]] = {}
    for u, v, key in multi.edges(keys=True):
        pair = (u, v) if u <= v else (v, u)
        grouped.setdefault(pair, []).append(key)

    for (u, v), link_names in grouped.items():
        lengths, diameters, roughness = [], [], []
        for name in link_names:
            link = wn.get_link(name)
            lengths.append(getattr(link, "length", 0.0) or 0.0)
            diameters.append(getattr(link, "diameter", 0.0) or 0.0)
            roughness.append(getattr(link, "roughness", 0.0) or 0.0)

        if len(link_names) == 1:
            length, diameter, rough = lengths[0], diameters[0], roughness[0]
        else:
            # 병렬관 등가화: 단면적(D^2) 합산 기준 등가관경, 연장은 최단 경로 기준 최소값.
            # 실제 수리해석은 원본 wn(EPANET 엔진)으로 수행하므로 여기서는 분할 가중치용 근사치.
            length = min(lengths) if lengths else 0.0
            diameter = (sum(d**2 for d in diameters)) ** 0.5 if diameters else 0.0
            rough = max(roughness) if roughness else 0.0

        G.add_edge(u, v, length=length, diameter=diameter, roughness=rough, pipe_ids=list(link_names))

    return G


def _inject_node_attributes(wn: "wntr.network.WaterNetworkModel", G: nx.Graph) -> None:
    for node_name in G.nodes:
        node = wn.get_node(node_name)

        elevation = getattr(node, "elevation", None)
        if elevation is None:
            elevation = getattr(node, "base_head", 0.0)  # Reservoir

        base_demand = 0.0
        for demand in getattr(node, "demand_timeseries_list", []) or []:
            base_demand += demand.base_value

        G.nodes[node_name]["elevation"] = float(elevation or 0.0)
        G.nodes[node_name]["base_demand"] = float(base_demand)
        G.nodes[node_name]["coordinates"] = node.coordinates
        G.nodes[node_name]["node_type"] = type(node).__name__
        G.nodes[node_name]["connections"] = None
        G.nodes[node_name]["is_large_consumer_csv"] = False


def _load_connections(G: nx.Graph, connections_csv: Optional[str | Path]) -> None:
    if not connections_csv:
        logger.warning(
            "connections.csv가 제공되지 않았습니다. 급수전수 기반 제약"
            "(constraints.connections_per_block)은 이번 실행에서 평가되지 않습니다."
        )
        return

    df = pd.read_csv(connections_csv, dtype={"node_id": str})
    missing = {"node_id", "connections"} - set(df.columns)
    if missing:
        raise ValueError(f"connections.csv에 필수 컬럼이 없습니다: {sorted(missing)}")

    unknown = 0
    for _, row in df.iterrows():
        node_id = str(row["node_id"])
        if node_id not in G.nodes:
            unknown += 1
            continue
        G.nodes[node_id]["connections"] = float(row["connections"])
        if "is_large_consumer" in df.columns and not pd.isna(row["is_large_consumer"]):
            G.nodes[node_id]["is_large_consumer_csv"] = bool(row["is_large_consumer"])

    if unknown:
        logger.warning("connections.csv의 %d개 행이 네트워크 노드와 일치하지 않아 무시되었습니다", unknown)