import os
import sys

import pydot
import pytest

import ray
from ray.dag import InputNode, MultiOutputNode
from ray.tests.conftest import *  # noqa


@pytest.fixture
def cleanup_files():
    """Clean up files generated during the test."""

    def _cleanup_files(filename: str):
        for ext in ["", ".png", ".pdf", ".jpeg", ".dot"]:
            file_path = filename + ext
            if os.path.exists(file_path):
                os.remove(file_path)

    return _cleanup_files


def test_visualize_basic(ray_start_regular, cleanup_files):
    """
    Expect output or dot_source:
        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
            0 -> 1 [label=SharedMemoryType]
            1 -> 2 [label=SharedMemoryType]
    """

    @ray.remote
    class Actor:
        def echo(self, x):
            return x

    actor = Actor.remote()

    with InputNode() as i:
        dag = actor.echo.bind(i)

    compiled_dag = dag.experimental_compile()

    # Call the visualize method
    dot_source = compiled_dag.visualize()

    graphs = pydot.graph_from_dot_data(dot_source)
    graph = graphs[0]

    node_names = {node.get_name() for node in graph.get_nodes()}
    edge_pairs = {
        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
    }

    expected_nodes = {"0", "1", "2"}
    assert expected_nodes.issubset(
        node_names
    ), f"Expected nodes {expected_nodes} not found."

    expected_edges = {("0", "1"), ("1", "2")}
    assert expected_edges.issubset(
        edge_pairs
    ), f"Expected edges {expected_edges} not found."

    cleanup_files("compiled_graph")


def test_visualize_multi_return(ray_start_regular, cleanup_files):
    """
    Expect output or dot_source:
        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
            0 -> 1 [label=SharedMemoryType]
            1 -> 2 [label=SharedMemoryType]
            1 -> 3 [label=SharedMemoryType]
            2 -> 4 [label=SharedMemoryType]
            3 -> 4 [label=SharedMemoryType]
    """

    @ray.remote
    class Actor:
        @ray.method(num_returns=2)
        def return_two(self, x):
            return x, x + 1

    actor = Actor.remote()

    with InputNode() as i:
        o1, o2 = actor.return_two.bind(i)
        dag = MultiOutputNode([o1, o2])

    compiled_dag = dag.experimental_compile()

    # Get the DOT source
    dot_source = compiled_dag.visualize()

    graphs = pydot.graph_from_dot_data(dot_source)
    graph = graphs[0]

    node_names = {node.get_name() for node in graph.get_nodes()}
    edge_pairs = {
        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
    }

    expected_nodes = {"0", "1", "2", "3", "4"}
    assert expected_nodes.issubset(
        node_names
    ), f"Expected nodes {expected_nodes} not found."

    expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
    assert expected_edges.issubset(
        edge_pairs
    ), f"Expected edges {expected_edges} not found."
    cleanup_files("compiled_graph")


def test_visualize_multi_return2(ray_start_regular, cleanup_files):
    """
    Expect output or dot_source:
        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
            0 -> 1 [label=SharedMemoryType]
            1 -> 2 [label=SharedMemoryType]
            1 -> 3 [label=SharedMemoryType]
            2 -> 4 [label=SharedMemoryType]
            3 -> 5 [label=SharedMemoryType]
            4 -> 6 [label=SharedMemoryType]
            5 -> 6 [label=SharedMemoryType]
    """

    @ray.remote
    class Actor:
        @ray.method(num_returns=2)
        def return_two(self, x):
            return x, x + 1

        def echo(self, x):
            return x

    a = Actor.remote()
    b = Actor.remote()
    with InputNode() as i:
        o1, o2 = a.return_two.bind(i)
        o3 = b.echo.bind(o1)
        o4 = b.echo.bind(o2)
        dag = MultiOutputNode([o3, o4])

    compiled_dag = dag.experimental_compile()

    # Get the DOT source
    dot_source = compiled_dag.visualize()

    graphs = pydot.graph_from_dot_data(dot_source)
    graph = graphs[0]

    node_names = {node.get_name() for node in graph.get_nodes()}
    edge_pairs = {
        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
    }

    expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
    assert expected_nodes.issubset(
        node_names
    ), f"Expected nodes {expected_nodes} not found."

    expected_edges = {
        ("0", "1"),
        ("1", "2"),
        ("1", "3"),
        ("2", "4"),
        ("3", "5"),
        ("4", "6"),
        ("5", "6"),
    }
    assert expected_edges.issubset(
        edge_pairs
    ), f"Expected edges {expected_edges} not found."
    cleanup_files("compiled_graph")


