"""
tfsm.extensions.diagrams
-------------------------------

Mermaid support for (nested) machines. This also includes partial views
of currently valid tfsm.
"""

import copy
import logging
from collections import defaultdict
from typing import TYPE_CHECKING, Any, Union

from .diagrams_base import BaseGraph, ContextManagerMachine
from .diagrams_graphviz import filter_states

if TYPE_CHECKING:
    from .diagrams import GraphMachine

_LOGGER = logging.getLogger(__name__)
_LOGGER.addHandler(logging.NullHandler())


class Graph(BaseGraph):
    """Graph creation for tfsm.core.Machine.
    Attributes:
        custom_styles (dict): A dictionary of styles for the current graph
    """

    def __init__(self, machine: Union["GraphMachine", ContextManagerMachine]) -> None:
        self.custom_styles: dict[str, dict[str, Any]] = {}
        self.reset_styling()
        super().__init__(machine)

    def set_previous_transition(self, src: str, dst: str) -> None:
        self.custom_styles["edge"][src][dst] = "previous"
        self.set_node_style(src, "previous")

    def set_node_style(self, state: Any, style: str) -> None:
        self.custom_styles["node"][state.name if hasattr(state, "name") else state] = style

    def reset_styling(self) -> None:
        self.custom_styles = {
            "edge": defaultdict(lambda: defaultdict(str)),
            "node": defaultdict(str),
        }

    def _add_nodes(self, states: list[Any], container: list[str]) -> None:
        for state in states:
            container.append('state "{}" as {}'.format(self._convert_state_attributes(state), state["name"]))
            container.append("Class {} s_{}".format(state["name"], self.custom_styles["node"][state["name"]] or "default"))

    def _add_edges(self, transitions: list[dict[str, Any]], container: list[str]) -> None:
        edge_labels: defaultdict[str, defaultdict[str, list[str]]] = defaultdict(lambda: defaultdict(list))
        for transition in transitions:
            try:
                dst = transition["dest"]
            except KeyError:
                dst = transition["source"]
            edge_labels[transition["source"]][dst].append(self._transition_label(transition))
        for src, dests in edge_labels.items():
            for dst, labels in dests.items():
                container.append("{} --> {}: {}".format(src, dst, " | ".join(labels)))

    def generate(self) -> None:
        """Triggers the generation of a graph. With graphviz backend, this does nothing since graph trees need to be
        built from scratch with the configured styles.
        """
        # we cannot really generate a graph in advance with graphviz

    def get_graph(self, title: str | None = None, roi_state: Any | None = None) -> "DigraphMock":
        title = title if title else self.machine.title  # type: ignore[union-attr]

        fsm_graph = ["---", title, "---", "stateDiagram-v2"]
        fsm_graph.extend(_to_mermaid(self.machine.machine_attributes, " "))  # type: ignore[union-attr]

        for style_name, style_attrs in self.machine.style_attributes["node"].items():  # type: ignore[union-attr]
            if style_name:
                fsm_graph.append("classDef s_{} {}".format(style_name, ",".join(_to_mermaid(style_attrs, ":"))))
        fsm_graph.append("")
        states, transitions = self._get_elements()
        if roi_state:
            active_states = set()
            sep = getattr(self.machine.state_cls, "separator", None)
            for state in self._flatten(roi_state):
                active_states.add(state)
                if sep:
                    state = sep.join(state.split(sep)[:-1])
                    while state:
                        active_states.add(state)
                        state = sep.join(state.split(sep)[:-1])
            transitions = [t for t in transitions if t["source"] in active_states or self.custom_styles["edge"][t["source"]][t["dest"]]]
            active_states = active_states.union({t for trans in transitions for t in [trans["source"], trans.get("dest", trans["source"])]})
            active_states = active_states.union({k for k, style in self.custom_styles["node"].items() if style})
            states = filter_states(copy.deepcopy(states), active_states, self.machine.state_cls)
        self._add_nodes(states, fsm_graph)
        fsm_graph.append("")
        self._add_edges(transitions, fsm_graph)
        if self.machine.initial and (roi_state is None or roi_state == self.machine.initial):  # type: ignore[union-attr]
            fsm_graph.append(f"[*] --> {self.machine.initial}")  # type: ignore[union-attr]

        indent = 0
        for i in range(len(fsm_graph)):
            next_indent = indent
            if fsm_graph[i].startswith("stateDiagram") or fsm_graph[i].endswith("{"):
                next_indent += 2
            elif fsm_graph[i].startswith("}"):
                indent -= 2
                next_indent -= 2
            fsm_graph[i] = " " * indent + fsm_graph[i]
            indent = next_indent

        return DigraphMock("\n".join(fsm_graph))

    def _convert_state_attributes(self, state: dict[str, Any]) -> str:
        label: Any = state.get("label", state["name"])
        if self.machine.show_state_attributes:
            if "tags" in state:
                label += " [" + ", ".join(state["tags"]) + "]"
            if "on_enter" in state:
                label += r"\n- enter:\n  + " + r"\n  + ".join(state["on_enter"])
            if "on_exit" in state:
                label += r"\n- exit:\n  + " + r"\n  + ".join(state["on_exit"])
            if "timeout" in state:
                label += r"\n- timeout(" + state["timeout"] + "s) -> (" + ", ".join(state["on_timeout"]) + ")"
        # end each label with a left-aligned newline
        return str(label)


