"""
Tests to ensure ray DAG can correctly mark its input(s) to take user
request, for all DAGNode types.
"""

from typing import Any, TypeVar

import pytest

import ray
from ray.dag.dag_node import DAGNode
from ray.dag.input_node import InputNode

RayHandleLike = TypeVar("RayHandleLike")


def test_no_args_to_input_node(shared_ray_instance):
    @ray.remote
    def f(input):
        return input

    with pytest.raises(
        ValueError, match="InputNode should not take any args or kwargs"
    ):
        with InputNode(0) as dag_input:
            f.bind(dag_input)
    with pytest.raises(
        ValueError,
        match="InputNode should not take any args or kwargs",
    ):
        with InputNode(key=1) as dag_input:
            f.bind(dag_input)


def test_simple_func(shared_ray_instance):
    @ray.remote
    def a(input: str):
        return f"{input} -> a"

    @ray.remote
    def b(a: "RayHandleLike"):
        # At runtime, a is replaced with execution result of a.
        return f"{a} -> b"

    # input -> a - > b -> ouput
    with InputNode() as dag_input:
        a_node = a.bind(dag_input)
        dag = b.bind(a_node)

    assert ray.get(dag.execute("input")) == "input -> a -> b"
    assert ray.get(dag.execute("test")) == "test -> a -> b"


def test_func_dag(shared_ray_instance):
    @ray.remote
    def a(user_input):
        return user_input

    @ray.remote
    def b(x):
        return x * 2

    @ray.remote
    def c(x):
        return x + 1

    @ray.remote
    def d(x, y):
        return x + y

    with InputNode() as dag_input:
        a_ref = a.bind(dag_input)
        b_ref = b.bind(a_ref)
        c_ref = c.bind(a_ref)
        d_ref = d.bind(b_ref, c_ref)
        d1_ref = d.bind(d_ref, d_ref)
        d2_ref = d.bind(d1_ref, d_ref)
        dag = d.bind(d2_ref, d_ref)

    # [(2*2 + 2+1) + (2*2 + 2+1)] + [(2*2 + 2+1) + (2*2 + 2+1)]
    assert ray.get(dag.execute(2)) == 28
    # [(3*2 + 3+1) + (3*2 + 3+1)] + [(3*2 + 3+1) + (3*2 + 3+1)]
    assert ray.get(dag.execute(3)) == 40


def test_multi_input_func_dag(shared_ray_instance):
    @ray.remote
    def a(user_input):
        return user_input * 2

    @ray.remote
    def b(user_input):
        return user_input + 1

    @ray.remote
    def c(x, y):
        return x + y

    with InputNode() as dag_input:
        a_ref = a.bind(dag_input)
        b_ref = b.bind(dag_input)
        dag = c.bind(a_ref, b_ref)

    # (2*2) + (2*1)
    assert ray.get(dag.execute(2)) == 7
    # (3*2) + (3*1)
    assert ray.get(dag.execute(3)) == 10


def test_invalid_input_node_as_class_constructor(shared_ray_instance):
    @ray.remote
    class Actor:
        def __init__(self, val):
            self.val = val

        def get(self):
            return self.val

    with pytest.raises(
        ValueError,
        match=(
            "InputNode handles user dynamic input the DAG, and "
            "cannot be used as args, kwargs, or other_args_to_resolve "
            "in ClassNode constructor because it is not available at "
            "class construction or binding time."
        ),
    ):
        with InputNode() as dag_input:
            Actor.bind(dag_input)


def test_class_method_input(shared_ray_instance):
    @ray.remote
    class Model:
        def __init__(self, weight: int):
            self.weight = weight

        def forward(self, input: "RayHandleLike"):
            return self.weight * input

    @ray.remote
    class FeatureProcessor:
        def __init__(self, scale):
            self.scale = scale

        def process(self, input: int):
            return input * self.scale

    with InputNode() as dag_input:
        preprocess = FeatureProcessor.bind(0.5)
        feature = preprocess.process.bind(dag_input)
        model = Model.bind(4)
        dag = model.forward.bind(feature)

    # 2 * 0.5 * 4
    assert ray.get(dag.execute(2)) == 4
    # 6 * 0.5 * 4
    assert ray.get(dag.execute(6)) == 12


def test_multi_class_method_input(shared_ray_instance):
    """
    Test a multiple class methods can all be used as inputs in a dag.
    """

    @ray.remote
    class Model:
        def __init__(self, weight: int):
            self.weight = weight

        def forward(self, input: int):
            return self.weight * input

    @ray.remote
    def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
        return m1 + m2

    with InputNode() as dag_input:
        m1 = Model.bind(2)
        m2 = Model.bind(3)

        m1_output = m1.forward.bind(dag_input)
        m2_output = m2.forward.bind(dag_input)

        dag = combine.bind(m1_output, m2_output)

    # 1*2 + 1*3
    assert ray.get(dag.execute(1)) == 5
    # 2*2 + 2*3
    assert ray.get(dag.execute(2)) == 10


