Container

class AutoWired:
    """Collects the rules automatically generated by implicit wiring.

    Used to report back to the
    [AssemblyNode][fmu_manipulation_toolbox.assembly.AssemblyNode] which
    inputs, outputs, and links were created by auto-wiring, so they
    can be recorded in the assembly topology.

    Attributes:
        rule_input (list[list[str]]): Auto-generated input rules
            `[exposed_name, fmu_name, port_name]`.
        rule_output (list[list[str]]): Auto-generated output rules
            `[fmu_name, port_name, exposed_name]`.
        rule_link (list[list[str]]): Auto-generated link rules
            `[from_fmu, from_port, to_fmu, to_port]`.
        nb_param (int): Number of auto-exposed parameters (subset of inputs).
    """

    def __init__(self):
        self.rule_input = []
        self.rule_output = []
        self.rule_link = []
        self.nb_param = 0

    def __repr__(self):
        return (f"{self.nb_param} parameters, {len(self.rule_input) - self.nb_param} inputs,"
                f" {len(self.rule_output)} outputs, {len(self.rule_link)} links.")

    def add_input(self, from_port, to_fmu, to_port):
        self.rule_input.append([from_port, to_fmu, to_port])

    def add_parameter(self, from_port, to_fmu, to_port):
        self.rule_input.append([from_port, to_fmu, to_port])
        self.nb_param += 1

    def add_output(self, from_fmu, from_port, to_port):
        self.rule_output.append([from_fmu, from_port, to_port])

    def add_link(self, from_fmu, from_port, to_fmu, to_port):
        self.rule_link.append([from_fmu, from_port, to_fmu, to_port])

class ClockList:
    """Tracks clocks that need to be scheduled by the FMI importer.

    Used for LS-BUS support where the container runtime needs to trigger
    countdown clocks on embedded FMUs.

    Attributes:
        clocks_per_fmu (dict[int, list[tuple[int, int]]]): Clock entries
            per FMU index: `(fmu_vr, local_vr)` pairs.
        fmu_index (dict[str, int]): Mapping from FMU name to its index
            in the container.
    """

    def __init__(self, involved_fmu: OrderedDict[str, EmbeddedFMU]):
        self.clocks_per_fmu: DefaultDict[int, List[Tuple[int, int]]] = defaultdict(list)
        self.fmu_index: Dict[str, int] = {}
        for i, fmu_name in enumerate(involved_fmu):
            self.fmu_index[fmu_name] = i

    def append(self, cport: ContainerPort, vr: int):
        """Register a clock for importer scheduling.

        Args:
            cport (ContainerPort): The clocked port on the embedded FMU.
            vr (int): The local value reference of the clock.
        """
        self.clocks_per_fmu[self.fmu_index[cport.fmu.name]].append((cport.port.vr, vr))

    def write_txt(self, txt_file):
        """Write the clock scheduling table to the `container.txt` file.

        Args:
            txt_file: Writable text file handle.
        """
        print(f"# importer CLOCKS: <FMU_INDEX> <NB> <FMU_VR> <VR> [<FMU_VR> <VR>]", file=txt_file)
        nb_total_clocks = 0
        for clocks in self.clocks_per_fmu.values():
            nb_total_clocks += len(clocks)

        print(f"{len(self.clocks_per_fmu)} {nb_total_clocks}", file=txt_file)
        for index, clocks in self.clocks_per_fmu.items():
            clocks_str = " ".join([f"{clock[0]} {clock[1]}" for clock in clocks])
            print(f"{index} {len(clocks)} {clocks_str}", file=txt_file)

class ContainerInput:
    """Represents an input port exposed by the container.

    A single container input can fan out to multiple embedded FMU input ports,
    provided they share the same type and causality.

    Attributes:
        name (str): Exposed name of the container input.
        type_name (str): Container-internal type name (e.g. `"real64"`).
        causality (str): Port causality (`"input"` or `"parameter"`).
        cport_list (list[ContainerPort]): List of embedded FMU ports connected
            to this input.
        vr (int | None): Value reference assigned by the container.
    """

    def __init__(self, name: str, cport_to: ContainerPort):
        self.name = name
        self.type_name = cport_to.port.type_name
        self.causality = cport_to.port.causality
        self.cport_list = [cport_to]
        self.vr = None

    def add_cport(self, cport_to: ContainerPort):
        """Connect an additional embedded FMU port to this container input.

        Args:
            cport_to (ContainerPort): The embedded FMU port to connect.

        Raises:
            FMUContainerError: If the port is already connected, or if types
                or causalities do not match.
        """
        if cport_to in self.cport_list: # Cannot be reached ! (Assembly prevent this to happen)
            raise FMUContainerError(f"Duplicate INPUT {cport_to} already connected to {self.name}")

        if cport_to.port.type_name != self.type_name:
            raise FMUContainerError(f"Cannot connect {self.name} of type {self.type_name} to "
                                    f"{cport_to} of type {cport_to.port.type_name}")

        if cport_to.port.causality != self.causality:
            raise FMUContainerError(f"Cannot connect {self.causality.upper()} {self.name} to "
                                    f"{cport_to.port.causality.upper()} {cport_to}")

        self.cport_list.append(cport_to)

