"""Repeated-solve helper for batch analyses (fire flow, criticality, VSP,
calibration ...).

`Network` opens one INP with the EPANET 2.3 DLL and keeps it open, so an
analysis can change values (demands, statuses, settings, controls) and
re-solve hydraulics many times without re-reading the file. Results are
returned at the reporting periods only ([TIMES] REPORT START / REPORT STEP),
the same periods the main /api/analyze endpoint reports, as numpy arrays.

Runs inside the API process pool (one Network per call, never shared).
"""
from __future__ import annotations

import ctypes
import os
import tempfile
from dataclasses import dataclass, field

import numpy as np

from .epanet import EN, EpanetAPI, EpanetException

_api: EpanetAPI | None = None


def get_api() -> EpanetAPI:
    """One DLL handle per process (ctypes handles must not cross processes)."""
    global _api
    if _api is None:
        _api = EpanetAPI()
    return _api


class BatchError(Exception):
    """User-facing input problem (unknown ID, bad parameter)."""


@dataclass
class SimResult:
    times: list[int]
    nodes: dict[int, np.ndarray] = field(default_factory=dict)  # prop -> [period, node]
    links: dict[int, np.ndarray] = field(default_factory=dict)  # prop -> [period, link]
    warnings: list[tuple[int, int]] = field(default_factory=list)  # (time, code)


def _decode(raw: bytes) -> str:
    try:
        return raw.decode("utf-8")
    except UnicodeDecodeError:
        return raw.decode("cp949", errors="replace")


