from collections import defaultdict
from io import BytesIO
import logging
from pydoc import locate
from typing import Any, Dict, Optional

import ray
from ray import cloudpickle
from ray.experimental.gradio_utils import type_to_string

from ray.dag import (
    DAGNode,
    InputNode,
    InputAttributeNode,
)
from ray.dag.utils import _DAGNodeNameGenerator
from ray.dag.vis_utils import _dag_to_dot

from ray.serve.handle import RayServeHandle
from ray.serve._private.deployment_function_executor_node import (
    DeploymentFunctionExecutorNode,
)
from ray.serve._private.deployment_method_executor_node import (
    DeploymentMethodExecutorNode,
)


logger = logging.getLogger(__name__)
_gradio = None


def lazy_import_gradio():
    global _gradio
    if _gradio is None:
        try:
            import gradio
        except ModuleNotFoundError:
            logger.error(
                "Gradio isn't installed. Run `pip install gradio` to use Gradio to "
                "visualize a Serve deployment graph."
            )
            raise

        _gradio = gradio
    return _gradio


def get_block(node, **kwargs):
    gr = lazy_import_gradio()

    result_type = node.get_result_type()
    node_name = _DAGNodeNameGenerator().get_node_name(node)
    if result_type is None:
        logger.warning(
            f'Return type for "{node_name}" node was not provided. '
            "Defaulting to gr.Textbox."
        )
        return gr.Textbox(**kwargs)

    # Type string output generated by ray.experimental.gradio_utils.type_to_string()
    if result_type == "int":
        return gr.Number(precision=0, **kwargs)
    elif result_type == "float":
        return gr.Number(**kwargs)
    elif result_type == "str":
        return gr.Textbox(**kwargs)
    elif result_type == "bool":
        return gr.Checkbox(**kwargs)
    elif result_type == "pandas.core.frame.DataFrame":
        return gr.Dataframe(**kwargs)
    elif (
        result_type == "list"
        or result_type == "dict"
        or result_type == "typing.List"
        or result_type == "typing.Dict"
        or result_type == "numpy.ndarray"
    ):
        return gr.JSON(**kwargs)
    elif result_type == "torch.tensor":
        return gr.Image(**kwargs)
    else:
        # Lazy import of PIL. Since PIL is a dependency of Gradio, PIL should be
        # installed if Gradio was successfully imported.
        from PIL import ImageFile

        if issubclass(locate(result_type), ImageFile.ImageFile):
            return gr.Image(**kwargs)

    logger.warning(
        f"Return type for {node_name} node is not valid. Defaulting to gr.Textbox."
    )
    return gr.Textbox(**kwargs)


def postprocessing(data):
    """Add support for types that are not supported by Gradio.

    Some data types like PyTorch tensors, cannot be processed and displayed through
    Gradio. Thus we extend support to these data types by transforming them into a form
    that Gradio can process and display.
    """

    if type_to_string(type(data)) == "torch.Tensor":
        try:
            import torch
            from torchvision import transforms

            # By default Torch tensors are displayed as images. To display them as JSON,
            # the user can simply convert them to numpy arrays.
            transformer = transforms.ToPILImage()
            return transformer(torch.squeeze(data))
        except ModuleNotFoundError:
            logger.warning(
                "Module `torchvision` isn't installed, unable to process torch tensor."
            )
            return data

    return data