def test_visualize_multi_input_nodes(ray_start_regular, cleanup_files):
    """
    Expect output or dot_source:
        MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
            0 -> 1
            0 -> 2
            0 -> 3
            1 -> 4
            2 -> 5
            3 -> 6
            4 -> 7
            5 -> 7
            6 -> 7
    """

    @ray.remote
    class Actor:
        def echo(self, x):
            return x

    actor = Actor.remote()

    with InputNode() as inp:
        o1 = actor.echo.bind(inp.x)
        o2 = actor.echo.bind(inp.y)
        o3 = actor.echo.bind(inp.z)
        dag = MultiOutputNode([o1, o2, o3])

    compiled_dag = dag.experimental_compile()

    # Get the DOT source
    dot_source = compiled_dag.visualize()

    graphs = pydot.graph_from_dot_data(dot_source)
    graph = graphs[0]

    node_names = {node.get_name() for node in graph.get_nodes()}
    edge_pairs = {
        (edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
    }

    expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7"}
    assert expected_nodes.issubset(
        node_names
    ), f"Expected nodes {expected_nodes} not found."

    expected_edges = {
        ("0", "1"),
        ("0", "2"),
        ("0", "3"),
        ("1", "4"),
        ("2", "5"),
        ("3", "6"),
        ("4", "7"),
        ("5", "7"),
        ("6", "7"),
    }
    assert expected_edges.issubset(
        edge_pairs
    ), f"Expected edges {expected_edges} not found."
    cleanup_files("compiled_graph")


class TestVisualizationAscii:

    """Tests for the visualize_ascii method of compiled DAGs."""

    @staticmethod
    def parse_ascii_visualization(ascii_visualization: str):
        """
        Parses the ASCII visualization output to extract node names and edge pairs.

        Args:
            ascii_visualization: The ASCII visualization
                output generated by the `visualize` function.

        Returns:
            tuple: A tuple containing:
                - node_names: A set of strings representing node names.
                - edge_pairs: A set of tuples representing edge
                pairs with type hints.
        """
        import re

        # Sets to store unique nodes and edges
        node_names = set()
        edge_pairs = set()

        # Extract nodes from "Nodes Information" section
        node_pattern = re.compile(r'^(\d+) \[label="Task \d+')
        edge_pattern = re.compile(r"^(\d+) (--->|\+\+\+>) (\d+)")

        lines = ascii_visualization.splitlines()
        in_nodes_section = False
        in_edges_section = False

        for line in lines:
            line = line.strip()

            # Check for nodes section
            if line.startswith("Nodes Information:"):
                in_nodes_section = True
                in_edges_section = False
                continue

            # Check for edges section
            if line.startswith("Edges Information:"):
                in_edges_section = True
                in_nodes_section = False
                continue

            # Collect nodes
            if in_nodes_section:
                node_match = node_pattern.match(line)
                if node_match:
                    node_id = node_match.group(1)
                    node_names.add(node_id)

            # Collect edges
            if in_edges_section:
                edge_match = edge_pattern.match(line)
                if edge_match:
                    from_node, _, to_node = edge_match.groups()
                    edge_pairs.add((from_node, to_node))

        return node_names, edge_pairs

    def test_visualize_ascii_basic(self, ray_start_regular):
        """
        Expect output:
            Nodes Information:
            0 [label="Task 0  InputNode"]
            1 [label="Task 1  Actor: d6c5c4... Method: echo"]
            2 [label="Task 2  MultiOutputNode"]

            Edges Information:
            0 ---> 1
            1 ---> 2

            Legend:
            +++> : Represents Nccl-type data channels
            ---> : Represents Shared Memory data channels

            Experimental Graph:
            0:InputNode
            |
            1:Actor_d6c5c4:echo
            |
            2:MultiOutputNode
        """

        @ray.remote
        class Actor:
            def echo(self, x):
                return x

        actor = Actor.remote()

        with InputNode() as i:
            dag = actor.echo.bind(i)

        compiled_dag = dag.experimental_compile()

        # Call the visualize method
        ascii_visualization = compiled_dag.visualize(format="ascii")

        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
            ascii_visualization
        )
        print(node_names, edge_pairs)
        expected_nodes = {"0", "1", "2"}
        assert expected_nodes.issubset(
            node_names
        ), f"Expected nodes {expected_nodes} not found."

        expected_edges = {("0", "1"), ("1", "2")}
        assert expected_edges.issubset(
            edge_pairs
        ), f"Expected edges {expected_edges} not found."

    def test_visualize_ascii_multi_return(self, ray_start_regular):
        """
        Expect output:
            Nodes Information:
            0 [label="Task 0  InputNode"]
            1 [label="Task 1  Actor: 885f1d... Method: return_two"]
            2 [label="Task 2  ClassMethodOutputNode[0]"]
            3 [label="Task 3  ClassMethodOutputNode[1]"]
            4 [label="Task 4  MultiOutputNode"]

            Edges Information:
            0 ---> 1
            1 ---> 2
            1 ---> 3
            2 ---> 4
            3 ---> 4

            Legend:
            +++> : Represents Nccl-type data channels
            ---> : Represents Shared Memory data channels

            Graph Built:
            0:InputNode
            |
            1:Actor_885f1d:return_two
            |---------------------------->|
            2:Output[0]                   3:Output[1]
            |<----------------------------|
            4:MultiOutputNode
        """

        @ray.remote
        class Actor:
            @ray.method(num_returns=2)
            def return_two(self, x):
                return x, x + 1

        actor = Actor.remote()

        with InputNode() as i:
            o1, o2 = actor.return_two.bind(i)
            dag = MultiOutputNode([o1, o2])

        compiled_dag = dag.experimental_compile()

        ascii_visualization = compiled_dag.visualize(format="ascii")

        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
            ascii_visualization
        )

        expected_nodes = {"0", "1", "2", "3", "4"}
        assert expected_nodes.issubset(
            node_names
        ), f"Expected nodes {expected_nodes} not found."

        expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
        assert expected_edges.issubset(
            edge_pairs
        ), f"Expected edges {expected_edges} not found."

    def test_visualize_ascii_multi_return2(self, ray_start_regular):
        """
        Expect output:
            Nodes Information:
            0 [label="Task 0  InputNode"]
            1 [label="Task 1  Actor: f3e919... Method: return_two"]
            2 [label="Task 2  ClassMethodOutputNode[0]"]
            3 [label="Task 3  ClassMethodOutputNode[1]"]
            4 [label="Task 4  Actor: 15ec69... Method: echo"]
            5 [label="Task 5  Actor: 15ec69... Method: echo"]
            6 [label="Task 6  MultiOutputNode"]

            Edges Information:
            0 ---> 1
            1 ---> 2
            1 ---> 3
            2 ---> 4
            3 ---> 5
            4 ---> 6
            5 ---> 6

            Legend:
            +++> : Represents Nccl-type data channels
            ---> : Represents Shared Memory data channels

            Graph Built:
            0:InputNode
            |
            1:Actor_f3e919:return_two
            |---------------------------->|
            2:Output[0]                   3:Output[1]
            |                             |
            4:Actor_15ec69:echo           5:Actor_15ec69:echo
            |<----------------------------|
            6:MultiOutputNode
        """

        @ray.remote
        class Actor:
            @ray.method(num_returns=2)
            def return_two(self, x):
                return x, x + 1

            def echo(self, x):
                return x

        a = Actor.remote()
        b = Actor.remote()
        with InputNode() as i:
            o1, o2 = a.return_two.bind(i)
            o3 = b.echo.bind(o1)
            o4 = b.echo.bind(o2)
            dag = MultiOutputNode([o3, o4])

        compiled_dag = dag.experimental_compile()

        ascii_visualization = compiled_dag.visualize(format="ascii")

        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
            ascii_visualization
        )

        expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
        assert expected_nodes.issubset(
            node_names
        ), f"Expected nodes {expected_nodes} not found."

        expected_edges = {
            ("0", "1"),
            ("1", "2"),
            ("1", "3"),
            ("2", "4"),
            ("3", "5"),
            ("4", "6"),
            ("5", "6"),
        }
        assert expected_edges.issubset(
            edge_pairs
        ), f"Expected edges {expected_edges} not found."

    def test_visualize_ascii_complicate(self, ray_start_regular):
        """
        Expect output:
            Nodes Information:
            0 [label="Task 0  InputNode"]
            1 [label="Task 1  Actor: 54777d... Method: return_three"]
            2 [label="Task 2  ClassMethodOutputNode[0]"]
            3 [label="Task 3  ClassMethodOutputNode[1]"]
            4 [label="Task 4  ClassMethodOutputNode[2]"]
            5 [label="Task 5  Actor: c927c9... Method: echo"]
            6 [label="Task 6  Actor: c927c9... Method: echo"]
            7 [label="Task 7  Actor: c927c9... Method: return_two"]
            8 [label="Task 8  MultiOutputNode"]
            9 [label="Task 9  ClassMethodOutputNode[0]"]
            10 [label="Task 10  ClassMethodOutputNode[1]"]

            Edges Information:
            0 ---> 1
            1 ---> 2
            1 ---> 3
            1 ---> 4
            2 ---> 5
            3 ---> 6
            4 ---> 7
            5 ---> 8
            6 ---> 8
            9 ---> 8
            10 ---> 8
            7 ---> 9
            7 ---> 10

            Legend:
            +++> : Represents Nccl-type data channels
            ---> : Represents Shared Memory data channels

            Graph Built:
            0:InputNode
            |
            1:Actor_54777d:return_three
            |---------------------------->|---------------------------->|                                                  # noqa
            2:Output[0]                   3:Output[1]                   4:Output[2]                                        # noqa
            |                             |                             |                                                  # noqa
            5:Actor_c927c9:echo           6:Actor_c927c9:echo           7:Actor_c927c9:return_two                          # noqa
            |                             |                             |---------------------------->|                    # noqa
            |                             |                             9:Output[0]                   10:Output[1]         # noqa
            |<----------------------------|-----------------------------|-----------------------------|                    # noqa
            8:MultiOutputNode
        """

        @ray.remote
        class Actor:
            @ray.method(num_returns=3)
            def return_three(self, x):
                return x, x + 1, x + 2

            def echo(self, x):
                return x

            @ray.method(num_returns=2)
            def return_two(self, x):
                return x, x + 1

        a = Actor.remote()
        b = Actor.remote()
        with InputNode() as i:
            o1, o2, o3 = a.return_three.bind(i)
            o4 = b.echo.bind(o1)
            o5 = b.echo.bind(o2)
            o6, o7 = b.return_two.bind(o3)
            dag = MultiOutputNode([o4, o5, o6, o7])

        compiled_dag = dag.experimental_compile()

        ascii_visualization = compiled_dag.visualize(format="ascii")

        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
            ascii_visualization
        )

        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
        assert expected_nodes.issubset(
            node_names
        ), f"Expected nodes {expected_nodes} not found."

        expected_edges = {
            ("0", "1"),
            ("1", "2"),
            ("1", "3"),
            ("1", "4"),
            ("2", "5"),
            ("3", "6"),
            ("4", "7"),
            ("5", "8"),
            ("6", "8"),
            ("9", "8"),
            ("10", "8"),
            ("7", "9"),
            ("7", "10"),
        }
        assert expected_edges.issubset(
            edge_pairs
        ), f"Expected edges {expected_edges} not found."

    def test_visualize_ascii_cross_line(self, ray_start_regular):
        """
        Expect output:
            Nodes Information:
            0 [label="Task 0  InputNode"]
            1 [label="Task 1  Actor: 84835a... Method: return_three"]
            2 [label="Task 2  ClassMethodOutputNode[0]"]
            3 [label="Task 3  ClassMethodOutputNode[1]"]
            4 [label="Task 4  ClassMethodOutputNode[2]"]
            5 [label="Task 5  Actor: 02a6a1... Method: echo"]
            6 [label="Task 6  Actor: 02a6a1... Method: return_two"]
            7 [label="Task 7  Actor: 02a6a1... Method: echo"]
            8 [label="Task 8  MultiOutputNode"]
            9 [label="Task 9  ClassMethodOutputNode[0]"]
            10 [label="Task 10  ClassMethodOutputNode[1]"]

            Edges Information:
            0 ---> 1
            1 ---> 2
            1 ---> 3
            1 ---> 4
            2 ---> 5
            3 ---> 6
            4 ---> 7
            5 ---> 8
            7 ---> 8
            9 ---> 8
            10 ---> 8
            6 ---> 9
            6 ---> 10

            Legend:
            +++> : Represents Nccl-type data channels
            ---> : Represents Shared Memory data channels

            Graph Built:
            0:InputNode
            |
            1:Actor_84835a:return_three
            |---------------------------->|---------------------------->|                            # noqa
            2:Output[0]                   3:Output[1]                   4:Output[2]                  # noqa
            |                             |                             |                            # noqa
            5:Actor_02a6a1:echo           6:Actor_02a6a1:return_two     7:Actor_02a6a1:echo          # noqa
            |                             |---------------------------->|                            # noqa
            |                             9:Output[0]                   10:Output[1]                 # noqa
            |<----------------------------------------------------------|                            # noqa
            8:MultiOutputNod
        """

        @ray.remote
        class Actor:
            @ray.method(num_returns=3)
            def return_three(self, x):
                return x, x + 1, x + 2

            def echo(self, x):
                return x

            @ray.method(num_returns=2)
            def return_two(self, x):
                return x, x + 1

        a = Actor.remote()
        b = Actor.remote()
        with InputNode() as i:
            o1, o2, o3 = a.return_three.bind(i)
            o4 = b.echo.bind(o1)
            o5 = b.echo.bind(o3)
            o6, o7 = b.return_two.bind(o2)
            dag = MultiOutputNode([o4, o5, o6, o7])

        compiled_dag = dag.experimental_compile()

        ascii_visualization = compiled_dag.visualize(format="ascii")

        node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
            ascii_visualization
        )

        expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
        assert expected_nodes.issubset(
            node_names
        ), f"Expected nodes {expected_nodes} not found."

        expected_edges = {
            ("0", "1"),
            ("1", "2"),
            ("1", "3"),
            ("1", "4"),
            ("2", "5"),
            ("3", "6"),
            ("4", "7"),
            ("5", "8"),
            ("7", "8"),
            ("9", "8"),
            ("10", "8"),
            ("6", "9"),
            ("6", "10"),
        }
        assert expected_edges.issubset(
            edge_pairs
        ), f"Expected edges {expected_edges} not found."


if __name__ == "__main__":
    if os.environ.get("PARALLEL_CI"):
        sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
    else:
        sys.exit(pytest.main(["-sv", __file__]))