"""Visualise HUGR using graphviz."""

import html
from collections.abc import Iterable
from dataclasses import dataclass, field

import graphviz as gv  # type: ignore[import-untyped]
from graphviz import Digraph
from typing_extensions import assert_never

from hugr.hugr import Hugr
from hugr.ops import AsExtOp, Case
from hugr.tys import CFKind, ConstKind, FunctionKind, Kind, OrderKind, ValueKind

from .node_port import InPort, Node, OutPort


@dataclass(frozen=True)
class Palette:
    """A set of colours used for rendering."""

    background: str
    node: str
    edge: str
    entrypoint_edge: str
    dark: str
    const: str
    discard: str
    node_border: str
    port_border: str

    @classmethod
    def named(cls, name: str) -> "Palette":
        return PALETTE[name]


PALETTE: dict[str, Palette] = {
    "default": Palette(
        background="white",
        node="#ACCBF9",
        edge="#1CADE4",
        entrypoint_edge="#F4A261",
        dark="black",
        const="#77CEEF",
        discard="#ff8888",
        node_border="white",
        port_border="#1CADE4",
    ),
    "nb": Palette(
        background="white",
        node="#7952B3",
        edge="#FFC107",
        entrypoint_edge="#00CFC1",
        dark="#343A40",
        const="#7c55b4",
        discard="#ff8888",
        node_border="#9d80c7",
        port_border="#ffd966",
    ),
    "zx": Palette(
        background="white",
        node="#629DD1",
        edge="#297FD5",
        entrypoint_edge="#FF8243",
        dark="#112D4E",
        const="#a1eea1",
        discard="#ff8888",
        node_border="#D8F8D8",
        port_border="#E8A5A5",
    ),
}


@dataclass
class RenderConfig:
    """Configuration for rendering a HUGR to a graphviz dot file."""

    #: The palette to use for rendering. See :obj:`PALETTE` for the included options.
    palette: Palette = field(default_factory=lambda: PALETTE["default"])
    #: If true prepend extension name to operation name.
    qualify_op_name: bool = False
    #: If true display node metadata.
    display_metadata: bool = True
    #: If true display node numbering.
    display_node_id: bool = False
    #: If true display full type name on the edges.
    display_link_label: bool = True
    #: Max length for edge labels. None means no truncation.
    max_edge_label_length: int | None = 24
    #: Max length for node labels. None means no truncation.
    max_node_label_length: int | None = 24
    #: Max length for metadata display. None means no truncation.
    max_metadata_length: int | None = 20