class ContainerPort:
    """References a specific port of an embedded FMU within a container.

    Wraps an [EmbeddedFMUPort][fmu_manipulation_toolbox.container.EmbeddedFMUPort]
    together with its parent
    [EmbeddedFMU][fmu_manipulation_toolbox.container.EmbeddedFMU], and tracks
    the value reference assigned by the container.

    Attributes:
        fmu (EmbeddedFMU): The embedded FMU owning this port.
        port (EmbeddedFMUPort): The port descriptor.
        vr (int | None): Value reference assigned by the container, or `None`
            if not yet assigned.

    Raises:
        FMUContainerError: If the port name does not exist in the FMU.
    """

    def __init__(self, fmu: EmbeddedFMU, port_name: str):
        self.fmu = fmu
        try:
            self.port = fmu.ports[port_name]
        except KeyError:
            raise FMUContainerError(f"Port '{fmu.name}/{port_name}' does not exist")
        self.vr = None

    def __repr__(self):
        return f"Port {self.fmu.name}/{self.port.name}"

    def __hash__(self):
        return hash(str(self))

    def __eq__(self, other):
        return str(self) == str(other)

class EmbeddedFMU(OperationAbstract):
    """Represents an FMU loaded and analysed for embedding inside a container.

    Parses the `modelDescription.xml` of an FMU to extract its ports,
    capabilities, step size, platform support, and co-simulation metadata.
    Implements
    [OperationAbstract][fmu_manipulation_toolbox.operations.OperationAbstract]
    to process the FMU descriptor via the visitor pattern.

    Attributes:
        capability_list (tuple[str, ...]): FMI capability flags tracked by the container.
        fmu (FMU): The underlying
            [FMU][fmu_manipulation_toolbox.operations.FMU] object.
        name (str): Filename of the FMU (e.g. `"model.fmu"`).
        id (str): Lowercase stem of the filename, used as an identifier.
        terminals (Terminals): FMI Terminals defined by this FMU.
        ls (LayeredStandard): LS-BUS layered standard information.
        step_size (float | None): Preferred step size in seconds, or `None`.
        start_time (float | None): Default experiment start time.
        stop_time (float | None): Default experiment stop time.
        model_identifier (str | None): Co-simulation model identifier.
        guid (str | None): GUID (FMI 2.0) or instantiation token (FMI 3.0).
        fmi_version (int | None): FMI version (`2` or `3`).
        platforms (set[str]): Supported operating systems (e.g. `{"Windows", "Linux"}`).
        ports (dict[str, EmbeddedFMUPort]): Ports of the FMU, keyed by name.
        has_event_mode (bool): Whether the FMU supports event mode (FMI 3.0).
        capabilities (dict[str, str]): FMI capability flags and their values.

    Raises:
        FMUContainerError: If the FMU does not implement Co-Simulation mode.
    """

    capability_list = ("needsExecutionTool",
                       "canBeInstantiatedOnlyOncePerProcess",
                       "canHandleVariableCommunicationStepSize")

    def __init__(self, filename):
        self.fmu = FMU(filename)
        self.name = Path(filename).name
        self.id = Path(filename).stem.lower()

        logger.debug(f"Analysing {self.name}")
        self.terminals = Terminals(self.fmu.tmp_directory)
        self.ls = LayeredStandard(self.fmu.tmp_directory)

        self.step_size = None
        self.start_time = None
        self.stop_time = None
        self.model_identifier = None
        self.guid = None
        self.fmi_version = None
        self.platforms = set()
        self.ports: Dict[str, EmbeddedFMUPort] = {}

        self.has_event_mode = False
        self.capabilities: Dict[str, str] = {}
        self.current_port = None  # used during apply_operation()

        self.fmu.apply_operation(self)  # Should be the last command in constructor!
        if self.model_identifier is None:
            raise FMUContainerError(f"FMU '{self.name}' does not implement Co-Simulation mode.")

    def fmi_attrs(self, attrs):
        fmi_version = attrs['fmiVersion']
        if fmi_version == "2.0":
            self.guid = attrs['guid']
            self.fmi_version = 2
        if fmi_version.startswith("3."):
            self.guid = attrs['instantiationToken']
            self.fmi_version = 3

    def cosimulation_attrs(self, attrs: Dict[str, str]):
        self.model_identifier = attrs['modelIdentifier']
        if attrs.get("hasEventMode", "false") == "true":
            self.has_event_mode = True
        for capability in self.capability_list:
            self.capabilities[capability] = attrs.get(capability, "false")

    def experiment_attrs(self, attrs: Dict[str, str]):
        try:
            self.step_size = float(attrs['stepSize'])
        except KeyError:
            logger.warning(f"FMU '{self.name}' does not specify preferred step size")
        self.start_time = float(attrs.get("startTime", 0.0))
        self.stop_time = float(attrs.get("stopTime", self.start_time + 1.0))

    def port_attrs(self, fmu_port: FMUPort):
        # Container will manage Enumeration as Integer
        if fmu_port.fmi_type == "Enumeration":
            if self.fmi_version == 2:
                fmu_port.fmi_type = "Integer"
            else:
                fmu_port.fmi_type = "Int32"
        port = EmbeddedFMUPort(fmu_port.fmi_type, fmu_port, fmi_version=self.fmi_version)
        self.ports[port.name] = port

    def closure(self):
        osname = {
            "win64": "Windows",
            "linux64": "Linux",
            "darwin64": "Darwin",
            "x86_64-windows": "Windows",
            "x86_64-linux": "Linux",
            "aarch64-darwin": "Darwin"
        }
        try:
            for directory in (Path(self.fmu.tmp_directory) / "binaries").iterdir():
                if directory.is_dir() and str(directory.stem) in osname:
                    self.platforms.add(osname[str(directory.stem)])
        except FileNotFoundError:
            pass  # no binaries

    def __repr__(self):
        properties = f"{len(self.ports)} variables, ts={self.step_size}s"
        if len(self.terminals) > 0:
            properties += f", {len(self.terminals)} terminals"
        if len(self.ls) > 0:
            properties += f", {self.ls}"
        return f"'{self.name}' ({properties})"

