# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

from __future__ import annotations

import subprocess
import tempfile
import io
from typing import TypeVar, Callable, cast, TYPE_CHECKING

from rustworkx import PyDiGraph, PyGraph


try:
    from PIL import Image  # type: ignore

    HAS_PILLOW = True
except ImportError:
    HAS_PILLOW = False

if TYPE_CHECKING:
    from PIL import Image  # type: ignore

_S = TypeVar("_S")
_T = TypeVar("_T")


__all__ = ["graphviz_draw"]

METHODS = {"twopi", "neato", "circo", "fdp", "sfdp", "dot"}
IMAGE_TYPES = {
    "canon",
    "cmap",
    "cmapx",
    "cmapx_np",
    "dia",
    "dot",
    "fig",
    "gd",
    "gd2",
    "gif",
    "hpgl",
    "imap",
    "imap_np",
    "ismap",
    "jpe",
    "jpeg",
    "jpg",
    "mif",
    "mp",
    "pcl",
    "pdf",
    "pic",
    "plain",
    "plain-ext",
    "png",
    "ps",
    "ps2",
    "svg",
    "svgz",
    "vml",
    "vmlz",
    "vrml",
    "vtx",
    "wbmp",
    "xdor",
    "xlib",
}


def graphviz_draw(
    graph: "PyDiGraph[_S, _T] | PyGraph[_S, _T]",  # noqa
    node_attr_fn: Callable[[_S], dict[str, str]] | None = None,
    edge_attr_fn: Callable[[_T], dict[str, str]] | None = None,
    graph_attr: dict[str, str] | None = None,
    filename: str | None = None,
    image_type: str | None = None,
    method: str | None = None,
) -> Image | None:
    """Draw a :class:`~rustworkx.PyGraph` or :class:`~rustworkx.PyDiGraph` object
    using graphviz

    .. note::

        This requires that pydot, pillow, and graphviz be installed. Pydot can
        be installed via pip with ``pip install pydot pillow`` however graphviz
        will need to be installed separately. You can refer to the
        Graphviz
        `documentation <https://graphviz.org/download/#executable-packages>`__
        for instructions on how to install it.

    :param graph: The rustworkx graph object to draw, can be a
        :class:`~rustworkx.PyGraph` or a :class:`~rustworkx.PyDiGraph`
    :param node_attr_fn: An optional callable object that will be passed the
        weight/data payload for every node in the graph and expected to return
        a dictionary of Graphviz node attributes to be associated with the node
        in the visualization. The key and value of this dictionary **must** be
        a string.
    :param edge_attr_fn: An optional callable that will be passed the
        weight/data payload for each edge in the graph and expected to return a
        dictionary of Graphviz edge attributes to be associated with the edge
        in the visualization file. The key and value of this dictionary
        must be a string.
    :param dict graph_attr: An optional dictionary that specifies any Graphviz
        graph attributes for the visualization. The key and value of this
        dictionary must be a string.
    :param str filename: An optional path to write the visualization to. If
        specified the return type from this function will be ``None`` as the
        output image is saved to disk.
    :param str image_type: The image file format to use for the generated
        visualization. The support image formats are:
        ``'canon'``, ``'cmap'``, ``'cmapx'``, ``'cmapx_np'``, ``'dia'``,
        ``'dot'``, ``'fig'``, ``'gd'``, ``'gd2'``, ``'gif'``, ``'hpgl'``,
        ``'imap'``, ``'imap_np'``, ``'ismap'``, ``'jpe'``, ``'jpeg'``,
        ``'jpg'``, ``'mif'``, ``'mp'``, ``'pcl'``, ``'pdf'``, ``'pic'``,
        ``'plain'``, ``'plain-ext'``, ``'png'``, ``'ps'``, ``'ps2'``,
        ``'svg'``, ``'svgz'``, ``'vml'``, ``'vmlz'``, ``'vrml'``, ``'vtx'``,
        ``'wbmp'``, ``'xdot'``, ``'xlib'``. It's worth noting that while these
        formats can all be used for generating image files when the ``filename``
        kwarg is specified, the Pillow library used for the returned object can
        not work with all these formats.
    :param str method: The layout method/Graphviz command method to use for
        generating the visualization. Available options are ``'dot'``,
        ``'twopi'``, ``'neato'``, ``'circo'``, ``'fdp'``, and ``'sfdp'``.
        You can refer to the
        `Graphviz documentation <https://graphviz.org/documentation/>`__ for
        more details on the different layout methods. By default ``'dot'`` is
        used.

    :returns: A ``PIL.Image`` object of the generated visualization, if
        ``filename`` is not specified. If ``filename`` is specified then
        ``None`` will be returned as the visualization was written to the
        path specified in ``filename``
    :rtype: PIL.Image

    .. jupyter-execute::

        import rustworkx as rx
        from rustworkx.visualization import graphviz_draw

        def node_attr(node):
          if node == 0:
            return {'color': 'yellow', 'fillcolor': 'yellow', 'style': 'filled'}
          if node % 2:
            return {'color': 'blue', 'fillcolor': 'blue', 'style': 'filled'}
          else:
            return {'color': 'red', 'fillcolor': 'red', 'style': 'filled'}

        graph = rx.generators.directed_star_graph(weights=list(range(32)))
        graphviz_draw(graph, node_attr_fn=node_attr, method='sfdp')

    """
    if not HAS_PILLOW:
        raise ImportError(
            "Pillow is necessary to use graphviz_draw() "
            "it can be installed with 'pip install pydot pillow'"
        )
    try:
        subprocess.run(
            ["dot", "-V"],
            cwd=tempfile.gettempdir(),
            check=True,
            capture_output=True,
        )
    except Exception:
        raise RuntimeError(
            "Graphviz could not be found or run. This function requires that "
            "Graphviz is installed. If you need to install Graphviz you can "
            "refer to: https://graphviz.org/download/#executable-packages for "
            "instructions."
        )

    dot_str = cast(str, graph.to_dot(node_attr_fn, edge_attr_fn, graph_attr))
    if image_type is None:
        output_format = "png"
    else:
        if image_type not in IMAGE_TYPES:
            raise ValueError(
                "The specified value for the image_type argument, "
                f"'{image_type}' is not a valid choice. It must be one of: "
                f"{IMAGE_TYPES}"
            )
        output_format = image_type

    if method is None:
        prog = "dot"
    else:
        if method not in METHODS:
            raise ValueError(
                f"The specified value for the method argument, '{method}' is "
                f"not a valid choice. It must be one of: {METHODS}"
            )
        prog = method

    if not filename:
        dot_result = subprocess.run(
            [prog, "-T", output_format],
            input=dot_str.encode("utf-8"),
            capture_output=True,
            encoding=None,
            check=True,
            text=False,
        )
        dot_bytes_image = io.BytesIO(dot_result.stdout)
        image = Image.open(dot_bytes_image)
        return image
    else:
        subprocess.run(
            [prog, "-T", output_format, "-o", filename],
            input=dot_str,
            check=True,
            encoding="utf8",
            text=True,
        )
        return None