class Network:
    def __init__(self, inp_bytes: bytes):
        self.api = get_api()
        self.lib = self.api.lib
        self._tmp = tempfile.TemporaryDirectory()
        inp = os.path.join(self._tmp.name, "n.inp")
        with open(inp, "wb") as f:
            f.write(inp_bytes)
        self.ph = self.api.EN_Project()
        self.lib.EN_createproject(ctypes.byref(self.ph))
        self._opened = False
        try:
            self._opened = True
            self.lib.EN_open(self.ph, inp.encode(), os.path.join(self._tmp.name, "n.rpt").encode(), b"")
        except EpanetException as e:
            self.close()
            raise BatchError(str(e)) from e

        n, l = ctypes.c_int(), ctypes.c_int()
        self.lib.EN_getcount(self.ph, EN.NODECOUNT, ctypes.byref(n))
        self.lib.EN_getcount(self.ph, EN.LINKCOUNT, ctypes.byref(l))
        self.node_count, self.link_count = n.value, l.value
        buf = ctypes.create_string_buffer(256)
        t = ctypes.c_int()
        self.node_ids: list[str] = []
        self.node_types: list[int] = []
        for i in range(1, self.node_count + 1):
            self.lib.EN_getnodeid(self.ph, i, buf)
            self.node_ids.append(_decode(buf.value))
            self.lib.EN_getnodetype(self.ph, i, ctypes.byref(t))
            self.node_types.append(t.value)
        self.link_ids: list[str] = []
        self.link_types: list[int] = []
        self.link_nodes: list[tuple[int, int]] = []
        a, b = ctypes.c_int(), ctypes.c_int()
        for i in range(1, self.link_count + 1):
            self.lib.EN_getlinkid(self.ph, i, buf)
            self.link_ids.append(_decode(buf.value))
            self.lib.EN_getlinktype(self.ph, i, ctypes.byref(t))
            self.link_types.append(t.value)
            self.lib.EN_getlinknodes(self.ph, i, ctypes.byref(a), ctypes.byref(b))
            self.link_nodes.append((a.value, b.value))
        self._node_index = {nid: i + 1 for i, nid in enumerate(self.node_ids)}
        self._link_index = {lid: i + 1 for i, lid in enumerate(self.link_ids)}

        self.duration = self.time_param(EN.DURATION)
        self.hyd_step = self.time_param(EN.HYDSTEP)
        self.report_step = max(1, self.time_param(EN.REPORTSTEP))
        self.report_start = self.time_param(EN.REPORTSTART)
        self.pattern_step = max(1, self.time_param(EN.PATTERNSTEP))

    # ---- lifecycle ----
    def close(self) -> None:
        if getattr(self, "ph", None) is None:
            return
        try:
            if self._opened:
                self.lib.EN_close(self.ph)
        except EpanetException:
            pass
        try:
            self.lib.EN_deleteproject(self.ph)
        except EpanetException:
            pass
        self.ph = None
        self._tmp.cleanup()

    def __enter__(self) -> "Network":
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # ---- lookups ----
    def node_index(self, node_id: str) -> int:
        try:
            return self._node_index[node_id]
        except KeyError:
            raise BatchError(f"절점 {node_id}가 모델에 없습니다.") from None

    def link_index(self, link_id: str) -> int:
        try:
            return self._link_index[link_id]
        except KeyError:
            raise BatchError(f"링크 {link_id}가 모델에 없습니다.") from None

    def junction_indices(self) -> list[int]:
        return [i + 1 for i, t in enumerate(self.node_types) if t == EN.JUNCTION]

    # ---- values ----
    def time_param(self, code: int) -> int:
        v = ctypes.c_long()
        self.lib.EN_gettimeparam(self.ph, code, ctypes.byref(v))
        return v.value

    def set_time_param(self, code: int, value: int) -> None:
        self.lib.EN_settimeparam(self.ph, code, int(value))

    def get_node(self, index: int, prop: int) -> float:
        v = ctypes.c_double()
        self.lib.EN_getnodevalue(self.ph, index, prop, ctypes.byref(v))
        return v.value

    def set_node(self, index: int, prop: int, value: float) -> None:
        self.lib.EN_setnodevalue(self.ph, index, prop, float(value))

    def get_link(self, index: int, prop: int) -> float:
        v = ctypes.c_double()
        self.lib.EN_getlinkvalue(self.ph, index, prop, ctypes.byref(v))
        return v.value

    def set_link(self, index: int, prop: int, value: float) -> None:
        self.lib.EN_setlinkvalue(self.ph, index, prop, float(value))

    def add_pattern(self, pid: str, values: list[float]) -> None:
        """Create the pattern, or overwrite its values if the ID already exists."""
        idx = ctypes.c_int()
        try:
            self.lib.EN_getpatternindex(self.ph, pid.encode(), ctypes.byref(idx))
        except EpanetException:
            self.lib.EN_addpattern(self.ph, pid.encode())
            self.lib.EN_getpatternindex(self.ph, pid.encode(), ctypes.byref(idx))
        arr = (ctypes.c_double * len(values))(*values)
        self.lib.EN_setpattern(self.ph, idx.value, arr, len(values))

    def window_pattern(self, pid: str, start_s: int, end_s: int) -> None:
        """Pattern that is 1 on [start, end) and 0 elsewhere, over the duration."""
        # EPANET uses period (t + PATTERN START) // PATTERN STEP, so period k
        # starts at simulation time k * step - pattern_start.
        pstart = self.time_param(EN.PATTERNSTART)
        n = max(1, -(-(max(self.duration, 1) + pstart) // self.pattern_step) + 1)
        vals = []
        for k in range(n):
            t0 = k * self.pattern_step - pstart
            vals.append(1.0 if start_s <= t0 < end_s else 0.0)
        self.add_pattern(pid, vals)

    def add_demand(self, index: int, base: float, pattern: str = "", name: str = "") -> int:
        """Adds a demand category; returns its 1-based category index."""
        self.lib.EN_adddemand(self.ph, index, float(base), pattern.encode(), name.encode())
        k = ctypes.c_int()
        self.lib.EN_getnumdemands(self.ph, index, ctypes.byref(k))
        return k.value

    def delete_demand(self, index: int, category: int) -> None:
        self.lib.EN_deletedemand(self.ph, index, category)

    def add_timer_control(self, link_index: int, setting: float, at_s: int) -> int:
        """setting: 0 = CLOSED, 1 = OPEN (pipes) or a speed/valve setting."""
        out = ctypes.c_int()
        self.lib.EN_addcontrol(self.ph, EN.TIMER, link_index, float(setting), 0, float(at_s), ctypes.byref(out))
        return out.value

    def delete_control(self, index: int) -> None:
        self.lib.EN_deletecontrol(self.ph, index)

    # ---- solve ----
    def run_quality(self, nodes: list[int] | None = None) -> tuple[list[int], np.ndarray]:
        """Quality run on hydraulics saved by EN_solveH; QUALITY at report periods [period, node]."""
        node_idx = nodes if nodes is not None else list(range(1, self.node_count + 1))
        t, tstep = ctypes.c_long(), ctypes.c_long(1)
        v = ctypes.c_double()
        times: list[int] = []
        rows: list[list[float]] = []
        ph, get_n = self.ph, self.lib.EN_getnodevalue
        self.lib.EN_openQ(ph)
        try:
            self.lib.EN_initQ(ph, EN.NOSAVE)
            while True:
                self.lib.EN_runQ(ph, ctypes.byref(t))
                if self.is_report_time(t.value) and (not times or times[-1] != t.value):
                    times.append(t.value)
                    row = []
                    for i in node_idx:
                        get_n(ph, i, EN.QUALITY, ctypes.byref(v))
                        row.append(v.value)
                    rows.append(row)
                self.lib.EN_nextQ(ph, ctypes.byref(tstep))
                if tstep.value <= 0:
                    break
        finally:
            self.lib.EN_closeQ(ph)
        return times, np.asarray(rows, dtype=float).reshape(len(times), len(node_idx))

    def pattern_value(self, pattern_index: int, t: int, pattern_start: int = 0) -> float:
        """Multiplier of a pattern at simulation time t (1.0 without a pattern)."""
        if pattern_index <= 0:
            return 1.0
        n = ctypes.c_int()
        self.lib.EN_getpatternlen(self.ph, pattern_index, ctypes.byref(n))
        if n.value <= 0:
            return 1.0
        period = ((t + pattern_start) // self.pattern_step) % n.value
        v = ctypes.c_double()
        self.lib.EN_getpatternvalue(self.ph, pattern_index, period + 1, ctypes.byref(v))
        return v.value

    def option(self, code: int) -> float:
        v = ctypes.c_double()
        self.lib.EN_getoption(self.ph, code, ctypes.byref(v))
        return v.value

    def is_report_time(self, t: int) -> bool:
        return t >= self.report_start and (t - self.report_start) % self.report_step == 0

    def simulate(
        self,
        node_props: tuple[int, ...] = (EN.PRESSURE,),
        link_props: tuple[int, ...] = (),
        nodes: list[int] | None = None,
        links: list[int] | None = None,
        max_periods: int = 20000,
    ) -> SimResult:
        """Run an EPS and record the given properties at reporting periods.

        `nodes`/`links` are 1-based indices (default: all). Arrays are
        [period, element] in the order of the given index lists.
        """
        node_idx = nodes if nodes is not None else list(range(1, self.node_count + 1))
        link_idx = links if links is not None else list(range(1, self.link_count + 1))
        rows_n: dict[int, list[list[float]]] = {p: [] for p in node_props}
        rows_l: dict[int, list[list[float]]] = {p: [] for p in link_props}
        times: list[int] = []
        warnings: list[tuple[int, int]] = []
        t, tstep = ctypes.c_long(), ctypes.c_long(1)
        v = ctypes.c_double()
        get_n, get_l, ph = self.lib.EN_getnodevalue, self.lib.EN_getlinkvalue, self.ph
        self.lib.EN_openH(ph)
        try:
            self.lib.EN_initH(ph, EN.NOSAVE)
            while True:
                code = self.lib.EN_runH(ph, ctypes.byref(t))
                if 0 < code <= 10:
                    warnings.append((t.value, code))
                if self.is_report_time(t.value):
                    times.append(t.value)
                    for p in node_props:
                        row = []
                        for i in node_idx:
                            get_n(ph, i, p, ctypes.byref(v))
                            row.append(v.value)
                        rows_n[p].append(row)
                    for p in link_props:
                        row = []
                        for i in link_idx:
                            get_l(ph, i, p, ctypes.byref(v))
                            row.append(v.value)
                        rows_l[p].append(row)
                    if len(times) > max_periods:
                        raise BatchError(f"보고 시간대가 {max_periods}개를 넘습니다.")
                self.lib.EN_nextH(ph, ctypes.byref(tstep))
                if tstep.value <= 0:
                    break
        finally:
            self.lib.EN_closeH(ph)
        return SimResult(
            times=times,
            nodes={p: np.asarray(r, dtype=float).reshape(len(times), len(node_idx)) for p, r in rows_n.items()},
            links={p: np.asarray(r, dtype=float).reshape(len(times), len(link_idx)) for p, r in rows_l.items()},
            warnings=warnings,
        )


def hhmm(seconds: int) -> str:
    return f"{seconds // 3600}:{(seconds % 3600) // 60:02d}"


def parse_time(value, default: int = 0) -> int:
    """Seconds from a number of hours, 'h:mm', or None."""
    if value is None or value == "":
        return default
    if isinstance(value, (int, float)):
        return int(round(float(value) * 3600))
    s = str(value).strip()
    if ":" in s:
        parts = [float(x) for x in s.split(":")]
        while len(parts) < 3:
            parts.append(0.0)
        return int(round(parts[0] * 3600 + parts[1] * 60 + parts[2]))
    return int(round(float(s) * 3600))