from pypdevs.DEVS import CoupledDEVS, Port


class FlowchartGenerator:
    def __init__(self, model: CoupledDEVS) -> None:
        self.model = model

        self._tab_level: int = 0
        self._result: str = ""

    def _add(self, text: str) -> None:
        """ Add text to the result, with the appropriate tab level"""
        self._result += ('\t' * self._tab_level) + text + '\n'

    def _tab(self):
        """ Increase tab level """
        self._tab_level += 1

    def _untab(self):
        """ Decrease tab level """
        self._tab_level = max(0, self._tab_level - 1)

    @staticmethod
    def _get_port_name(port: Port) -> (str, str):
        """ Output port names, first element is id, second is shorthand name"""
        names = [port.name]
        parent = port.host_DEVS
        while parent is not None:
            names.append(parent.name)
            try:
                parent = parent.host_DEVS
            except AttributeError:
                parent = None

        identifier = '_'.join(reversed(names))
        shorthand = names[0]
        return identifier, shorthand

    def _add_component(self, component):
        self._add(f"subgraph {component.name}")
        self._tab()
        for port in set(component.IPorts) | set(component.OPorts):
            port_name = self._get_port_name(port)
            self._add(f"{port_name[0]}(\"{port_name[1]}\")")

        # Create invisible links between the ports
        for i, iport in enumerate(component.IPorts):
            iport_name = self._get_port_name(iport)
            if len(component.OPorts) <= i:
                break
            oport_name = self._get_port_name(component.OPorts[i])
            self._add(f"{iport_name[0]} ~~~ {oport_name[0]}")

        if hasattr(component, "component_set"):
            for sub_component in component.component_set:
                self._add_component(sub_component)

        self._untab()
        self._add(f"end")

        self._add_component_connections(component)

    def _add_component_connections(self, component):
        # Go over all connections and connect them to the appropriate ends.
        for port_out in component.OPorts + component.IPorts:
            port_out_name = self._get_port_name(port_out)
            for port_in in port_out.outline:
                port_in_name = self._get_port_name(port_in)
                # Add this connection
                self._add(f"{port_out_name[0]} --> {port_in_name[0]}")

    def generate(self, orientation: str = "LR") -> str:
        """
        Generates a mermaid flowchart for your coupled-DEVS model.
        :return: a string containing the mermaid flowchart.
        """
        self._result = ""  # clear previous result

        self._add("---")
        self._add(f"title: {self.model.name}")
        self._add("---")

        self._add(f"flowchart {orientation}")
        self._tab()

        # Declare all submodels with their ports
        for component in self.model.component_set:
            self._add_component(component)

        self._untab()
        return self._result

    def generate_file(self, filepath: str = "flowchart.md", orientation: str = "LR") -> None:
        """ Generates a flowchart for a DEVS model and writes it to a file."""
        try:
            with open(filepath, 'w', encoding='utf-8') as file:
                file.write("```mermaid\n")
                file.write(self.generate(orientation))
                file.write("```")
        except OSError as e:
            print(f"Error writing to file: {e}")
            raise


if __name__ == '__main__':
    from devs_models.fischertechnik_factory import FischertechnikFactory

    mdl = FischertechnikFactory("FischertechnikFactory")
    fg = FlowchartGenerator(mdl)

    chart = fg.generate(orientation="LR")
    print(chart)