# This file is part of OpenModelica.
#
# Copyright (c) 1998-2026, Open Source Modelica Consortium (OSMC),
# c/o Linköpings universitet, Department of Computer and Information Science,
# SE-58183 Linköping, Sweden.
#
# All rights reserved.
#
# THIS PROGRAM IS PROVIDED UNDER THE TERMS OF AGPL VERSION 3 LICENSE OR
# THIS OSMC PUBLIC LICENSE (OSMC-PL) VERSION 1.8.
# ANY USE, REPRODUCTION OR DISTRIBUTION OF THIS PROGRAM CONSTITUTES
# RECIPIENT'S ACCEPTANCE OF THE OSMC PUBLIC LICENSE OR THE GNU AGPL
# VERSION 3, ACCORDING TO RECIPIENTS CHOICE.
#
# The OpenModelica software and the OSMC (Open Source Modelica Consortium)
# Public License (OSMC-PL) are obtained from OSMC, either from the above
# address, from the URLs:
# http://www.openmodelica.org or
# https://github.com/OpenModelica/ or
# http://www.ida.liu.se/projects/OpenModelica,
# and in the OpenModelica distribution.
#
# GNU AGPL version 3 is obtained from:
# https://www.gnu.org/licenses/licenses.html#GPL
#
# This program is distributed WITHOUT ANY WARRANTY; without
# even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE, EXCEPT AS EXPRESSLY SET FORTH
# IN THE BY RECIPIENT SELECTED SUBSIDIARY LICENSE CONDITIONS OF OSMC-PL.
#
# See the full OSMC Public License conditions for more details.

import logging
import warnings
from lxml import etree as ET
from pathlib import Path

from OMSimulator.connection import Connection, ConnectionGeometry, LinearTransformation
from OMSimulator.connector import Connector, ConnectorGeometry
from OMSimulator.unit import Unit
from OMSimulator.variable import Causality, SignalType, Float64, Float32, Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64
from OMSimulator.elementgeometry import ElementGeometry
from OMSimulator import namespace, CRef
from collections import defaultdict

logger = logging.getLogger(__name__)

# Custom warning formatter: only show category + message
def _simple_warning(message, category, filename, lineno, line=None):
  return f"{category.__name__}: {message}\n"
warnings.formatwarning = _simple_warning

def _setParameters(parameterValues: dict, obj):
  if len(parameterValues) > 0:
    for key, (value, type, unit, description) in parameterValues.items():
      obj.value.setValue(key, value, type, unit, description)

def parseDefaultExperiment(node, root):
  default_experiment = node.find("ssd:DefaultExperiment", namespaces=namespace.ns)
  if default_experiment is None:
    return
  root.startTime = default_experiment.get("startTime")
  root.stopTime = default_experiment.get("stopTime")
  ##TODO parse ssd:annotation


def parseElements(node, resources = None):
  """Extract components from <ssd:Elements> section"""
  from OMSimulator.system import System, SystemGeometry
  from OMSimulator.component import Component
  from OMSimulator.componenttable import ComponentTable

  elements = {}
  elements_node = node.find("ssd:Elements", namespaces=namespace.ns)
  if elements_node is None:
    return elements

  # Parse the sub system <ssd:system> in the <ssd:Elements> section
  for system in elements_node.findall("ssd:System", namespaces=namespace.ns):
    name = CRef(system.get("name"))
    description = system.get("description")
    elements[name] = System(name)
    elements[name].description = description
    elements[name].connectors = Connector.importFromNode(system)
    elements[name].elementgeometry = ElementGeometry.importFromNode(system)
    elements[name].systemgeometry = SystemGeometry.importFromNode(system)
    parseParameterBindings(system, elements[name], resources)
    ## check for meta tag
    parseMetaData(system, elements[name], resources)
    solvers = parseAnnotations(system)
    if solvers:
      for solver in solvers:
        elements[name].solver = solver.get("name")
    elements[name].elements = parseElements(system, resources)  # recursively parse nested elements in the sub-system
    Connection.importFromNode(system, elements[name]) # parse connections for the sub-system

  for component in elements_node.findall("ssd:Component", namespaces=namespace.ns):
    comp_type = component.get("type")
    name = CRef(component.get("name"))
    source = component.get("source")
    description = component.get("description")
    if (comp_type == "application/x-fmu-sharedlibrary"):
      ## check implementation
      match component.get("implementation", "any"):
        case "any":
          implementation = "me_cs"
        case "ModelExchange":
          implementation = "me"
        case "CoSimulation":
          implementation = "cs"
        case _:
          raise ValueError(f"Unknown FMU implementation type: {implementation}")

      elements[name] = Component(name, source)
      elements[name].implementation = implementation
      elements[name].description = description
      elements[name].connectors = Connector.importFromNode(component)
      elements[name].elementgeometry = ElementGeometry.importFromNode(component)
      parseParameterBindings(component, elements[name], resources)
      ## check for meta tag
      parseMetaData(component, elements[name], resources)
      solvers = parseAnnotations(component)
      if solvers:
        for solver in solvers:
          elements[name].solver = solver.get("name")
    elif (comp_type == "application/table" or comp_type == "text/csv"):
      elements[name] = ComponentTable(name, source)
      elements[name].connectors = Connector.importFromNode(component)

  return elements