class NestedGraph(Graph):
    """Graph creation support for tfsm.extensions.nested.HierarchicalGraphMachine."""

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self._cluster_states: list[str] = []
        super().__init__(*args, **kwargs)

    def set_node_style(self, state: Any, style: str) -> None:
        for state_name in self._get_state_names(state):
            super().set_node_style(state_name, style)

    def set_previous_transition(self, src: str, dst: str) -> None:
        self.custom_styles["edge"][src][dst] = "previous"
        self.set_node_style(src, "previous")

    def _add_nodes(self, states: list[Any], container: list[str]) -> None:
        self._add_nested_nodes(states, container, prefix="", default_style="default")

    def _add_nested_nodes(self, states: list[Any], container: list[str], prefix: str, default_style: str) -> None:
        for state in states:
            name = prefix + state["name"]
            container.append(f'state "{self._convert_state_attributes(state)}" as {name}')
            if state.get("final", False):
                container.append(f"{name} --> [*]")
            if not prefix:
                container.append("Class {} s_{}".format(name.replace(" ", ""), self.custom_styles["node"][name] or default_style))
            if state.get("children", None) is not None:
                container.append(f"state {name} {{")
                self._cluster_states.append(name)
                # with container.subgraph(name=cluster_name, graph_attr=attr) as sub:
                initial = state.get("initial", "")
                is_parallel = isinstance(initial, list)
                if is_parallel:
                    for child in state["children"]:
                        self._add_nested_nodes(
                            [child],
                            container,
                            default_style="parallel",
                            prefix=prefix + state["name"] + getattr(self.machine.state_cls, "separator", "_"),
                        )
                        container.append("--")
                    if state["children"]:
                        container.pop()
                else:
                    if initial:
                        container.append(
                            "[*] --> {}".format(prefix + state["name"] + getattr(self.machine.state_cls, "separator", "_") + initial)
                        )
                    self._add_nested_nodes(
                        state["children"],
                        container,
                        default_style="default",
                        prefix=prefix + state["name"] + getattr(self.machine.state_cls, "separator", "_"),
                    )
                container.append("}")

    def _add_edges(self, transitions: list[dict[str, Any]], container: list[str]) -> None:
        edges_attr: defaultdict[str, defaultdict[str, dict[str, Any]]] = defaultdict(lambda: defaultdict(dict))

        for transition in transitions:
            # enable customizable labels
            src = transition["source"]
            dst = transition.get("dest", src)
            if edges_attr[src][dst]:
                attr = edges_attr[src][dst]
                attr["label"] = " | ".join([edges_attr[src][dst]["label"], self._transition_label(transition)])
            else:
                edges_attr[src][dst] = self._create_edge_attr(src, dst, transition)

        for custom_src, dests in self.custom_styles["edge"].items():
            for custom_dst, style in dests.items():
                if style and (custom_src not in edges_attr or custom_dst not in edges_attr[custom_src]):
                    edges_attr[custom_src][custom_dst] = self._create_edge_attr(custom_src, custom_dst, {"trigger": "", "dest": ""})

        for src, dests in edges_attr.items():
            for dst, attr in dests.items():
                if not attr["label"]:
                    continue
                container.append("{source} --> {dest}: {label}".format(**attr))

    def _create_edge_attr(self, src: str, dst: str, transition: dict[str, Any]) -> dict[str, Any]:
        return {"source": src, "dest": dst, "label": self._transition_label(transition)}


class DigraphMock:
    def __init__(self, source: str) -> None:
        self.source: str = source

    # pylint: disable=redefined-builtin,unused-argument
    def draw(self, filename: Any, format: str | None = None, prog: str = "dot", args: str = "") -> str | None:
        """
        Generates and saves an image of the state machine using graphviz. Note that `prog` and `args` are only part
        of the signature to mimic `Agraph.draw` and thus allow to easily switch between graph backends.
        Args:
            filename (str or file descriptor or stream or None): path and name of image output, file descriptor,
            stream object or None
            format (str): ignored
            prog (str): ignored
            args (str): ignored
        Returns:
            None or str: Returns a binary string of the graph when the first parameter (`filename`) is set to None.
        """

        if filename is None:
            return self.source
        if isinstance(filename, str):
            with open(filename, "w") as f:
                f.write(self.source)
        else:
            filename.write(self.source.encode())
        return None


invalid = {"style", "shape", "peripheries", "strict", "directed"}
convertible = {"fillcolor": "fill", "rankdir": "direction"}


def _to_mermaid(style_attrs: dict[str, Any], sep: str) -> list[str]:
    params = []
    for k, v in style_attrs.items():
        if k in invalid:
            continue
        if k in convertible:
            k = convertible[k]
        params.append(f"{k}{sep}{v}")
    return params
