import ctypes
import platform
import os
import sys
import json
import base64
from datetime import datetime
import subprocess
import logging
from typing import Any
import uuid

sys.path.append("..")
from app.native.wndb import project
from app.native.wndb import inp_out


def _verify_platform():
    _platform = platform.system()
    if _platform not in ["Windows", "Linux"]:
        raise Exception(f"Platform {_platform} unsupported (not yet)")


if __name__ == "__main__":
    _verify_platform()


class Output:
    def __init__(self, path: str) -> None:
        self._path = path

        if platform.system() == "Windows":
            self._lib = ctypes.CDLL(
                os.path.join(os.path.dirname(__file__), "windows", "epanet-output.dll")
            )
        else:
            self._lib = ctypes.CDLL(
                os.path.join(os.path.dirname(__file__), "linux", "libepanet-output.so")
            )

        self._handle = ctypes.c_void_p()
        self._check(self._lib.ENR_init(ctypes.byref(self._handle)))

        self._check(
            self._lib.ENR_open(self._handle, ctypes.c_char_p(self._path.encode()))
        )

    def __del__(self):
        # throw exception in destructor ? :)
        self._check(self._lib.ENR_close(ctypes.byref(self._handle)))

    def _check(self, result):
        if result != 0 and result != 10:
            msg = ctypes.c_char_p()
            code = self._lib.ENR_checkError(self._handle, ctypes.byref(msg))
            assert code == result

            error = f"Failed to read project [{self._path}] output, message [{msg.value.decode()}]"

            self._lib.ENR_free(ctypes.byref(msg))

            raise Exception(error)

    def version(self) -> int:
        v = ctypes.c_int()
        self._check(self._lib.ENR_getVersion(self._handle, ctypes.byref(v)))
        return v.value

    def net_size(self) -> dict[str, int]:
        element_count = ctypes.POINTER(ctypes.c_int)()
        length = ctypes.c_int()
        self._check(
            self._lib.ENR_getNetSize(
                self._handle, ctypes.byref(element_count), ctypes.byref(length)
            )
        )
        assert length.value == 5
        category = ["node", "tank", "link", "pump", "valve"]
        sizes = {}
        for i in range(length.value):
            sizes[category[i]] = element_count[i]
        self._lib.ENR_free(ctypes.byref(element_count))
        return sizes

    def units(self) -> dict[str, str]:
        f_us = ["CFS", "GPM", "MGD", "IMGD", "AFD", "LPS", "LPM", "MLD", "CMH", "CMD"]
        p_us = ["PSI", "MTR", "KPA"]
        q_us = ["NONE", "MGL", "UGL", "HOURS", "PRCNT"]
        f, p, q = ctypes.c_int(1), ctypes.c_int(2), ctypes.c_int(3)
        f_u, p_u, q_u = ctypes.c_int(), ctypes.c_int(), ctypes.c_int()
        self._check(self._lib.ENR_getUnits(self._handle, f, ctypes.byref(f_u)))
        self._check(self._lib.ENR_getUnits(self._handle, p, ctypes.byref(p_u)))
        self._check(self._lib.ENR_getUnits(self._handle, q, ctypes.byref(q_u)))
        return {
            "flow": f_us[f_u.value],
            "pressure": p_us[p_u.value],
            "quality": q_us[q_u.value],
        }

    def times(self) -> dict[str, int]:
        ts = []
        for i in range(1, 5):
            t = ctypes.c_int(1)
            self._check(
                self._lib.ENR_getTimes(self._handle, ctypes.c_int(i), ctypes.byref(t))
            )
            ts.append(t.value)
        d = {}
        category = ["report_start", "report_step", "sim_duration", "num_periods"]
        for i in range(4):
            d[category[i]] = ts[i]
        return d

    def element_name(self) -> dict[str, list[str]]:
        sizes = self.net_size()

        node_type = ctypes.c_int(1)
        nodes = []
        for i in range(1, sizes["node"] + 1):
            name = ctypes.c_char_p()
            name_len = ctypes.c_int()
            self._check(
                self._lib.ENR_getElementName(
                    self._handle,
                    node_type,
                    ctypes.c_int(i),
                    ctypes.byref(name),
                    ctypes.byref(name_len),
                )
            )
            nodes.append(name.value.decode())
            self._lib.ENR_free(ctypes.byref(name))

        link_type = ctypes.c_int(2)
        links = []
        for i in range(1, sizes["link"] + 1):
            name = ctypes.c_char_p()
            name_len = ctypes.c_int()
            self._check(
                self._lib.ENR_getElementName(
                    self._handle,
                    link_type,
                    ctypes.c_int(i),
                    ctypes.byref(name),
                    ctypes.byref(name_len),
                )
            )
            links.append(name.value.decode())
            self._lib.ENR_free(ctypes.byref(name))

        return {"nodes": nodes, "links": links}

    def energy_usage(self) -> list[dict[str, Any]]:
        size = self.net_size()["pump"]
        usages = []
        category = [
            "utilization",
            "avg.efficiency",
            "avg.kW/flow",
            "avg.kwatts",
            "max.kwatts",
            "cost/day",
        ]
        links = self.element_name()["links"]
        for i in range(1, size + 1):
            index = ctypes.c_int()
            values = ctypes.POINTER(ctypes.c_float)()
            length = ctypes.c_int()
            self._check(
                self._lib.ENR_getEnergyUsage(
                    self._handle,
                    ctypes.c_int(i),
                    ctypes.byref(index),
                    ctypes.byref(values),
                    ctypes.byref(length),
                )
            )
            assert length.value == 6
            d = {"pump": links[index.value - 1]}
            for j in range(length.value):
                d |= {category[j]: values[j]}
            usages.append(d)
            self._lib.ENR_free(ctypes.byref(values))
        return usages

    def reactions(self) -> dict[str, float]:
        values = ctypes.POINTER(ctypes.c_float)()
        length = ctypes.c_int()
        self._check(
            self._lib.ENR_getNetReacts(
                self._handle, ctypes.byref(values), ctypes.byref(length)
            )
        )
        assert length.value == 4
        category = ["bulk", "wall", "tank", "source"]
        d = {}
        for i in range(4):
            d[category[i]] = values[i]
        self._lib.ENR_free(ctypes.byref(values))
        return d

    def node_results(self) -> list[dict[str, Any]]:
        size = self.net_size()["node"]
        num_periods = self.times()["num_periods"]
        nodes = self.element_name()["nodes"]
        category = ["demand", "head", "pressure", "quality"]
        ds = []
        for i in range(1, size + 1):
            d = {"node": nodes[i - 1], "result": []}
            for j in range(num_periods):
                values = ctypes.POINTER(ctypes.c_float)()
                length = ctypes.c_int()
                self._check(
                    self._lib.ENR_getNodeResult(
                        self._handle, j, i, ctypes.byref(values), ctypes.byref(length)
                    )
                )
                assert length.value == len(category)
                attributes = {}
                for k in range(length.value):
                    attributes[category[k]] = values[k]
                d["result"].append(attributes)
                self._lib.ENR_free(ctypes.byref(values))
            ds.append(d)
        return ds

    def link_results(self) -> list[dict[str, Any]]:
        size = self.net_size()["link"]
        num_periods = self.times()["num_periods"]
        links = self.element_name()["links"]
        category = [
            "flow",
            "velocity",
            "headloss",
            "quality",
            "status",
            "setting",
            "reaction",
            "friction",
        ]
        ds = []

        for i in range(1, size + 1):
            d = {"link": links[i - 1], "result": []}
            for j in range(num_periods):
                values = ctypes.POINTER(ctypes.c_float)()
                length = ctypes.c_int()
                self._check(
                    self._lib.ENR_getLinkResult(
                        self._handle, j, i, ctypes.byref(values), ctypes.byref(length)
                    )
                )
                assert length.value == len(category)
                attributes = {}
                for k in range(length.value):
                    if category[k] == "status":
                        if values[k] == 2.0:
                            attributes[category[k]] = "CLOSED"
                        else:
                            attributes[category[k]] = "OPEN"
                        continue
                    attributes[category[k]] = values[k]
                d["result"].append(attributes)
                self._lib.ENR_free(ctypes.byref(values))
            ds.append(d)
        return ds

    def dump(self) -> dict[str, Any]:
        data = {}
        data |= {"version": self.version()}
        data |= {"net_size": self.net_size()}
        data |= {"units": self.units()}
        data |= {"times": self.times()}
        data |= {"element_name": self.element_name()}
        data |= {"energy_usage": self.energy_usage()}
        data |= {"reactions": self.reactions()}
        data |= {"node_results": self.node_results()}
        data |= {"link_results": self.link_results()}
        return data


