import time
from typing import TYPE_CHECKING, List, Optional

import pandas as pd
import pytest

import ray
from ray.data import ActorPoolStrategy, TaskPoolStrategy
from ray.data._internal.datasource.range_datasource import RangeDatasource
from ray.data._internal.logical.operators.read_operator import Read
from ray.data._internal.util import rows_same
from ray.data.block import Block, BlockMetadata
from ray.data.context import DataContext
from ray.data.datasource import Datasource, ReadTask
from ray.data.tests.conftest import *  # noqa
from ray.tests.conftest import *  # noqa

if TYPE_CHECKING:
    from ray.actor import ActorHandle
    from ray.data._internal.logical.interfaces import LogicalOperator


def find_read_op(op: "LogicalOperator") -> Optional[Read]:
    """Find the Read operator in the logical plan."""
    if isinstance(op, Read):
        return op
    if hasattr(op, "input_dependencies"):
        for input_op in op.input_dependencies:
            result = find_read_op(input_op)
            if result:
                return result
    return None


class TestDatasource(Datasource):
    """Unified datasource that captures actor_id and optionally tracks concurrency."""

    def __init__(self, n: int, concurrency_counter: Optional["ActorHandle"] = None):
        super().__init__()
        self._n = int(n)
        self._concurrency_counter = concurrency_counter

    def get_name(self) -> str:
        return "TestDatasource"

    def estimate_inmemory_data_size(self) -> Optional[int]:
        """Return an estimate of the in-memory data size."""
        # 2 columns (value, actor_id), 8 bytes per value
        return 8 * self._n * 2

    def get_read_tasks(
        self,
        parallelism: int,
        per_task_row_limit: Optional[int] = None,
        data_context: Optional["DataContext"] = None,
    ) -> List[ReadTask]:
        if self._n == 0:
            return []

        read_tasks: List[ReadTask] = []
        n = self._n
        block_size = max(1, n // parallelism)
        counter_ref = self._concurrency_counter

        def make_block(start: int, count: int, counter) -> Block:
            import pyarrow as pa

            # Track concurrency if counter is provided
            if counter is not None:
                ray.get(counter.inc.remote())  # type: ignore
            try:
                # Simulate some work when tracking concurrency
                if counter is not None:
                    time.sleep(0.1)

                # Capture actor_id during execution
                runtime_context = ray.get_runtime_context()
                actor_id = runtime_context.get_actor_id()

                return pa.Table.from_arrays(
                    [
                        pa.array(range(start, start + count)),
                        pa.array([actor_id] * count),
                    ],
                    names=["value", "actor_id"],
                )
            finally:
                # Decrement concurrency counter if provided
                if counter is not None:
                    ray.get(counter.decr.remote())  # type: ignore

        i = 0
        while i < n:
            count = min(block_size, n - i)
            meta = BlockMetadata(
                num_rows=count,
                size_bytes=8 * count * 2,  # Rough estimate: 2 columns
                input_files=None,
                exec_stats=None,
            )

            def read_fn(start=i, count=count, counter=counter_ref):
                yield make_block(start, count, counter)

            read_tasks.append(
                ReadTask(
                    read_fn,
                    meta,
                    schema=None,
                    per_task_row_limit=per_task_row_limit,
                )
            )
            i += block_size

        return read_tasks


@pytest.mark.parametrize(
    "concurrency,compute,expected_strategy_type",
    [
        (None, None, TaskPoolStrategy),
        (1, None, TaskPoolStrategy),
        (2, None, TaskPoolStrategy),
        (None, TaskPoolStrategy(), TaskPoolStrategy),
        (None, TaskPoolStrategy(size=4), TaskPoolStrategy),
        (None, ActorPoolStrategy(size=2), ActorPoolStrategy),
        (
            1,
            ActorPoolStrategy(size=4),
            TaskPoolStrategy,
        ),  # concurrency takes precedence
    ],
)
def test_read_datasource_compute_strategy(
    ray_start_regular_shared_2_cpus,
    concurrency,
    compute,
    expected_strategy_type,
    target_max_block_size_infinite_or_default,
):
    """Test that compute strategy is correctly set based on concurrency and compute parameters."""
    datasource = RangeDatasource(n=100)

    ds = ray.data.read_datasource(
        datasource,
        concurrency=concurrency,
        compute=compute,
        override_num_blocks=4,
    )

    # Get the logical plan to inspect the compute strategy on the logical operator
    logical_plan = ds._logical_plan
    read_op = find_read_op(logical_plan.dag)

    # Verify the compute strategy type on the logical operator
    assert read_op is not None, "Could not find Read operator in logical plan"
    assert isinstance(
        read_op.compute, expected_strategy_type
    ), f"Expected {expected_strategy_type}, got {type(read_op.compute)}"

    # If concurrency was specified, verify it takes precedence
    if concurrency is not None:
        assert isinstance(read_op.compute, TaskPoolStrategy)
        assert read_op.compute.size == concurrency


@pytest.mark.parametrize(
    "compute,expect_actor_execution",
    [
        (TaskPoolStrategy(), False),
        (ActorPoolStrategy(size=2), True),
        (ActorPoolStrategy(min_size=1, max_size=4), True),
    ],
)
def test_read_datasource_actor_execution(
    ray_start_regular_shared_2_cpus,
    compute,
    expect_actor_execution,
    target_max_block_size_infinite_or_default,
):
    """Test that ReadTasks execute in actors when using ActorPoolStrategy."""
    datasource = TestDatasource(n=100)

    ds = ray.data.read_datasource(
        datasource,
        compute=compute,
        override_num_blocks=4,
    )

    # Materialize to trigger execution
    result = ds.take_all()

    # Extract actor_ids from the data (they're included in each row)
    actor_ids_in_data = {row.get("actor_id") for row in result}

    if expect_actor_execution:
        # Should have actor_ids in the data (not None)
        assert len(actor_ids_in_data) > 0, "Expected actor_ids in data"
        # All actor_ids should be non-None
        assert (
            None not in actor_ids_in_data
        ), "Expected all actor_ids to be non-None for ActorPoolStrategy"
        # With ActorPoolStrategy(size=2), we should have at most 2 unique actor_ids
        # When size is specified, min_size == max_size == size
        if (
            isinstance(compute, ActorPoolStrategy)
            and compute.min_size == compute.max_size
        ):
            assert (
                len(actor_ids_in_data) <= compute.min_size
            ), f"Expected at most {compute.min_size} unique actor_ids, got {len(actor_ids_in_data)}"
    else:
        # Should not have actor_ids (all None)
        assert (
            actor_ids_in_data == {None} or len(actor_ids_in_data) == 0
        ), f"Expected all actor_ids to be None for TaskPoolStrategy, got {actor_ids_in_data}"


@pytest.mark.parametrize(
    "compute_strategy",
    [
        TaskPoolStrategy(),
        TaskPoolStrategy(size=2),
        ActorPoolStrategy(size=2),
        ActorPoolStrategy(min_size=1, max_size=4),
    ],
)
def test_read_datasource_basic_functionality(
    ray_start_regular_shared_2_cpus,
    compute_strategy,
):
    """Test that read_datasource works correctly with different compute strategies."""
    datasource = RangeDatasource(n=100)

    ds = ray.data.read_datasource(
        datasource,
        compute=compute_strategy,
        override_num_blocks=4,
    )

    df = ds.to_pandas()

    expected_df = pd.DataFrame({"value": list(range(100))})
    assert rows_same(df, expected_df)


if __name__ == "__main__":
    import sys

    import pytest

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