def parseMetaData(node, obj, resources):
  ssc_meta_node = node.find("ssc:MetaData", namespaces=namespace.ns)
  if ssc_meta_node is not None:
    source = ssc_meta_node.get("source", "")
    kind = ssc_meta_node.get("kind", "general")
    type = ssc_meta_node.get("type", "application/octet-stream")
    ## TODO check for other attributes
    obj.metaDataResources.append({"source": source, "kind":kind, "type":type})

def parseParameterBindings(node, obj, resources):
  """Extract and print system parameters"""
  parameter_bindings = node.find("ssd:ParameterBindings", namespaces=namespace.ns)
  if parameter_bindings is not None:
    for binding in parameter_bindings.findall("ssd:ParameterBinding", namespaces=namespace.ns):
      source = binding.get("source")
      if source is not None:
        ## use the instantiated ssv class to set the parameter Resources
        if source not in resources:
          logger.warning(f"SSV file not found: {source}")
        parameter_mapping = binding.find("ssd:ParameterMapping", namespaces=namespace.ns)
        mapped_source = parameter_mapping.get("source") if parameter_mapping is not None else None
        if mapped_source is not None and  mapped_source not in resources:
          logger.warning(f"SSM file not found: {mapped_source}")
        obj.parameterResources.append({source: mapped_source})
      else:
        values = binding.find("ssd:ParameterValues", namespaces=namespace.ns)
        if values is not None:
          param_set = values.find("ssv:ParameterSet", namespaces=namespace.ns)
          if param_set is not None:
            parameters = param_set.find("ssv:Parameters", namespaces=namespace.ns)
            parameterValues = parseParameterBindingHelper(parameters)
            _setParameters(parameterValues, obj)
          ## parse parameter mapping
          parameter_mapping_ssd = binding.find("ssd:ParameterMapping", namespaces=namespace.ns)
          if parameter_mapping_ssd is not None:
            parameter_mapping_ssm = parameter_mapping_ssd.find("ssm:ParameterMapping", namespaces=namespace.ns)
            for mapping in parameter_mapping_ssm.findall("ssm:MappingEntry", namespaces=namespace.ns):
              obj.parameterMapping.mapParameter(mapping.get("source"), mapping.get("target"))

          Unit.importFromNode(param_set, obj, tagname="ssv:Units")

def parseSSV(filename):
  tree = ET.parse(filename)
  root = tree.getroot()
  validateSSP(root, filename, "SystemStructureParameterValues.xsd")
  parameters = root.find("ssv:Parameters", namespaces=namespace.ns)
  return parseParameterBindingHelper(parameters)