def _dump_output(path: str) -> dict[str, Any]:
    opt = Output(path)
    data = opt.dump()
    return data


def dump_output(path: str) -> str:
    data = _dump_output(path)
    return json.dumps(data)


def dump_report(path: str) -> str:
    return open(path, "r").read()


def dump_output_binary(path: str) -> str:
    with open(path, "rb") as f:
        data = f.read()
        bast64_data = base64.b64encode(data)
        return str(bast64_data, "utf-8")


def _safe_remove(path: str) -> None:
    try:
        if os.path.exists(path):
            os.remove(path)
    except Exception:
        logging.warning("failed to remove temp file: %s", path)


def _make_isolated_run_paths(base_name: str, cwd: str) -> tuple[str, str, str]:
    # 确保保存临时文件的目录存在
    db_inp_dir = os.path.join(cwd, "db_inp")
    temp_dir = os.path.join(cwd, "temp")
    os.makedirs(db_inp_dir, exist_ok=True)
    os.makedirs(temp_dir, exist_ok=True)

    # 进程号 + UUID 生成唯一后缀，避免并发进程互相覆盖临时文件。
    token = f"{os.getpid()}_{uuid.uuid4().hex}"
    inp = os.path.join(db_inp_dir, f"{base_name}.db.{token}.inp")
    rpt = os.path.join(temp_dir, f"{base_name}.db.{token}.rpt")
    opt = os.path.join(temp_dir, f"{base_name}.db.{token}.opt")
    return inp, rpt, opt