class GraphVisualizer:
    def __init__(self):
        lazy_import_gradio()
        self._reset_state()

    def _reset_state(self):
        """Resets state for each new RayServeHandle representing a new DAG."""
        self.cache = {}
        self.resolved_nodes = 0
        self.finished_last_inference = True

        # Whether user created InputAttributeNodes when building the graph. It is
        # assumed that input is either given through the single InputNode, or
        # exclusively through InputAttributeNodes.
        self.take_input_through_attribute_nodes = False

        # maps DAGNode to unique instance of a gradio block
        self.node_to_block: Dict[DAGNode, Any] = {}
        # maps input nodes to unique instance of interactive gradio block
        self.input_node_to_block: Dict[int, Any] = {}

    def clear_cache(self):
        self.cache = {}

    def _make_blocks(self, depths: Dict[str, int]) -> None:
        """Instantiates Gradio blocks for each graph node stored in depths.

        Nodes of depth 0 will be rendered in the top row, depth 1 in the next row,
        and so forth. Note that this will render either the single InputNode, or all
        of the InputAttributeNodes; it will not render a mix of the two.

        Args:
            depths: maps uuids of nodes in the DAG to their depth
        """
        gr = lazy_import_gradio()
        input_node_type_to_render = (
            InputAttributeNode if self.take_input_through_attribute_nodes else InputNode
        )

        levels = {}
        for node in depths:
            if isinstance(
                node,
                (
                    input_node_type_to_render,
                    DeploymentMethodExecutorNode,
                    DeploymentFunctionExecutorNode,
                ),
            ):
                levels.setdefault(depths[node], []).append(node)

        node_names = _DAGNodeNameGenerator()

        def render_level(level):
            for node in levels[level]:
                name = node_names.get_node_name(node)

                if isinstance(node, input_node_type_to_render):
                    self.input_node_to_block[node] = get_block(node, label=name)
                else:
                    self.node_to_block[node] = get_block(
                        node, label=name, interactive=False
                    )

        for level in sorted(levels.keys()):
            with gr.Row():
                render_level(level)

    def _fetch_depths(self, node: DAGNode, depths: Dict[str, int]) -> DAGNode:
        """Gets the node's depth.

        Calculates graph node's depth, which is determined by the longest distance
        between that node and the InputNode. The single InputNode in the graph will have
        depth 0, and any InputAttributeNodes, if they exist, will have depth 1. The
        node's depth is cached in the passed-in depths dictionary.

        Args:
            node: the graph node to process
            depths: map between DAGNode uuid to the current longest found distance
                between the DAGNode and any input nodes

        Returns:
            The original node. After apply_recursive is done, the cache will store
            an uuid -> node map, which will be used in make_blocks.
        """
        if isinstance(node, InputAttributeNode):
            self.take_input_through_attribute_nodes = True

        uuid = node.get_stable_uuid()
        for child_node in node._get_all_child_nodes():
            depths[uuid] = max(depths[uuid], depths[child_node.get_stable_uuid()] + 1)

        return node

    async def _get_result(self, node_uuid: str):
        """Retrieves DAGNode's last execution output.

        This function should only be called after a request has been sent through
        self._send_request() separately.
        """
        try:
            result = await self.cache[node_uuid]
            self.resolved_nodes += 1
            if self.resolved_nodes == len(self.node_to_block):
                self.finished_last_inference = True
            return postprocessing(result)
        except Exception:
            self.resolved_nodes += 1
            if self.resolved_nodes == len(self.node_to_block):
                self.finished_last_inference = True
            raise

    async def _send_request(self, trigger_value: int, *input_values) -> int:
        """Sends request to the graph and gets results.

        Sends a request to the root DAG node through self.handle and retrieves the
        cached object refs pointing to return values of each executed node in the DAG.
        Will not run if the last inference process has not finished (if all nodes in
        DAG have been resolved).

        Args:
            trigger_value: The current value of the `trigger` Gradio block.
            *input_values: The user input values from Gradio input blocks.

        Returns:
            The current value of the `trigger` Gradio block incremented by 1. The
            `trigger` Gradio block will change to this value, and trigger the other
            Gradio blocks to run their respective _get_result functions.
        """
        if not self.finished_last_inference:
            logger.warning("Last inference has not finished yet.")
            return trigger_value

        # Extract positional args
        args, kwargs = [], {}
        if self.take_input_through_attribute_nodes:
            # Assumes self.input_node_to_block is an ordered dictionary
            input_keys = [node._key for node in self.input_node_to_block]

            indices = [i for i in input_keys if isinstance(i, int)]
            if len(indices):
                max_index = max(indices)
                for i in range(max_index + 1):
                    try:
                        loc = input_keys.index(i)
                        args.append(input_values[loc])
                    except ValueError:
                        args.append(None)

            # Extract keyword args
            kwargs = {
                input_keys[i]: input_values[i]
                for i in range(len(input_values))
                if isinstance(input_keys[i], str)
            }
        else:
            args = input_values

        self.handle.predict.remote(*args, _ray_cache_refs=True, **kwargs)
        self.cache = await self.handle.get_intermediate_object_refs.remote()

        # Set state to track the inference process
        self.resolved_nodes = 0
        self.finished_last_inference = False

        return trigger_value + 1

    def visualize_with_gradio(
        self,
        driver_handle: RayServeHandle,
        port: Optional[int] = None,
        _launch: bool = True,
        _block: bool = True,
    ):
        """Starts deployment graph's Gradio UI.

        Launches a Gradio UI that allows interactive request dispatch and displays
        the evaluated outputs of each node in a deployment graph in real time.

        Args:
            driver_handle: The handle to a DAGDriver deployment obtained through a call
                to serve.run(). The DAG rooted at that DAGDriver deployment will be
                visualized through Gradio.
            port: The port on which to start the Gradio app. If None, will default to
                Gradio's default.
            _launch: Whether to launch the Gradio app. Used for unit testing purposes.
            _block: Whether to block the main thread while the Gradio server is running.
                Used for unit testing purposes.
        """
        gr = lazy_import_gradio()

        self._reset_state()
        self.handle = driver_handle

        # Load the root DAG node from handle.
        pickled_dag_node = ray.get(self.handle.get_pickled_dag_node.remote())
        self.dag = cloudpickle.loads(pickled_dag_node)

        # Get level for each node in dag
        uuid_to_depths = defaultdict(lambda: 0)

        def depths_fn(node):
            return self._fetch_depths(node, uuid_to_depths)

        self.dag.apply_recursive(depths_fn)
        node_to_depths = {
            depths_fn.cache[uuid]: uuid_to_depths[uuid] for uuid in uuid_to_depths
        }

        with gr.Blocks() as demo:
            from PIL import Image

            try:
                graph = _dag_to_dot(self.dag)
                gr.Image(
                    label="Ray Serve Deployment Graph",
                    value=Image.open(BytesIO(graph.create(graph.prog, format="png"))),
                )
            except ImportError:
                gr.Markdown(
                    "## Warning: cannot show graph illustration.\n"
                    "Python module `pydot` and package `graphviz` is needed to show "
                    "graph illustration. Install pydot with `pip install pydot` and "
                    "graphviz with either `brew install pydot` or `sudo apt install "
                    "graphviz`."
                )

            self._make_blocks(node_to_depths)

            with gr.Row():
                submit = gr.Button("Run").style()
                trigger = gr.Number(0, visible=False)
                clear = gr.Button("Clear").style()

            # Add event listener that sends the request to the deployment graph
            submit.click(
                fn=self._send_request,
                inputs=[trigger] + list(self.input_node_to_block.values()),
                outputs=trigger,
            )
            # Add event listeners that resolve object refs for each of the nodes
            for node, block in self.node_to_block.items():
                trigger.change(
                    self._get_result, gr.Variable(node.get_stable_uuid()), block
                )

            # Resets all blocks if Clear button is clicked
            all_blocks = [*self.node_to_block.values()] + [
                *self.input_node_to_block.values()
            ]
            clear.click(
                lambda: self.clear_cache() or [None] * len(all_blocks), [], all_blocks
            )

        if _launch:
            return demo.launch(
                show_error=True, server_port=port, prevent_thread_lock=not _block
            )
