import ctypes
import os


class EN:
    # EN_CountType
    NODECOUNT = 0
    LINKCOUNT = 2

    # EN_NodeProperty
    ELEVATION = 0
    BASEDEMAND = 1
    EMITTER = 3
    TANKLEVEL = 8
    DEMAND = 9
    HEAD = 10
    PRESSURE = 11
    QUALITY = 12
    # EN_LinkProperty
    DIAMETER = 0
    LENGTH = 1
    ROUGHNESS = 2
    INITSTATUS = 4
    INITSETTING = 5
    KWALL = 7
    STATUS = 11
    SETTING = 12
    FLOW = 8
    VELOCITY = 9
    HEADLOSS = 10
    ENERGY = 13
    LINKQUAL = 14
    PUMP_EFFIC = 17
    # EN_InitHydOption
    NOSAVE = 0
    SAVE = 1
    
    # EN_Option
    PRESS_UNITS = 25

    # EN_TimeParameter
    DURATION = 0
    HYDSTEP = 1
    REPORTSTEP = 5
    REPORTSTART = 6

    # EN_TimeParameter (cont.)
    PATTERNSTEP = 3
    PATTERNSTART = 4

    # EN_ControlType
    TIMER = 2

    # EN_QualityType / quality time step
    QUAL_TRACE = 3
    QUALSTEP = 2

    # EN_Option (energy)
    GLOBALEFFIC = 8
    GLOBALPRICE = 9
    GLOBALPATTERN = 10
    DEMANDCHARGE = 11

    # EN_LinkProperty (pump energy)
    PUMP_ECOST = 21
    PUMP_EPAT = 22

    # EN_NodeType / EN_LinkType
    JUNCTION = 0
    RESERVOIR = 1
    TANK = 2
    CVPIPE = 0
    PIPE = 1
    PUMP = 2

FLOW_UNITS = {
    0: "CFS", 1: "GPM", 2: "MGD", 3: "IMGD", 4: "AFD",
    5: "LPS", 6: "LPM", 7: "MLD", 8: "CMH", 9: "CMD", 10: "CMS"
}

PRESS_UNITS = {
    0: "psi", 1: "kPa", 2: "m", 3: "bar", 4: "ft"
}
QUAL_TYPES = {
    0: "NONE", 1: "CHEMICAL", 2: "AGE", 3: "TRACE"
}


class EpanetException(Exception):
    pass