# DingZQ, 2025-02-04, 返回dict[str, Any]
def run_project_return_dict(name: str, readable_output: bool = True) -> dict[str, Any]:
    if not project.have_project(name):
        raise Exception(f"Not found project [{name}]")

    cwd = os.path.abspath(os.getcwd())

    inp, rpt, opt = _make_isolated_run_paths(name, cwd)
    inp_out.dump_inp(name, inp, "2")

    if platform.system() == "Windows":
        exe = os.path.join(os.path.dirname(__file__), "windows", "runepanet.exe")
    else:
        exe = os.path.join(os.path.dirname(__file__), "linux", "runepanet")

    if platform.system() != "Windows":
        if not os.access(exe, os.X_OK):
            os.chmod(exe, 0o755)

    data = {}

    # 设置环境变量以包含库文件路径
    env = os.environ.copy()
    if platform.system() == "Linux":
        lib_dir = os.path.dirname(exe)
        env["LD_LIBRARY_PATH"] = f"{lib_dir}:{env.get('LD_LIBRARY_PATH', '')}"

    process = subprocess.run([exe, inp, rpt, opt], env=env, capture_output=True, text=True)
    result = process.returncode

    if result != 0:
        logging.error(f"EPANET failed with return code {result}")
        logging.error(f"EPANET stdout: {process.stdout}")
        logging.error(f"EPANET stderr: {process.stderr}")
        data["simulation_result"] = "failed"
        data["error_code"] = result
        data["stdout"] = process.stdout
        data["stderr"] = process.stderr
    else:
        data["simulation_result"] = "successful"
        if readable_output:
            data["output"] = _dump_output(opt)
        else:
            data["output"] = dump_output_binary(opt)

    data["input_file"] = inp
    data["report_file"] = rpt
    data["output_file"] = opt
    
    if os.path.exists(rpt):
        data["report"] = dump_report(rpt)
    else:
        logging.error(f"EPANET report file not found: {rpt}")
        data["report"] = f"Error: EPANET report file not found. Simulation return code: {result}. Check server logs for stdout/stderr."

    # 返回内容后删除仿真临时文件，避免临时文件堆积。
    _safe_remove(inp)
    _safe_remove(rpt)
    _safe_remove(opt)

    return data


