"""Step 2 — 기준 수리해석, Step 9 — 분할안 수리적 검증."""

from __future__ import annotations

import copy
import logging
import math
from dataclasses import dataclass, field
from typing import Optional

import networkx as nx
import wntr

from .config import Config


def _run_epanet(wn):
    """Run wntr's EPANET simulator in a private temp dir.

    EpanetSimulator writes temp.inp/temp.rpt/temp.bin into the *current
    directory* by default, which left stray files in backend/ on every run.
    """
    import os
    import tempfile

    with tempfile.TemporaryDirectory(prefix="spipe_blocksys_") as tmp:
        return wntr.sim.EpanetSimulator(wn).run_sim(file_prefix=os.path.join(tmp, "sim"))

logger = logging.getLogger(__name__)


@dataclass
class BaselineResult:
    node_pressure_mean: dict[str, float]
    node_pressure_min: dict[str, float]
    link_flow_mean_abs: dict[str, float]
    link_flow_max_abs: dict[str, float]
    node_demand_mean: dict[str, float]
    node_age_max: dict[str, float]  # hours


@dataclass
class ValidationResult:
    pressure_low_violations: list[str] = field(default_factory=list)
    pressure_high_violations: list[str] = field(default_factory=list)
    isolated_nodes: list[str] = field(default_factory=list)
    water_age_violations: list[str] = field(default_factory=list)
    min_pressure_m: float = math.nan
    max_water_age_h: float = math.nan
    pressure_drop: dict[str, float] = field(default_factory=dict)
    status: str = "ok"  # ok | violated | rejected
    rejection_reason: Optional[str] = None

    @property
    def is_ok(self) -> bool:
        return self.status == "ok"


def _configure_time(wn: "wntr.network.WaterNetworkModel", cfg: Config) -> None:
    duration = cfg.hydraulics.duration_hours * 3600
    step = cfg.hydraulics.timestep_minutes * 60
    wn.options.time.duration = duration
    wn.options.time.hydraulic_timestep = step
    wn.options.time.report_timestep = step
    wn.options.time.quality_timestep = min(wn.options.time.quality_timestep, step)


def baseline_simulation(wn: "wntr.network.WaterNetworkModel", cfg: Config) -> BaselineResult:
    wn_sim = copy.deepcopy(wn)
    _configure_time(wn_sim, cfg)

    results = _run_epanet(wn_sim)
    pressure = results.node["pressure"]
    flow = results.link["flowrate"].abs()
    demand = results.node["demand"]

    return BaselineResult(
        node_pressure_mean=pressure.mean().to_dict(),
        node_pressure_min=pressure.min().to_dict(),
        link_flow_mean_abs=flow.mean().to_dict(),
        link_flow_max_abs=flow.max().to_dict(),
        node_demand_mean=demand.mean().to_dict(),
        node_age_max=_run_water_age(wn, cfg),
    )


def _run_water_age(wn: "wntr.network.WaterNetworkModel", cfg: Config) -> dict[str, float]:
    wn_age = copy.deepcopy(wn)
    _configure_time(wn_age, cfg)
    wn_age.options.quality.parameter = "AGE"

    results = _run_epanet(wn_age)
    age_hours = results.node["quality"] / 3600.0  # EPANET AGE는 초 단위(SI)로 보고됨
    return age_hours.max().to_dict()


def _open_link_graph(wn: "wntr.network.WaterNetworkModel") -> nx.Graph:
    G = nx.Graph()
    G.add_nodes_from(wn.node_name_list)
    for _, link in wn.links():
        if link.initial_status == wntr.network.LinkStatus.Closed:
            continue
        G.add_edge(link.start_node_name, link.end_node_name)
    return G


def _find_isolated_nodes(G: nx.Graph, source_names: list[str]) -> list[str]:
    sources = [s for s in source_names if s in G]
    if not sources:
        return []
    reachable: set[str] = set()
    for source in sources:
        reachable |= nx.node_connected_component(G, source)
    return sorted(n for n in G.nodes if n not in reachable)


def hydraulic_validation(
    wn_blocked: "wntr.network.WaterNetworkModel",
    baseline: BaselineResult,
    cfg: Config,
) -> ValidationResult:
    source_names = list(wn_blocked.reservoir_name_list) + list(wn_blocked.tank_name_list)
    isolated = _find_isolated_nodes(_open_link_graph(wn_blocked), source_names)

    if isolated and not cfg.hydraulics.allow_isolated_nodes:
        return ValidationResult(
            isolated_nodes=isolated,
            status="rejected",
            rejection_reason=f"{len(isolated)}개 노드가 수원으로부터 고립되었습니다 (구축절차 ④ 위반)",
        )

    try:
        wn_sim = copy.deepcopy(wn_blocked)
        _configure_time(wn_sim, cfg)
        results = _run_epanet(wn_sim)
        pressure = results.node["pressure"]
        age_max = _run_water_age(wn_blocked, cfg)
    except Exception as exc:  # EPANET 엔진 수렴 실패 등
        logger.warning("수리해석 실패: %s", exc)
        return ValidationResult(
            isolated_nodes=isolated,
            status="rejected",
            rejection_reason=f"수리해석 실패: {exc}",
        )

    pressure_min = pressure.min()
    pressure_max = pressure.max()

    pressure_low = sorted(pressure_min[pressure_min < cfg.hydraulics.min_pressure_m].index)
    pressure_high = sorted(pressure_max[pressure_max > cfg.hydraulics.max_pressure_m].index)
    age_violations = sorted(n for n, v in age_max.items() if v > cfg.hydraulics.max_water_age_hours)

    pressure_min_dict = pressure_min.to_dict()
    pressure_drop = {
        n: baseline.node_pressure_min[n] - pressure_min_dict[n]
        for n in pressure_min_dict
        if n in baseline.node_pressure_min
    }

    violated = bool(pressure_low or pressure_high or age_violations or (isolated and cfg.hydraulics.allow_isolated_nodes))

    return ValidationResult(
        pressure_low_violations=pressure_low,
        pressure_high_violations=pressure_high,
        isolated_nodes=isolated,
        water_age_violations=age_violations,
        min_pressure_m=float(pressure_min.min()) if len(pressure_min) else math.nan,
        max_water_age_h=max(age_max.values()) if age_max else math.nan,
        pressure_drop=pressure_drop,
        status="violated" if violated else "ok",
    )