class EpanetAPI:
    def __init__(self):
        library_name = "epanet2.dll" if os.name == "nt" else "libepanet2.so"
        dll_path = os.environ.get("EPANET_LIBRARY_PATH") or os.path.join(os.path.dirname(os.path.dirname(__file__)), "64bit", library_name)
        if not os.path.exists(dll_path):
            raise RuntimeError(f"DLL not found: {dll_path}")
        
        self.lib = ctypes.CDLL(dll_path)
        self.EN_Project = ctypes.c_void_p
        self._setup()
        
    def _setup(self):
        def _errcheck(result, func, args):
            if result > 10:
                err_msg = ctypes.create_string_buffer(256)
                self.lib.EN_geterror(result, err_msg, 256)
                raise EpanetException(f"EPANET Error {result}: {err_msg.value.decode('utf-8').strip()}")
            return result
            
        self.lib.EN_geterror.argtypes = [ctypes.c_int, ctypes.c_char_p, ctypes.c_int]
        self.lib.EN_geterror.restype = ctypes.c_int
        
        self.lib.EN_getversion.argtypes = [ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getversion.restype = ctypes.c_int
        self.lib.EN_getversion.errcheck = _errcheck
        
        self.lib.EN_createproject.argtypes = [ctypes.POINTER(self.EN_Project)]
        self.lib.EN_createproject.restype = ctypes.c_int
        self.lib.EN_createproject.errcheck = _errcheck
        
        self.lib.EN_deleteproject.argtypes = [self.EN_Project]
        self.lib.EN_deleteproject.restype = ctypes.c_int
        self.lib.EN_deleteproject.errcheck = _errcheck
        
        self.lib.EN_open.argtypes = [self.EN_Project, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_char_p]
        self.lib.EN_open.restype = ctypes.c_int
        self.lib.EN_open.errcheck = _errcheck
        
        self.lib.EN_close.argtypes = [self.EN_Project]
        self.lib.EN_close.restype = ctypes.c_int
        self.lib.EN_close.errcheck = _errcheck
        
        self.lib.EN_getcount.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getcount.restype = ctypes.c_int
        self.lib.EN_getcount.errcheck = _errcheck
        
        self.lib.EN_getnodeid.argtypes = [self.EN_Project, ctypes.c_int, ctypes.c_char_p]
        self.lib.EN_getnodeid.restype = ctypes.c_int
        self.lib.EN_getnodeid.errcheck = _errcheck
        
        self.lib.EN_getnodetype.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getnodetype.restype = ctypes.c_int
        self.lib.EN_getnodetype.errcheck = _errcheck
        
        self.lib.EN_getnodevalue.argtypes = [self.EN_Project, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
        self.lib.EN_getnodevalue.restype = ctypes.c_int
        self.lib.EN_getnodevalue.errcheck = _errcheck
        
        self.lib.EN_getcoord.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
        self.lib.EN_getcoord.restype = ctypes.c_int
        self.lib.EN_getcoord.errcheck = _errcheck
        
        self.lib.EN_getlinkid.argtypes = [self.EN_Project, ctypes.c_int, ctypes.c_char_p]
        self.lib.EN_getlinkid.restype = ctypes.c_int
        self.lib.EN_getlinkid.errcheck = _errcheck
        
        self.lib.EN_getlinktype.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getlinktype.restype = ctypes.c_int
        self.lib.EN_getlinktype.errcheck = _errcheck
        
        self.lib.EN_getlinknodes.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getlinknodes.restype = ctypes.c_int
        self.lib.EN_getlinknodes.errcheck = _errcheck
        
        self.lib.EN_getlinkvalue.argtypes = [self.EN_Project, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
        self.lib.EN_getlinkvalue.restype = ctypes.c_int
        self.lib.EN_getlinkvalue.errcheck = _errcheck
        
        self.lib.EN_getvertexcount.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getvertexcount.restype = ctypes.c_int
        self.lib.EN_getvertexcount.errcheck = _errcheck
        
        self.lib.EN_getvertex.argtypes = [self.EN_Project, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
        self.lib.EN_getvertex.restype = ctypes.c_int
        self.lib.EN_getvertex.errcheck = _errcheck
        
        self.lib.EN_getflowunits.argtypes = [self.EN_Project, ctypes.POINTER(ctypes.c_int)]
        self.lib.EN_getflowunits.restype = ctypes.c_int
        self.lib.EN_getflowunits.errcheck = _errcheck
        
        self.lib.EN_getoption.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
        self.lib.EN_getoption.restype = ctypes.c_int
        self.lib.EN_getoption.errcheck = _errcheck
        
        self.lib.EN_gettimeparam.argtypes = [self.EN_Project, ctypes.c_int, ctypes.POINTER(ctypes.c_long)]
        self.lib.EN_gettimeparam.restype = ctypes.c_int
        self.lib.EN_gettimeparam.errcheck = _errcheck

        self.lib.EN_openH.argtypes = [self.EN_Project]
        self.lib.EN_openH.restype = ctypes.c_int
        self.lib.EN_openH.errcheck = _errcheck
        
        self.lib.EN_initH.argtypes = [self.EN_Project, ctypes.c_int]
        self.lib.EN_initH.restype = ctypes.c_int
        self.lib.EN_initH.errcheck = _errcheck
        
        self.lib.EN_runH.argtypes = [self.EN_Project, ctypes.POINTER(ctypes.c_long)]
        self.lib.EN_runH.restype = ctypes.c_int
        self.lib.EN_runH.errcheck = _errcheck
        

        self.lib.EN_nextH.argtypes = [self.EN_Project, ctypes.POINTER(ctypes.c_long)]
        self.lib.EN_nextH.restype = ctypes.c_int
        self.lib.EN_nextH.errcheck = _errcheck
        self.lib.EN_closeH.argtypes = [self.EN_Project]
        self.lib.EN_closeH.restype = ctypes.c_int
        self.lib.EN_closeH.errcheck = _errcheck

        # Water Quality Analysis Bindings
        self.lib.EN_openQ.argtypes = [self.EN_Project]
        self.lib.EN_openQ.restype = ctypes.c_int
        self.lib.EN_openQ.errcheck = _errcheck

        self.lib.EN_initQ.argtypes = [self.EN_Project, ctypes.c_int]
        self.lib.EN_initQ.restype = ctypes.c_int
        self.lib.EN_initQ.errcheck = _errcheck

        self.lib.EN_runQ.argtypes = [self.EN_Project, ctypes.POINTER(ctypes.c_long)]
        self.lib.EN_runQ.restype = ctypes.c_int
        self.lib.EN_runQ.errcheck = _errcheck

        self.lib.EN_nextQ.argtypes = [self.EN_Project, ctypes.POINTER(ctypes.c_long)]
        self.lib.EN_nextQ.restype = ctypes.c_int
        self.lib.EN_nextQ.errcheck = _errcheck

        self.lib.EN_closeQ.argtypes = [self.EN_Project]
        self.lib.EN_closeQ.restype = ctypes.c_int
        self.lib.EN_closeQ.errcheck = _errcheck

        self.lib.EN_getqualinfo.argtypes = [
            self.EN_Project,
            ctypes.POINTER(ctypes.c_int),
            ctypes.c_char_p,
            ctypes.c_char_p,
            ctypes.POINTER(ctypes.c_int),
        ]
        self.lib.EN_getqualinfo.restype = ctypes.c_int
        self.lib.EN_getqualinfo.errcheck = _errcheck
        # ---- Bindings used by batch analyses (backend/batch.py) ----
        P, I, D, L, S = self.EN_Project, ctypes.c_int, ctypes.c_double, ctypes.c_long, ctypes.c_char_p

        def bind(name, argtypes):
            fn = getattr(self.lib, name)
            fn.argtypes = argtypes
            fn.restype = ctypes.c_int
            fn.errcheck = _errcheck

        bind("EN_getnodeindex", [P, S, ctypes.POINTER(I)])
        bind("EN_getlinkindex", [P, S, ctypes.POINTER(I)])
        bind("EN_setnodevalue", [P, I, I, D])
        bind("EN_setlinkvalue", [P, I, I, D])
        bind("EN_settimeparam", [P, I, L])
        bind("EN_addpattern", [P, S])
        bind("EN_getpatternindex", [P, S, ctypes.POINTER(I)])
        bind("EN_setpattern", [P, I, ctypes.POINTER(D), I])
        bind("EN_adddemand", [P, I, D, S, S])
        bind("EN_getnumdemands", [P, I, ctypes.POINTER(I)])
        bind("EN_deletedemand", [P, I, I])
        bind("EN_addcontrol", [P, I, I, D, I, D, ctypes.POINTER(I)])
        bind("EN_deletecontrol", [P, I])
        bind("EN_solveH", [P])
        bind("EN_setqualtype", [P, I, S, S, S])
        bind("EN_getpatternlen", [P, I, ctypes.POINTER(I)])
        bind("EN_getpatternvalue", [P, I, I, ctypes.POINTER(D)])
        bind("EN_getcontrol", [P, I, ctypes.POINTER(I), ctypes.POINTER(I), ctypes.POINTER(D), ctypes.POINTER(I), ctypes.POINTER(D)])