class DotRenderer:
    """Render a HUGR to a graphviz dot file.

    Args:
        config: Render config
    """

    config: RenderConfig

    def __init__(self, config: RenderConfig | None = None) -> None:
        self.config = config or RenderConfig()
        self.nodes: set[Node] = set()

    def render(self, hugr: Hugr, root: Node | None = None) -> Digraph:
        """Render a HUGR to a graphviz dot object.

        Args:
            hugr: The HUGR to render.
            root: Root node defining the set of nodes to render. By default this is the
                module root and all nodes are rendered. If this is a container node, all
                nodes under it are rendered. Every incoming edge to the rendered set and
                outgoing edge from it is also shown, with its other endpoint labelled
                with its node index.
        """
        root = root or hugr.module_root
        graph_attr = {
            "rankdir": "",
            "ranksep": "0.1",
            "nodesep": "0.15",
            "margin": "0",
            "bgcolor": self.config.palette.background,
        }
        if name := hugr[root].metadata.get("name", None):
            name = html.escape(str(name))
        else:
            name = ""

        graph = gv.Digraph(name, strict=False)
        graph.attr(**graph_attr)

        self._viz_node(root, hugr, graph)

        for src_port, tgt_port in hugr.links():
            kind = hugr.port_kind(src_port)
            self._viz_link(src_port, tgt_port, kind, graph)

        return graph

    def store(
        self, hugr: Hugr, filename: str, format: str = "svg", root: Node | None = None
    ) -> None:
        """Render a HUGR and save it to a file.

        Args:
            hugr: The HUGR to render.
            filename: Filename for saving the rendered graph.
            format: The format used for rendering ('pdf', 'png', etc.).
                Defaults to SVG.
            root: Root node defining the set of nodes to render. By default this is the
                module root and all nodes are rendered. If this is a container node, all
                nodes under it are rendered. Every incoming edge to the rendered set and
                outgoing edge from it is also shown, with its other endpoint labelled
                with its node index.
        """
        gv_graph = self.render(hugr, root=root)
        gv_graph.render(filename, format=format)

    _FONTFACE = "monospace"

    _HTML_LABEL_TEMPLATE = """
    <TABLE BORDER="{border_width}" CELLBORDER="0" CELLSPACING="1" CELLPADDING="1"
        BGCOLOR="{node_back_color}" COLOR="{border_colour}">
    {inputs_row}
    <TR>
        <TD>
        <TABLE BORDER="0" CELLBORDER="0">
            <TR><TD><FONT POINT-SIZE="{fontsize}" FACE="{fontface}"
                COLOR="{label_color}"><B>{node_label}</B>{node_data}</FONT></TD></TR>
        </TABLE>
        </TD>
    </TR>
    {outputs_row}
    </TABLE>
    """

    _HTML_PORTS_ROW_TEMPLATE = """
        <TR>
            <TD>
                <TABLE BORDER="0" CELLBORDER="0" CELLSPACING="3" CELLPADDING="2">
                    <TR>
                        {port_cells}
                    </TR>
                </TABLE>
            </TD>
        </TR>
    """

    _HTML_PORT_TEMPLATE = (
        '<TD BGCOLOR="{back_colour}" COLOR="{border_colour}"'
        ' PORT="{port_id}" BORDER="{border_width}">'
        '<FONT POINT-SIZE="10.0" FACE="{fontface}" COLOR="{font_colour}">{port}</FONT>'
        "</TD>"
    )

    _INPUT_PREFIX = "in."
    _OUTPUT_PREFIX = "out."

    def _format_html_label(self, **kwargs: str) -> str:
        _HTML_LABEL_DEFAULTS = {
            "label_color": self.config.palette.dark,
            "node_back_color": self.config.palette.node,
            "inputs_row": "",
            "outputs_row": "",
            "border_colour": self.config.palette.port_border,
            "border_width": "1",
            "fontface": self._FONTFACE,
            "fontsize": 11.0,
        }
        return self._HTML_LABEL_TEMPLATE.format(**{**_HTML_LABEL_DEFAULTS, **kwargs})

    def _html_ports(self, ports: Iterable[str], id_prefix: str) -> str:
        return self._HTML_PORTS_ROW_TEMPLATE.format(
            port_cells="".join(
                self._HTML_PORT_TEMPLATE.format(
                    port=port,
                    # differentiate input and output node identifiers
                    # with a prefix
                    port_id=id_prefix + port,
                    back_colour=self.config.palette.background,
                    font_colour=self.config.palette.dark,
                    border_width="1",
                    border_colour=self.config.palette.port_border,
                    fontface=self._FONTFACE,
                )
                for port in ports
            )
        )

    def _in_port_name(self, p: InPort) -> str:
        return f"{p.node.idx}:{self._INPUT_PREFIX}{p.offset}"

    def _out_port_name(self, p: OutPort) -> str:
        return f"{p.node.idx}:{self._OUTPUT_PREFIX}{p.offset}"

    def _in_order_name(self, n: Node) -> str:
        return f"{n.idx}:{self._INPUT_PREFIX}None"

    def _out_order_name(self, n: Node) -> str:
        return f"{n.idx}:{self._OUTPUT_PREFIX}None"

    def _viz_node(
        self,
        node: Node,
        hugr: Hugr,
        graph: Digraph,
        *,
        sibling_order: int | None = None,
    ) -> None:
        """Render a (possibly nested) node to a graphviz graph.

        Args:
            node: The node to render.
            hugr: The HUGR to render.
            graph: The graphviz graph to render the node to.
            sibling_order: The order of the node in the region's sibling list.
        """
        meta = hugr[node].metadata
        if len(meta) > 0 and self.config.display_metadata:
            if self.config.max_metadata_length is not None:
                data = "<BR/><BR/>" + "<BR/>".join(
                    html.escape(key)
                    + ": "
                    + html.escape(
                        _smart_truncate(str(value), self.config.max_metadata_length)
                    )
                    for key, value in meta.items()
                )
            else:
                data = "<BR/><BR/>" + "<BR/>".join(
                    html.escape(key) + ": " + html.escape(str(value))
                    for key, value in meta.items()
                )
        else:
            data = ""

        in_ports = [str(i) for i in range(hugr.num_in_ports(node))]
        out_ports = [str(i) for i in range(hugr.num_out_ports(node))]
        inputs_row = (
            self._html_ports(in_ports, self._INPUT_PREFIX) if len(in_ports) > 0 else ""
        )
        outputs_row = (
            self._html_ports(out_ports, self._OUTPUT_PREFIX)
            if len(out_ports) > 0
            else ""
        )

        op = hugr[node].op
        if isinstance(op, AsExtOp) and not self.config.qualify_op_name:
            op_name = op.op_def().name
        elif isinstance(op, Case) and sibling_order is not None:
            # Indicate the case number
            op_name = f"{op.name()}[{sibling_order}]"
        else:
            op_name = op.name()

        if self.config.max_node_label_length is not None:
            op_name = _smart_truncate(
                op_name, max_length=self.config.max_node_label_length
            )

        op_name = html.escape(op_name)
        node_label = (
            f"{op_name} <BR/>({node.idx}) " if self.config.display_node_id else op_name
        )

        label_config = {
            "node_back_color": self.config.palette.node,
            "node_label": node_label,
            "node_data": data,
            "inputs_row": inputs_row,
            "outputs_row": outputs_row,
            "border_colour": self.config.palette.background,
            "border_width": "1",
        }
        if hugr.children(node):
            # Some overrides when rendering a container node
            label_config["node_back_color"] = self.config.palette.edge
            label_config["border_colour"] = self.config.palette.port_border
        if node == hugr.entrypoint:
            label_config["node_label"] = "<b>[" + label_config["node_label"] + "]</b>"
            label_config["border_colour"] = self.config.palette.entrypoint_edge
            label_config["border_width"] = "2"

        if hugr.children(node):
            with graph.subgraph(name=f"cluster{node.idx}") as sub:
                for sibling_order, child in enumerate(hugr.children(node)):
                    self._viz_node(child, hugr, sub, sibling_order=sibling_order)
                html_label = self._format_html_label(**label_config)
                sub.node(f"{node.idx}", shape="plain", label=f"<{html_label}>")
                sub.attr(
                    label="",
                    margin="10",
                    color=label_config["border_colour"],
                    penwidth=label_config["border_width"],
                )
        else:
            html_label = self._format_html_label(**label_config)
            graph.node(f"{node.idx}", label=f"<{html_label}>", shape="plain")
        self.nodes.add(node)

    def _viz_link(
        self, src_port: OutPort, tgt_port: InPort, kind: Kind, graph: Digraph
    ) -> None:
        edge_attr = {
            "penwidth": "1.5",
            "arrowhead": "none",
            "arrowsize": "1.0",
            "fontname": self._FONTFACE,
            "fontsize": "9",
            "fontcolor": "black",
        }

        label = ""
        match kind:
            case ValueKind(ty):
                if self.config.display_link_label:
                    # TODO: Figure out why some nested types result in all type fields
                    # being printed instead of just the `__str__` representation.
                    label = html.escape(str(ty))
                color = self.config.palette.edge
            case OrderKind():
                color = self.config.palette.dark
            case ConstKind() | FunctionKind():
                color = self.config.palette.const
            case CFKind():
                color = self.config.palette.dark
            case _:
                assert_never(kind)

        src = self._out_port_name(src_port)
        tgt = self._in_port_name(tgt_port)

        unknown_src = src_port.node not in self.nodes
        unknown_tgt = tgt_port.node not in self.nodes
        if unknown_src and unknown_tgt:
            return
        if unknown_src:
            src = f"{src_port.node.idx}"
            html_label = self._format_html_label(
                node_back_color=self.config.palette.node,
                node_label=f"{src_port.node}",
                node_data="",
                border_colour=self.config.palette.background,
                border_width="1",
            )
            graph.node(src, label=f"<{html_label}>", shape="plain")
        if unknown_tgt:
            tgt = f"{tgt_port.node.idx}"
            html_label = self._format_html_label(
                node_back_color=self.config.palette.node,
                node_label=f"{tgt_port.node}",
                node_data="",
                border_colour=self.config.palette.background,
                border_width="1",
            )
            graph.node(tgt, label=f"<{html_label}>", shape="plain")

        if self.config.max_edge_label_length is not None:
            label = _smart_truncate(
                label=label, max_length=self.config.max_edge_label_length
            )

        graph.edge(src, tgt, xlabel=label, color=color, **edge_attr)