def parseSSM(filename):
  tree = ET.parse(filename)
  root = tree.getroot()
  validateSSP(root, filename, "SystemStructureParameterMapping.xsd")
  mappingEntry = defaultdict(list)
  for mapping in root.findall("ssm:MappingEntry", namespaces=namespace.ns):
    source = mapping.get("source")
    target = mapping.get("target")
    linearTransformation = None
    ## find ssc:LinearTransformation in the same mapping entry
    linearTransformation_node = mapping.find("ssc:LinearTransformation", namespaces=namespace.ns)
    if linearTransformation_node is not None:
      factor = linearTransformation_node.get("factor")
      offset = linearTransformation_node.get("offset")
      linearTransformation = LinearTransformation(factor, offset)
    ##TODO add support for other GTTransformation such as BooleanMappingTransformation, IntegerMappingTransformation and EnumerationMappingTransformation
    mappingEntry[source].append({"target": target, "linearTransformation": linearTransformation})
  return mappingEntry

def parseParameterBindingHelper(parameters):
  if parameters is not None:
    parameterValues={}
    for param in parameters.findall("ssv:Parameter", namespaces=namespace.ns):
      name = CRef(param.get("name"))
      description = param.get("description")
      value_types = {
                      "ssv:Real": float,
                      "ssv:Integer": int,
                      "ssv:Boolean": lambda v: v.lower() == "true",  # Convert "true"/"false" to bool
                      "ssv:String": str,
                      "ssv:Enumeration": int,
                       ## FMI-3.0 types
                      "ssv:Float64": Float64,
                      "ssv:Float32": Float32,
                      "ssv:Int8": Int8,
                      "ssv:Int16": Int16,
                      "ssv:Int32": Int32,
                      "ssv:Int64": Int64,
                      "ssv:UInt8": UInt8,
                      "ssv:UInt16": UInt16,
                      "ssv:UInt32": UInt32,
                      "ssv:UInt64": UInt64
                    }
      for value_type, cast_func in value_types.items():
        value_element = param.find(value_type, namespaces=namespace.ns)
        if value_element is not None:
          value = value_element.get("value")
          unit = value_element.get("unit")
           # Extract type name (remove namespace prefix)
          type_name = value_type.split(":")[1]
          if type_name == "Enumeration":
            type_name = "Integer"  # Treat Enumeration as Integer for type purposes
          parameterValues[name] = (cast_func(value), SignalType[type_name], unit, description)  # Convert to correct type
          break  # Stop after first found type
    return parameterValues

def validateSSP(root, filename : str, schema_file : str):
  module_dir = Path(__file__).parent
  schema = ET.XMLSchema(file=Path(module_dir, 'schema/ssp', schema_file))
  if not schema.validate(root):
    message = f"Failed to validate {Path(filename).name} against schemafile {schema_file}"
    for entry in schema.error_log:
      message += "\n%s (line %d, column %d): %s" % (entry.level_name, entry.line, entry.column, entry.message)
    # warn the user instead of raising an exception
    warnings.warn(message, UserWarning)

def exportAnnotations(node, solvers):
  """Export annotations to the XML node"""
  ssd_annotation_node = ET.SubElement(node, namespace.tag("ssd", "Annotations"))
  annotation_node = ET.SubElement(ssd_annotation_node, namespace.tag("ssc", "Annotation"))
  annotation_node.set("type", "org.openmodelica")
  oms_annotation_node = ET.SubElement(annotation_node, namespace.tag("oms", "Annotations"))
  if (isinstance(solvers, list)):
    for solver in solvers:
      oms_simulationInformation_node = ET.SubElement(oms_annotation_node, namespace.tag("oms", "SimulationInformation"))
      for key, value in solver.items():
        oms_simulationInformation_node.set(key, str(value))
  else:
    oms_simulationInformation_node = ET.SubElement(oms_annotation_node, namespace.tag("oms", "SimulationInformation"))
    oms_simulationInformation_node.set("name", solvers)

def parseAnnotations(node):
  """Extract and print system annotations"""
  annotations_node = node.find("ssd:Annotations", namespaces=namespace.ns)
  if annotations_node is None:
    return []

  solvers = []
  for annotation in annotations_node.findall("ssc:Annotation", namespaces=namespace.ns):
    type = annotation.get("type")
    if type == "org.openmodelica":
      oms_annotation = annotation.find("oms:Annotations", namespaces=namespace.ns)
      if oms_annotation is not None:
        oms_simulationInformation = oms_annotation.findall("oms:SimulationInformation", namespaces=namespace.ns)
        for sim in oms_simulationInformation:
          solvers.append(sim.attrib)

  return solvers
