import pytest

import ray
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
from ray.data._internal.execution.interfaces import (
    BlockEntry,
    PhysicalOperator,
    RefBundle,
)
from ray.data._internal.execution.interfaces.execution_options import (
    ExecutionOptions,
    ExecutionResources,
)
from ray.data._internal.execution.operators.union_operator import UnionOperator
from ray.data._internal.execution.resource_manager import (
    ResourceManager,
)
from ray.data._internal.execution.streaming_executor_state import (
    build_streaming_topology,
)
from ray.data.block import BlockMetadata
from ray.data.context import DataContext
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.conftest import noop_counter


def test_physical_operator_tracks_output_dependencies():
    input_op = PhysicalOperator("input", [], DataContext.get_current())
    downstream_op = PhysicalOperator(
        "downstream", [input_op], DataContext.get_current()
    )

    assert input_op.output_dependencies == [downstream_op]


def test_physical_apply_transform_rewires_all_input_output_dependencies():
    ctx = DataContext.get_current()
    left_input = PhysicalOperator("left_input", [], ctx)
    right_input = PhysicalOperator("right_input", [], ctx)
    root = PhysicalOperator("root", [left_input, right_input], ctx)
    left_replacement = PhysicalOperator("left_replacement", [], ctx)

    transformed_root = root._apply_transform(
        lambda op: left_replacement if op is left_input else op
    )

    assert transformed_root is not root
    assert transformed_root.id != root.id
    assert transformed_root.metrics is not root.metrics
    assert transformed_root.input_dependencies == [left_replacement, right_input]
    assert transformed_root in left_replacement.output_dependencies
    assert transformed_root in right_input.output_dependencies
    assert root not in left_input.output_dependencies
    assert root not in right_input.output_dependencies


def test_physical_apply_transform_rewires_when_current_node_is_replaced():
    ctx = DataContext.get_current()
    left_input = PhysicalOperator("left_input", [], ctx)
    right_input = PhysicalOperator("right_input", [], ctx)
    root = PhysicalOperator("root", [left_input, right_input], ctx)

    transformed_root = root._apply_transform(
        lambda op: PhysicalOperator("replacement", [left_input], ctx)
        if op is root
        else op
    )

    assert transformed_root is not root
    assert transformed_root in left_input.output_dependencies
    assert root not in left_input.output_dependencies
    assert root not in right_input.output_dependencies
    assert transformed_root not in right_input.output_dependencies


def test_physical_apply_transform_deep_chain_no_stale_downstream_refs():
    ctx = DataContext.get_current()
    leaf = PhysicalOperator("leaf", [], ctx)
    mid = PhysicalOperator("mid", [leaf], ctx)
    root = PhysicalOperator("root", [mid], ctx)

    def transform(op: PhysicalOperator) -> PhysicalOperator:
        if op is leaf:
            return PhysicalOperator("leaf_replacement", [], ctx)
        if op.name == "root":
            return PhysicalOperator("root_replacement", op.input_dependencies, ctx)
        return op

    transformed_root = root._apply_transform(transform)
    transformed_mid = transformed_root.input_dependencies[0]
    transformed_leaf = transformed_mid.input_dependencies[0]

    assert transformed_root.name == "root_replacement"
    assert transformed_mid is not mid
    assert transformed_leaf.name == "leaf_replacement"
    assert root not in transformed_mid.output_dependencies
    assert transformed_mid.output_dependencies == [transformed_root]


def test_physical_apply_transform_rejects_in_place_input_mutation():
    ctx = DataContext.get_current()
    old_input = PhysicalOperator("old_input", [], ctx)
    new_input = PhysicalOperator("new_input", [], ctx)
    root = PhysicalOperator("root", [old_input], ctx)

    def transform(op: PhysicalOperator) -> PhysicalOperator:
        if op is root:
            op._input_dependencies = [new_input]
            return op
        return op

    with pytest.raises(
        AssertionError,
        match="In-place input mutation is not supported; return a new node instead.",
    ):
        root._apply_transform(transform)