def _smart_truncate(label: str, max_length: int = 24) -> str:
    """Truncate a string for display.

    Rules:
    - If the string fits, return it unchanged.
        - If truncation is needed, truncate only at delimiters (space or bracket chars).
    - Never truncate the first word (the initial token before a delimiter), even if it
      exceeds `max_length`.
        - If truncation happens inside unmatched brackets, append enough closing
            brackets after the ellipsis to close any still-open ones.
            Supported bracket pairs are `()`, `{}`, `[]`, `<>`.
    """
    if len(label) <= max_length:
        return label

    ellipsis = "..."

    delimiters = {
        " ",
        "(",
        ")",
        "{",
        "}",
        "[",
        "]",
        "<",
        ">",
        ".",
    }
    opening_brackets = {"(", "{", "[", "<"}

    # Identify the protected prefix (never truncate this), which is at least the
    # first token, and optionally also the next token if the string starts with
    # `Token<token` / `Token(token` / etc.
    first_delim_idx = next(
        (i for i, ch in enumerate(label) if ch in delimiters),
        len(label),
    )
    protected_end = first_delim_idx
    if first_delim_idx < len(label) and label[first_delim_idx] in opening_brackets:
        second_start = first_delim_idx + 1
        second_delim_idx = next(
            (
                i
                for i, ch in enumerate(label[second_start:], start=second_start)
                if ch in delimiters
            ),
            len(label),
        )
        protected_end = second_delim_idx

    prefix = label[:protected_end].rstrip()

    if prefix == label:
        return label

    truncated = prefix
    bracket_pairs = [("(", ")"), ("{", "}"), ("[", "]"), ("<", ">")]
    # we first try to close parenthesis without the ellipsis
    # and check if it is the same as the original label
    for open_ch, close_ch in bracket_pairs:
        balance = prefix.count(open_ch) - prefix.count(close_ch)
        if balance > 0:
            truncated += close_ch * balance

    if truncated == label:
        return label
    # If not, we add the ellipsis and then close any still-open parenthesis
    truncated_w_ellipsis = prefix + ellipsis

    for open_ch, close_ch in bracket_pairs:
        balance = prefix.count(open_ch) - prefix.count(close_ch)
        if balance > 0:
            truncated_w_ellipsis += close_ch * balance
    return truncated_w_ellipsis
