from OMSimulator.capi import Capi, Status
from OMSimulator.component import Component
from OMSimulator.componenttable import ComponentTable
from OMSimulator.cref import CRef
from OMSimulator.system import System
from OMSimulator.values import Values
from OMSimulator.ssm import SSM
from OMSimulator.variable import Causality, SignalType
import json
import tempfile
from enum import Enum

class SolverType(Enum):
  '''Enumeration for solver method to map with c api.'''
  euler = 2
  cvode = 3

class SystemType(Enum):
  '''Enumeration for system type to map with c api.'''
  wc = 1
  sc = 2
  sc3 = 3

class InstantiatedModel:
  _suppress_path_set = False # Class variable to track if suppressPath has been set
  def __init__(self, json_description, system: System, resources: dict):
    config = json.loads(json_description)
    self.modelName = "model" ## create random name, but we cannot commits test as jenkins will gerate new model name
    self.apiCall = []
    self.mappedCrefs = {}  # Store mapped CRefs associated with their export names
    self.system = system
    self.resources = resources
    self.fmuInstantitated = False

    # Only set once
    if not InstantiatedModel._suppress_path_set:
      status = Capi.setCommandLineOption("--suppressPath=true")
      if status != Status.ok:
        raise RuntimeError(f"Failed to set command line option: {status}")
      InstantiatedModel._suppress_path_set = True
    # Set the temporary directory
    status = Capi.setTempDirectory(tempfile.mkdtemp())
    if status != Status.ok:
      raise RuntimeError(f"Failed to set temp directory: {status}")
    # Create a new model
    status = Capi.newModel("model")
    if status != Status.ok:
      raise RuntimeError(f"Failed to create new model: {status}")
    self.apiCall.append(f'oms_newModel("{self.modelName}")')

    # Extract all simulation units
    sim_units = config.get("simulation units", [])
    # Count the number of unique solvers for determining system type (WC/SC)
    solver_units = [unit for unit in sim_units if "solver" in unit]
    num_solvers = len(solver_units)

    # --- Add systems depending on number of solvers ---
    if num_solvers == 1:
      # Only one solver unit → SC system directly under root
      status = Capi.addSystem(f"{self.modelName}.root", SystemType.sc.value)  # 2 = oms_system_sc
      if status != Status.ok:
        raise RuntimeError(f"Failed to create root SC system: {status}")
      self.apiCall.append(f'oms_addSystem("{self.modelName}.root", "oms_system_sc")')
    else:
      # Multiple solver units or zero solvers top-level WC system
      status = Capi.addSystem(f"{self.modelName}.root", SystemType.wc.value)  # 1 = oms_system_wc
      if status != Status.ok:
        raise RuntimeError(f"Failed to create root WC system: {status}")
      self.apiCall.append(f'oms_addSystem("{self.modelName}.root", "oms_system_wc")')

    # Iterate over simulation units
    for unit in sim_units:
      ## check if unit has solver as key
      solvername = "None"
      solver = unit.get("solver")
      if num_solvers > 1  and solver:
        solvername = unit["solver"]['name']
        solver_path = f"{self.modelName}.root.{solvername}"
        self.apiCall.append(f'oms_addSystem("{solver_path}", "oms_system_sc")')
        status = Capi.addSystem(f"model.root.{solvername}", SystemType.sc.value)  # 2 = oms_system_sc
        if status != Status.ok:
          raise RuntimeError(f"Failed to add oms_addSystem: {status}")
      else:
        solver_path = f"{self.modelName}.root"

      ## set solver method and tolerance if provided
      if solver:
        method = solver.get("method")
        tolerance = solver.get("tolerance")
        stepSize = solver.get("stepSize")

        if method in ("cvode", "euler"):
          status = Capi.setSolver(solver_path, SolverType[method].value)  # 3=cvode, 2=euler
          if status != Status.ok:
            raise RuntimeError(f"Failed to set solver: {status}")
          self.apiCall.append(f'oms_setSolver("{solver_path}", "{method}")')

        if tolerance is not None:
          status = Capi.setTolerance(solver_path, float(tolerance))
          if status != Status.ok:
            raise RuntimeError(f"Failed to set tolerance: {status}")
          self.apiCall.append(f'oms_setTolerance("{solver_path}", {float(tolerance)})')

        if stepSize is not None:
          status = Capi.setFixedStepSize(solver_path, float(stepSize))
          if status != Status.ok:
            raise RuntimeError(f"Failed to set step size: {status}")
          self.apiCall.append(f'oms_setFixedStepSize("{solver_path}", {float(stepSize)})')

      listofsystems = []
      ## add components
      for comp in unit["components"]:
        comp_path = ".".join([solver_path] + [comp["name"][-1]])
        currentSystem = ".".join(comp["name"][:-1])
        if currentSystem not in listofsystems:
          listofsystems.append(currentSystem)
        self.apiCall.append(f'oms_addSubModel("{comp_path}", "{comp["path"]}")')
        status = Capi.addSubModel(comp_path, comp["path"])
        if status != Status.ok:
          raise RuntimeError(f"Failed to add oms_addSubModel: {status}")
        export_name = ".".join(comp["name"])
        #print(f"Setting export name for {comp_path} to {export_name}")

        ## parse connector geometry for the component if exist and set it to capi after adding the component, this is needed for proper mapping of connector geometry
        if "connectors" in comp:
          for connector in comp["connectors"]:
            if "geometry" in connector:
              connector_path = ".".join([comp_path, connector["name"]])
              geometry = connector["geometry"]
              x = geometry.get("x", 0.0)
              y = geometry.get("y", 0.0)
              status = Capi.setConnectorGeometry(connector_path, x, y)
              if status != Status.ok:
                raise RuntimeError(f"Failed to set connector geometry for {connector_path}: {status}")
              self.apiCall.append(f'oms_setConnectorGeometry("{connector_path}", {x}, {y})')

        if not export_name in self.mappedCrefs:
          self.mappedCrefs[export_name] = comp_path
          self.mappedCrefs[".".join(comp["name"][:-1])] = solver_path # map parent system too for connector lookup
        status = Capi.setExportName(comp_path, export_name)  # Set export name if provided
        if status != Status.ok:
          raise RuntimeError(f"Failed to set export name: {status}")

      ## add top level system connectors before adding connections
      for current_system in listofsystems:
        ## top level connectors:
        if current_system == self.system.name:
          self._addConnector(self.system.connectors, self.system.name)
        ## add top sub-system level connectors mapped with currentSystem
        self.addConnectorFromElements(self.system.elements, current_system)

      ## add connections
      for connection in unit["connections"]:
        start_element = ".".join(connection['start element'])
        end_element = ".".join(connection['end element'])
        start = self.mappedCrefs[start_element] + f".{connection['start connector']}"
        end = self.mappedCrefs[end_element] + f".{connection['end connector']}"
        self.apiCall.append(f'oms_addConnection("{start}", "{end}")')
        status = Capi.addConnection(start, end)
        if status != Status.ok:
          raise RuntimeError(f"Failed to add oms_addConnection: {status}")
        ## add linear transformation if exist
        if "linear transformation" in connection:
          factor = connection["linear transformation"]["factor"]
          offset = connection["linear transformation"]["offset"]
          self.apiCall.append(f'oms_setConnectionLinearTransformation("{start}", "{end}", {factor}, {offset})')
          status = Capi.setConnectionLinearTransformation(start, end, float(factor), float(offset))
          if status != Status.ok:
            raise RuntimeError(f"Failed to set connection linear transformation: {status}")
        ## add connection geometry if exist
        if "connection geometry" in connection:
          pointsX = connection["connection geometry"]["pointsX"]
          pointsY = connection["connection geometry"]["pointsY"]
          self.apiCall.append(f'oms_setConnectionGeometry("{start}", "{end}", {pointsX}, {pointsY})')
          status = Capi.setConnectionGeometry(start, end, pointsX, pointsY)
          if status != Status.ok:
            raise RuntimeError(f"Failed to set connection geometry: {status}")

    ## set start values
    self.setStartValues(self.system.value, self.system.name, self.system.parameterMapping)
    ## set start values from ssv files
    self.setStartValuesFromSSV(self.system.parameterResources, self.system.name)
    ## iterate start values from sub-system both inline and ssv files if exist
    self.setStartValuesFromElements(self.system.elements, self.system.name)

    self.apiCall.append(f'oms_instantiate("{self.modelName}")')
    status = Capi.instantiate(self.modelName)
    if status != Status.ok:
      raise RuntimeError(f"Failed to instantiate model: {status}")
    self.fmuInstantitated = True
    ## set start and stop time from ssp
    self.setStartTime(float(config["simulation settings"]['start time']))
    self.setStopTime(float(config["simulation settings"]['stop time']))

  def setStartValuesFromElements(self, elements, systemName):
    for key, element in elements.items():
      ## do not set start values for tables
      if isinstance(element, ComponentTable):
        continue
      new_system_name = ".".join([systemName, str(element.name)])
      self.setStartValues(element.value, new_system_name, element.parameterMapping)
      self.setStartValuesFromSSV(element.parameterResources, new_system_name)

      ## recursive for sub-system
      if isinstance(element, System):
        new_system_name = ".".join([systemName, str(element.name)])
        self.setStartValuesFromElements(element.elements, new_system_name)

  def setStartValuesFromSSV(self, parameterResources, systemName):
    if not parameterResources:
      return
    for resource in parameterResources:
      for (key, mapping) in resource.items():
        if key in self.resources:
          ssv = self.resources.get(key)
          ssm = self.resources.get(mapping, None)  # get the ssm mapping if exist
          self.setStartValues(ssv.value, systemName, ssm)
        else:
          raise KeyError(f"Missing required resource: '{key}'")

  def map_cref(self, system_name, var_str):
    ## map cref with the longest path
    ## e.g. default.sub-system.gain1.k -> model.root.gain1.k
    ## e.g. default.sub-system.gain1.R1.T. -> model.root.solver2.gain1.R1.T
    # Combine system name and variable
    cref = f"{system_name}.{var_str}"
    # Split into parts
    parts = cref.split(".")
    ## map with the longest path
    for i in range(len(parts), 0, -1):
      prefix = ".".join(parts[:i])
      if prefix in self.mappedCrefs:
        mapped_prefix = self.mappedCrefs[prefix]
        validate_path = self.validatePath(mapped_prefix.split("."), parts[i:])
        return validate_path

    raise KeyError(f"No mapping found for {cref}")

  def validatePath(self, mapped_prefix : list, suffix : list):
    ## remove common path from suffix
    ## e.g. mapped_prefix = model.root.solver2.gain1
    ##      suffix = gain1.R1.T
    ##      result = model.root.solver2.gain1.R1.T
    result = []
    for item in suffix:
      if item not in mapped_prefix:
        result.append(item)

    return ".".join(mapped_prefix + result)

  def setStartValues(self, value: Values, systemName: str, ssm: SSM | None):
    if value.empty():
      return

    if systemName not in self.mappedCrefs:
      raise KeyError(f"Missing required key: '{systemName}'")

    ## apply ssm mapping if exist and apply start values from mapped parameters
    if ssm and ssm.mappingEntry:
      for source, targets in ssm.mappingEntry.items():
        if CRef(source) in value.start_values:
          source_value, _, _ = value.start_values[CRef(source)]
          for entry in targets:
            target = entry["target"]
            linearTransformation = entry["linearTransformation"]
            value_path = self.map_cref(systemName, str(target))
            if linearTransformation:
              source_value = source_value * float(linearTransformation.factor) + float(linearTransformation.offset)
            self.apply_start_value(value_path, source_value)
    else:
      for key, (source_value, _, _) in value.start_values.items():
        value_path = self.map_cref(systemName, str(key))
        self.apply_start_value(value_path, source_value)

  def apply_start_value(self, value_path:str, value):
    # Determine the variable type
    type, status = Capi.getVariableType(value_path)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get variable type for {value_path}: {status}")
    ## TODO: handle FMi3 data types directly, like Float64, Int32,
    value_ = value.value if hasattr(value, 'value') else value

    match SignalType(type):
      case SignalType.Real:  # oms_signal_type_real
        self._setReal(value_path, float(value_))
      case SignalType.Integer:  # oms_signal_type_integer
        self._setInteger(value_path, int(value_))
      case SignalType.Boolean:  # oms_signal_type_boolean
        self._setBoolean(value_path, bool(value_))
      case SignalType.String:  # oms_signal_type_string
        self._setString(value_path, str(value_))
      case SignalType.Enumeration:  # oms_signal_type_enumeration
        self._setInteger(value_path, int(value_))  # Treat enumeration as integer
      case _:
        raise TypeError(f"Unsupported type: {type}")

  def _addConnector(self, connectors, systemName):
    for connector in connectors:
      name = [systemName] + [str(connector.name)]
      connector_path = ".".join([self.mappedCrefs[systemName], str(connector.name)])
      self.apiCall.append(f'oms_addConnector("{connector_path}", "{connector.causality}", {connector.signal_type})')
      status = Capi.addConnector(connector_path, connector.causality.value, connector.signal_type.value)
      if status != Status.ok:
        raise RuntimeError(f"Failed to add oms_addConnector:{status}")
      ## set connector geometry if exist
      if connector.connectorGeometry:
        x = connector.connectorGeometry.x
        y = connector.connectorGeometry.y
        self.apiCall.append(f'oms_setConnectorGeometry("{connector_path}", {x}, {y})')
        status = Capi.setConnectorGeometry(connector_path, x, y)
        if status != Status.ok:
          raise RuntimeError(f"Failed to set connector geometry for {connector_path}: {status}")

      export_name = systemName
      if not export_name in self.mappedCrefs:
        self.mappedCrefs[export_name] = connector_path
      status = Capi.setExportName(connector_path, export_name)  # Set export name if provided
      if status != Status.ok:
        raise RuntimeError(f"Failed to set export name: {status}")

  def addConnectorFromElements(self, elements, currentSystem):
    ## add connectors mapped with currentSystem
    for key, element in elements.items():
      connector_path = ".".join([self.system.name, str(element.name)])
      if currentSystem == connector_path:
        self._addConnector(element.connectors, connector_path)

      ## recurse into subsystem
      if isinstance(element, System):
        self.addConnectorFromElements(element.elements, currentSystem)

  def dumpApiCalls(self):
    """Returns the generated API calls as a string."""
    return "\n".join(self.apiCall)

  def setValue(self, cref: CRef, value):
    """Sets a value for a specific CRef in the model."""
    name = ".".join(cref.names[:-1])
    if name not in self.mappedCrefs:
      raise KeyError(f"Missing required key: '{name}'")

    value_path = ".".join([self.mappedCrefs[name], cref.names[-1]])

    # Determine the variable type
    type, status = Capi.getVariableType(value_path)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get variable type for {cref}: {status}")

    value_ = value.value if hasattr(value, 'value') else value
    ## TODO handle FMi3 data types directly, like Float64, Int32,etc..
    match SignalType(type):
      case SignalType.Real:  # oms_signal_type_real
        return self._setReal(value_path, float(value_))
      case SignalType.Integer:  # oms_signal_type_integer
        return self._setInteger(value_path, int(value_))
      case SignalType.Boolean:  # oms_signal_type_boolean
        return self._setBoolean(value_path, bool(value_))
      case SignalType.String:  # oms_signal_type_string
        return self._setString(value_path, str(value_))
      case SignalType.Enumeration:  # oms_signal_type_enumeration
        return self._setInteger(value_path, int(value_))  # Treat enumeration as integer
      case _:
        raise TypeError(f"Unsupported type: {type}")

  def _setReal(self, mapped_cref: str, value: float):
    self.apiCall.append(f'oms_setReal("{mapped_cref}, {value})')
    status = Capi.setReal(mapped_cref, value) # Get the value from the CAPI
    if status != Status.ok:
      raise RuntimeError(f"Failed to set value for {mapped_cref}: {status}")

  def _setInteger(self, mapped_cref: str, value: int):
    self.apiCall.append(f'oms_setInteger("{mapped_cref}, {value})')
    status = Capi.setInteger(mapped_cref, value) # Get the value from the CAPI
    if status != Status.ok:
      raise RuntimeError(f"Failed to set value for {mapped_cref}: {status}")

  def _setBoolean(self, mapped_cref: str, value: bool):
    self.apiCall.append(f'oms_setBoolean("{mapped_cref}, {value})')
    status = Capi.setBoolean(mapped_cref, value) # Get the value from the CAPI
    if status != Status.ok:
      raise RuntimeError(f"Failed to set value for {mapped_cref}: {status}")

  def _setString(self, mapped_cref: str, value: str):
    self.apiCall.append(f'oms_setString("{mapped_cref}, {value})')
    status = Capi.setString(mapped_cref, value) # Get the value from the CAPI
    if status != Status.ok:
      raise RuntimeError(f"Failed to set value for {mapped_cref}: {status}")

  def getValue(self, cref: CRef):
    name = ".".join(cref.names[:-1])
    if name not in self.mappedCrefs:
      raise KeyError(f"Missing required key: '{name}'")

    value_path = ".".join([self.mappedCrefs[name], cref.names[-1]])

    # Determine the variable type
    type, status = Capi.getVariableType(value_path)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get variable type for {cref}: {status}")

    match SignalType(type):
      case SignalType.Real:  # oms_signal_type_real
        return self._getReal(value_path)
      case SignalType.Integer:  # oms_signal_type_integer
        return self._getInteger(value_path)
      case SignalType.Boolean:  # oms_signal_type_boolean
        return self._getBoolean(value_path)
      case SignalType.String:  # oms_signal_type_string
        return self._getString(value_path)
      case SignalType.Enumeration:  # oms_signal_type_enumeration
        return self._getInteger(value_path)  # Treat enumeration as integer
      case _:
        raise TypeError(f"Unsupported type: {type}")

  def _getReal(self, cref: CRef):
    value, status = Capi.getReal(cref) # Get the value from the CAPI
    if status != Status.ok:
      raise RuntimeError(f"Failed to get value for {cref}: {status}")
    return value

  def _getInteger(self, cref: CRef):
    value, status = Capi.getInteger(cref)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get value for {cref}: {status}")
    return value

  def _getBoolean(self, cref: CRef):
    value, status = Capi.getBoolean(cref)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get value for {cref}: {status}")
    return value

  def _getString(self, cref: CRef):
    value, status = Capi.getString(cref)
    if status != Status.ok:
      raise RuntimeError(f"Failed to get value for {cref}: {status}")
    return value

  def initialize(self):
    status = Capi.initialize(self.modelName)
    if status != Status.ok:
      raise RuntimeError(f"Failed to initialize model: {status}")

  def simulate(self):
    status = Capi.simulate(self.modelName)
    if status != Status.ok:
      raise RuntimeError(f"Failed to simulate model: {status}")

  def stepUntil(self, stopTime):
    status = Capi.stepUntil(self.modelName, stopTime)
    if status != Status.ok:
      raise RuntimeError(f"Failed to step until {stopTime}: {status}")

  def setResultFile(self, filename: str):
    status = Capi.setResultFile("model", filename)
    if status !=Status.ok:
      raise RuntimeError(f"Failed to setResultFile {filename}: {status}")

  def setStartTime(self, startTime: float):
    status = Capi.setStartTime(self.modelName, startTime)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set start time: {status}")

  def setStopTime(self, stopTime: float):
    if self.fmuInstantitated is False:
      raise RuntimeError("FMU must be instantiated before setting stop time")

    status = Capi.setStopTime(self.modelName, stopTime)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set stop time: {status}")

  def setTolerance(self, tolerance: float):
    if self.fmuInstantitated is False:
      raise RuntimeError("FMU must be instantiated before setting tolerance")

    status = Capi.setTolerance(f"{self.modelName}.root", tolerance)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set tolerance: {status}")

  def setFixedStepSize(self, stepSize: float):
    if self.fmuInstantitated is False:
      raise RuntimeError("FMU must be instantiated before setting variable step size")

    status = Capi.setFixedStepSize(f"{self.modelName}.root", stepSize)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set fixed step size: {status}")

  def setVariableStepSize(self, initialStepSize: float, minimumStepSize: float, maximumStepSize: float):
    if self.fmuInstantitated is False:
      raise RuntimeError("FMU must be instantiated before setting variable step size")

    status = Capi.setVariableStepSize(f"{self.modelName}.root", initialStepSize, minimumStepSize, maximumStepSize)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set variable step size: {status}")

  def setLoggingLevel(self, level: int):
    status = Capi.setLoggingLevel(level)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set logging level: {status}")

  def setLogFile(self, filename: str):
    status = Capi.setLogFile(filename)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set log file: {status}")

  def setLoggingInterval(self, interval: float):
    status = Capi.setLoggingInterval(self.modelName, interval)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set logging interval: {status}")

  def addSignalsToResults(self, regex: str):
    status = Capi.addSignalsToResults(self.modelName, regex)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set regular expression : {status}")

  def removeSignalsFromResults(self, regex: str):
    status = Capi.removeSignalsFromResults(self.modelName, regex)
    if status != Status.ok:
      raise RuntimeError(f"Failed to set regular expression : {status}")

  def terminate(self):
    status = Capi.terminate(self.modelName)
    if status != Status.ok:
      raise RuntimeError(f"Failed to terminate model: {status}")

  def delete(self):
    status = Capi.delete(self.modelName)
    if status != Status.ok:
      raise RuntimeError(f"Failed to delete model: {status}")
    self.modelName = None  # Clear the model name after deletion
    self.apiCall.clear()  # Clear the API call history