def test_does_not_double_count_usage_from_union():
    """Regression test for https://github.com/ray-project/ray/pull/61040."""
    # Create a mock topology:
    #
    #   input1 ───┐
    #             ├─▶ union_op
    #   input2 ───┘
    input1 = PhysicalOperator("op1", [], DataContext.get_current())
    input2 = PhysicalOperator("op2", [], DataContext.get_current())
    union_op = UnionOperator(DataContext.get_current(), input1, input2)
    topology = build_streaming_topology(union_op, ExecutionOptions(), noop_counter())

    # Create a resource manager.
    total_resources = ExecutionResources(cpu=0, object_store_memory=2)
    resource_manager = ResourceManager(
        topology,
        ExecutionOptions(),
        lambda: total_resources,
        DataContext.get_current(),
        BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
    )

    # Create two 1-byte `RefBundle`s.
    block_ref1 = ray.ObjectRef(b"1" * 28)
    block_ref2 = ray.ObjectRef(b"2" * 28)
    block_metadata = BlockMetadata(
        num_rows=1, size_bytes=1, input_files=None, exec_stats=None
    )
    bundle1 = RefBundle(
        [BlockEntry(block_ref1, block_metadata)], owns_blocks=True, schema=None
    )
    bundle2 = RefBundle(
        [BlockEntry(block_ref2, block_metadata)], owns_blocks=True, schema=None
    )

    # Add two 1-byte `RefBundle` to the union operator.
    topology[union_op].add_output(bundle1)
    topology[union_op].add_output(bundle2)
    resource_manager.update_usages()

    # The total object store memory usage should be 2. If the resource manager double-
    # counts the usage from the union operator, the total object store memory usage can
    # be greater than 2.
    total_object_store_memory = sum(
        [
            resource_manager.get_op_usage(
                op, include_ineligible_downstream=True
            ).object_store_memory
            for op in topology.keys()
        ]
    )
    assert total_object_store_memory == 2, total_object_store_memory


def test_per_input_inqueue_attribution_for_union():
    """Test that per-input attribution correctly charges each upstream operator
    only for the blocks it produced in the union's internal input queue.

    When preserve_order=True, the union operator buffers blocks per-input.
    The resource manager should attribute each input buffer's memory only to
    the corresponding upstream operator, not to all upstream operators.
    """
    # Create a mock topology:
    #
    #   input1 ───┐
    #             ├─▶ union_op
    #   input2 ───┘
    input1 = PhysicalOperator("op1", [], DataContext.get_current())
    input2 = PhysicalOperator("op2", [], DataContext.get_current())
    union_op = UnionOperator(DataContext.get_current(), input1, input2)

    options = ExecutionOptions()
    options.preserve_order = True
    topology = build_streaming_topology(union_op, options, noop_counter())

    # Create a resource manager.
    total_resources = ExecutionResources(cpu=0, object_store_memory=200)
    resource_manager = ResourceManager(
        topology,
        options,
        lambda: total_resources,
        DataContext.get_current(),
        BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True),
    )

    # Create two 10-byte RefBundles with distinct block refs (simulates real execution
    # where each block from a source has its own ObjectRef).
    block_ref1 = ray.ObjectRef(b"1" * 28)
    block_ref2 = ray.ObjectRef(b"2" * 28)
    block_metadata = BlockMetadata(
        num_rows=1, size_bytes=10, input_files=None, exec_stats=None
    )
    bundle1 = RefBundle(
        [BlockEntry(block_ref1, block_metadata)], owns_blocks=True, schema=None
    )
    bundle2 = RefBundle(
        [BlockEntry(block_ref2, block_metadata)], owns_blocks=True, schema=None
    )

    # Add blocks only to input2's buffer inside the union operator.
    # With preserve_order=True, _add_input_inner routes to _input_buffers[input_index].
    union_op.add_input(bundle1, input_index=1)
    union_op.add_input(bundle2, input_index=1)

    resource_manager.update_usages()

    # input2 should be charged for its blocks in the union's input buffer (20 bytes).
    input2_usage = resource_manager.get_op_usage(
        input2, include_ineligible_downstream=True
    ).object_store_memory
    # input1 should NOT be charged for input2's blocks (0 bytes from union inqueue).
    input1_usage = resource_manager.get_op_usage(
        input1, include_ineligible_downstream=True
    ).object_store_memory

    assert input1_usage == 0
    assert input2_usage == 20


if __name__ == "__main__":
    import sys

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