# original code
def run_project(name: str, readable_output: bool = True) -> str:
    if not project.have_project(name):
        raise Exception(f"Not found project [{name}]")

    cwd = os.path.abspath(os.getcwd())

    inp, rpt, opt = _make_isolated_run_paths(name, cwd)
    inp_out.dump_inp(name, inp, "2")

    if platform.system() == "Windows":
        exe = os.path.join(os.path.dirname(__file__), "windows", "runepanet.exe")
    else:
        exe = os.path.join(os.path.dirname(__file__), "linux", "runepanet")

    logging.info(f"Run simulation at {datetime.now()}")
    logging.info("%s %s %s %s", exe, inp, rpt, opt)

    if platform.system() != "Windows":
        if not os.access(exe, os.X_OK):
            os.chmod(exe, 0o755)

    data = {}

    # 设置环境变量以包含库文件路径
    env = os.environ.copy()
    if platform.system() == "Linux":
        lib_dir = os.path.dirname(exe)
        env["LD_LIBRARY_PATH"] = f"{lib_dir}:{env.get('LD_LIBRARY_PATH', '')}"

    # DingZQ, 2025-06-02, 使用subprocess替代os.system
    process = subprocess.run([exe, inp, rpt, opt], env=env)
    result = process.returncode
    # logging.info(f"Simulation result: {result}")

    if result != 0:
        data["simulation_result"] = "failed"

        logging.error("simulation failed")

    else:
        data["simulation_result"] = "successful"
        logging.info("simulation successful")

        if readable_output:
            data["output"] = _dump_output(opt)
        else:
            data["output"] = dump_output_binary(opt)

    data["input_file"] = inp
    data["report_file"] = rpt
    data["output_file"] = opt
    data["report"] = dump_report(rpt)

    # 返回内容后删除仿真临时文件，避免临时文件堆积。
    _safe_remove(inp)
    _safe_remove(rpt)
    _safe_remove(opt)
    # logging.info(f"Report: {data['report']}")

    return json.dumps(data)


def run_inp(name: str, readable_output: bool = True) -> str:
    cwd = os.path.abspath(os.getcwd())

    if platform.system() == "Windows":
        exe = os.path.join(os.path.dirname(__file__), "windows", "runepanet.exe")
    else:
        exe = os.path.join(os.path.dirname(__file__), "linux", "runepanet")
    source_inp = os.path.join(cwd, "inp", name + ".inp")
    token = f"{os.getpid()}_{uuid.uuid4().hex}"
    rpt = os.path.join(cwd, "temp", f"{name}.{token}.rpt")
    opt = os.path.join(cwd, "temp", f"{name}.{token}.opt")

    if platform.system() != "Windows":
        if not os.access(exe, os.X_OK):
            os.chmod(exe, 0o755)

    data = {}

    # 设置环境变量以包含库文件路径
    env = os.environ.copy()
    if platform.system() == "Linux":
        lib_dir = os.path.dirname(exe)
        env["LD_LIBRARY_PATH"] = f"{lib_dir}:{env.get('LD_LIBRARY_PATH', '')}"

    process = subprocess.run([exe, source_inp, rpt, opt], env=env)
    result = process.returncode

    if result != 0:
        data["simulation_result"] = "failed"
    else:
        data["simulation_result"] = "successful"
        if readable_output:
            data["output"] = _dump_output(opt)
        else:
            data["output"] = dump_output_binary(opt)

    data["input_file"] = source_inp
    data["report_file"] = rpt
    data["output_file"] = opt
    data["report"] = dump_report(rpt)

    # 返回内容后删除仿真临时文件，避免临时文件堆积。
    _safe_remove(source_inp)
    _safe_remove(rpt)
    _safe_remove(opt)

    return json.dumps(data)