class EmbeddedFMUPort:
    """Represents a port of an FMU embedded inside a container.

    Handles the mapping between FMI-standard type names (e.g. `Real`, `Float64`)
    and internal container type names (e.g. `real64`), and generates the
    corresponding XML fragments for `modelDescription.xml`.

    Attributes:
        FMI_TO_CONTAINER (dict[int, dict[str, str]]): Mapping from FMI type names
            to container type names, keyed by FMI version.
        CONTAINER_TO_FMI (dict[int, dict[str, str]]): Reverse mapping from container
            type names to FMI type names, keyed by FMI version.
        ALL_TYPES (tuple[str, ...]): All container type names.
        causality (str): Port causality (`"input"`, `"output"`, `"local"`,
            `"parameter"`).
        variability (str | None): Port variability (`"continuous"`, `"discrete"`, etc.).
        name (str): Port name.
        vr (int): Value reference in the original FMU.
        type_name (str): Container-internal type name (e.g. `"real64"`).
        start_value (str | None): Start value, if defined.
        initial (str | None): Initial value attribute.
        clock (str | None): Clock reference for clocked ports.
        description (str | None): Human-readable description of the port.
    """

    FMI_TO_CONTAINER = {
        2: {
            'Real': 'real64',
            'Integer': 'integer32',
            'String': 'string',
            'Boolean': 'boolean'
        },
        3: {
            'Float64': 'real64',
            'Float32': 'real32',
            'Int8': 'integer8',
            'UInt8': 'uinteger8',
            'Int16': 'integer16',
            'UInt16': 'uinteger16',
            'Int32': 'integer32',
            'UInt32': 'uinteger32',
            'Int64': 'integer64',
            'UInt64': 'uinteger64',
            'String': 'string',
            'Boolean': 'boolean1',
            'Binary': 'binary',
            'Clock': 'clock'
        }
    }

    CONTAINER_TO_FMI = {
        2: {
            'real64': 'Real',
            'integer32': 'Integer',
            'string': 'String',
            'boolean': 'Boolean'
        },
        3: {
            'real64': 'Float64' ,
            'real32': 'Float32' ,
            'integer8': 'Int8' ,
            'uinteger8': 'UInt8' ,
            'integer16': 'Int16' ,
            'uinteger16': 'UInt16' ,
            'integer32': 'Int32' ,
            'uinteger32': 'UInt32' ,
            'integer64': 'Int64' ,
            'uinteger64': 'UInt64' ,
            'string': 'String' ,
            'boolean1': 'Boolean',
            'binary': 'Binary',
            'clock': 'Clock'
        }
    }

    ALL_TYPES = (
        "real64", "real32",
        "integer8", "uinteger8", "integer16", "uinteger16", "integer32", "uinteger32", "integer64", "uinteger64",
        "boolean", "boolean1",
        "string",
        "binary", "clock"
    )

    def __init__(self, fmi_type, attrs: Union[FMUPort, Dict[str, str]], fmi_version=0):
        self.causality = attrs.get("causality", "local")
        self.variability = attrs.get("variability", None)
        self.interval_variability = attrs.get("intervalVariability", None)
        self.name = attrs["name"]
        self.vr = int(attrs["valueReference"])
        self.description = attrs.get("description", None)

        if fmi_version > 0:
            self.type_name = self.FMI_TO_CONTAINER[fmi_version][fmi_type]
        else:
            self.type_name = fmi_type

        self.start_value = attrs.get("start", None)
        self.initial = attrs.get("initial", None)
        self.clock = attrs.get("clocks", None)
        self.interval_variability = attrs.get("intervalVariability", None)


    def xml(self, vr: int, name=None, causality=None, start=None, fmi_version=2) -> str:
        """Generate the XML element for this port in `modelDescription.xml`.

        Produces a `<ScalarVariable>` element (FMI 2.0) or a typed element
        like `<Float64>` (FMI 3.0).

        Args:
            vr (int): Value reference to use in the generated XML.
            name (str | None): Override port name. Defaults to `self.name`.
            causality (str | None): Override causality. Defaults to `self.causality`.
            start (str | None): Override start value. Defaults to `self.start_value`.
            fmi_version (int): FMI version (`2` or `3`).

        Returns:
            str: XML fragment string, or an empty string if the type is not
                compatible with the requested FMI version.
        """
        if name is None:
            name = self.name
        if causality is None:
            causality = self.causality
        if start is None:
            start = self.start_value
            if start is None and self.type_name == "binary" and self.initial == "exact":
                start = ""
        if self.variability is None:
            self.variability = "continuous" if "real" in self.type_name else "discrete"

        try:
            fmi_type = self.CONTAINER_TO_FMI[fmi_version][self.type_name]
        except KeyError:
            logger.error(f"Cannot expose ({causality}) '{name}' because type '{self.type_name}' is not compatible "
                         f"with FMI-{fmi_version}.0")
            return ""

        if fmi_version == 2:
            child_attrs =  {
                "start": start,
            }

            filtered_child_attrs = {key: value for key, value in child_attrs.items() if value is not None}
            child_str = (f"<{fmi_type} " +
                         " ".join([f'{key}="{value}"' for (key, value) in filtered_child_attrs.items()]) +
                         "/>")

            scalar_attrs = {
                "name": name,
                "valueReference": vr,
                "causality": causality,
                "variability": self.variability,
                "initial": self.initial,
                "description": self.description,
            }
            filtered_attrs = {key: value for key, value in scalar_attrs.items() if value is not None}
            scalar_attrs_str = " ".join([f'{key}="{value}"' for (key, value) in filtered_attrs.items()])
            return f'<ScalarVariable {scalar_attrs_str}>{child_str}</ScalarVariable>'

        elif fmi_version == 3:
            if fmi_type in ('String', 'Binary'):
                if start is not None:
                    child_str = f'<Start value="{start}"/>'
                else:
                    child_str = ''
                scalar_attrs = {
                    "name": name,
                    "valueReference": vr,
                    "causality": causality,
                    "variability": self.variability,
                    "initial": self.initial,
                    "description": self.description,
                }
                filtered_attrs = {key: value for key, value in scalar_attrs.items() if value is not None}
                scalar_attrs_str = " ".join([f'{key}="{value}"' for (key, value) in filtered_attrs.items()])
                return f'<{fmi_type} {scalar_attrs_str}>{child_str}</{fmi_type}>'
            else:
                scalar_attrs = {
                    "name": name,
                    "valueReference": vr,
                    "causality": causality,
                    "variability": self.variability,
                    "initial": self.initial,
                    "description": self.description,
                    "start": start,
                    "intervalVariability": self.interval_variability
                }
                filtered_attrs = {key: value for key, value in scalar_attrs.items() if value is not None}
                scalar_attrs_str = " ".join([f'{key}="{value}"' for (key, value) in filtered_attrs.items()])

                return f'<{fmi_type} {scalar_attrs_str}/>'
        else:
            logger.critical(f"Unknown version {fmi_version}. BUG?")
            return ''