def test_func_class_mixed_input(shared_ray_instance):
    """
    Test both class method and function are used as input in the
    same dag.
    """

    @ray.remote
    class Model:
        def __init__(self, weight: int):
            self.weight = weight

        def forward(self, input: int):
            return self.weight * input

    @ray.remote
    def model_func(input: int):
        return input * 2

    @ray.remote
    def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
        return m1 + m2

    with InputNode() as dag_input:
        m1 = Model.bind(3)
        m1_output = m1.forward.bind(dag_input)
        m2_output = model_func.bind(dag_input)

    dag = combine.bind(m1_output, m2_output)
    # 2*3 + 2*2
    assert ray.get(dag.execute(2)) == 10
    # 3*3 + 3*2
    assert ray.get(dag.execute(3)) == 15


def test_input_attr_partial_access(shared_ray_instance):
    @ray.remote
    class Model:
        def __init__(self, weight: int):
            self.weight = weight

        def forward(self, input: int):
            return self.weight * input

    @ray.remote
    def combine(a, b, c, d=None):
        if not d:
            return a + b + c
        else:
            return a + b + c + d["deep"]["nested"]

    # 1) Test default wrapping of args and kwargs into internal python object
    with InputNode() as dag_input:
        m1 = Model.bind(1)
        m2 = Model.bind(2)
        m1_output = m1.forward.bind(dag_input[0])
        m2_output = m2.forward.bind(dag_input[1])
        dag = combine.bind(m1_output, m2_output, dag_input.m3, dag_input.m4)
    # 1*1 + 2*2 + 3 + 4 = 12
    assert ray.get(dag.execute(1, 2, m3=3, m4={"deep": {"nested": 4}})) == 12

    # 2) Test user passed data object as only input to the dag.execute()
    class UserDataObj:
        user_object_field_0: Any
        user_object_field_1: Any
        field_3: Any

        def __init__(
            self, user_object_field_0: Any, user_object_field_1: Any, field_3: Any
        ) -> None:
            self.user_object_field_0 = user_object_field_0
            self.user_object_field_1 = user_object_field_1
            self.field_3 = field_3

    with InputNode() as dag_input:
        m1 = Model.bind(1)
        m2 = Model.bind(2)
        m1_output = m1.forward.bind(dag_input.user_object_field_0)
        m2_output = m2.forward.bind(dag_input.user_object_field_1)
        dag = combine.bind(m1_output, m2_output, dag_input.field_3)

    # 1*1 + 2*2 + 3
    assert ray.get(dag.execute(UserDataObj(1, 2, 3))) == 8

    # 3) Test user passed only one list object with regular list index accessor
    with InputNode() as dag_input:
        m1 = Model.bind(1)
        m2 = Model.bind(2)
        m1_output = m1.forward.bind(dag_input[0])
        m2_output = m2.forward.bind(dag_input[1])
        dag = combine.bind(m1_output, m2_output, dag_input[2])
    # 1*1 + 2*2 + 3 + 4 = 12
    assert ray.get(dag.execute([1, 2, 3])) == 8

    # 4) Test user passed only one dict object with key str accessor
    with InputNode() as dag_input:
        m1 = Model.bind(1)
        m2 = Model.bind(2)
        m1_output = m1.forward.bind(dag_input["m1"])
        m2_output = m2.forward.bind(dag_input["m2"])
        dag = combine.bind(m1_output, m2_output, dag_input["m3"])
    # 1*1 + 2*2 + 3 + 4 = 12
    assert ray.get(dag.execute({"m1": 1, "m2": 2, "m3": 3})) == 8

    with pytest.raises(
        AssertionError,
        match="Please only use int index or str as first-level key",
    ):
        with InputNode() as dag_input:
            m1 = Model.bind(1)
            dag = m1.forward.bind(dag_input[(1, 2)])


def test_ensure_in_context_manager(shared_ray_instance):
    # No enforcement on creation given __enter__ executes after __init__
    input = InputNode()
    with pytest.raises(
        AssertionError,
        match=(
            "InputNode is a singleton instance that should be only used "
            "in context manager"
        ),
    ):
        input.execute()

    @ray.remote
    def f(input):
        return input

    # No enforcement on creation given __enter__ executes after __init__
    dag = f.bind(InputNode())
    with pytest.raises(
        AssertionError,
        match=(
            "InputNode is a singleton instance that should be only used "
            "in context manager"
        ),
    ):
        dag.execute()


def test_ensure_input_node_singleton(shared_ray_instance):
    @ray.remote
    def f(input):
        return input

    @ray.remote
    def combine(a, b):
        return a + b

    with InputNode() as input_1:
        a = f.bind(input_1)
    with InputNode() as input_2:
        b = f.bind(input_2)
        dag = combine.bind(a, b)

    with pytest.raises(
        AssertionError, match="Each DAG should only have one unique InputNode"
    ):
        _ = ray.get(dag.execute(2))


def test_apply_recursive_caching(shared_ray_instance):
    @ray.remote
    def f(input):
        return input

    input = InputNode()
    f_node = f.bind(input)

    a, b = input, f_node
    for _ in range(10):
        a, b = f.bind(a, b), f.bind(a, b)

    counter = 0
    original_apply_recursive = DAGNode.apply_recursive

    def _apply_recursive_with_counter(self, fn):
        nonlocal counter
        counter += 1
        return original_apply_recursive(self, fn)

    DAGNode.apply_recursive = _apply_recursive_with_counter

    a.apply_recursive(lambda node: node)

    # Prior to #40337; count was 2559
    assert counter == 40
    DAGNode.apply_recursive = original_apply_recursive


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))