import pyarrow as pa
import pytest

import ray
from ray.data._internal.execution.interfaces import ExecutionOptions
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (  # noqa: E501
    ShuffleMapOp,
    make_partition_sentinel,
)
from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (  # noqa: E501
    ShuffleReduceOp,
)
from ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks import (
    _encode_partition_ipc,
    _get_shard_batch,
    _ipc_write_options,
)
from ray.data._internal.execution.util import make_ref_bundles
from ray.data.block import BlockMetadata
from ray.data.context import DataContext, ShuffleStrategy
from ray.data.tests.conftest import *  # noqa: F401, F403
from ray.data.tests.conftest import noop_counter
from ray.data.tests.util import run_op_tasks_sync
from ray.exceptions import GetTimeoutError
from ray.tests.conftest import *  # noqa: F401, F403


def _keys_per_block(ds, columns):
    """Return, for each output block, the set of distinct key tuples it holds.

    Used to assert the hash-shuffle co-location guarantee: a key must appear in
    exactly one block.
    """
    per_block = []
    for ref_bundle in ds.iter_internal_ref_bundles():
        for block_ref in ref_bundle.block_refs:
            block = ray.get(block_ref)
            cols = [block[c].to_pylist() for c in columns]
            per_block.append(set(zip(*cols)))
    return per_block


def _assert_keys_colocated(per_block):
    """Every key tuple appears in at most one block."""
    all_keys = [k for block in per_block for k in block]
    assert len(all_keys) == len(
        set(all_keys)
    ), f"A key landed in more than one block: {per_block}"


@pytest.mark.parametrize("num_partitions", [1, 4, 8])
def test_repartition_keys_preserves_rows(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
    num_partitions,
):
    """No rows are lost or duplicated; key totals are preserved."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    ds = ray.data.range(1000, override_num_blocks=10)
    out = ds.repartition(num_partitions, keys=["id"])
    assert out.count() == 1000
    assert out.sum("id") == sum(range(1000))


def test_repartition_block_number_matched(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """All-non-empty partitions => exactly num_partitions output blocks."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    # 1000 distinct keys over 8 buckets => all 8 partitions are non-empty.
    ds = ray.data.range(1000, override_num_blocks=20)
    out = ds.repartition(8, keys=["id"]).materialize()
    assert out.num_blocks() == 8


def test_same_key_lands_in_same_block(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """All rows sharing a key should end up in one block."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    ds = ray.data.range(500, override_num_blocks=10).map(
        lambda row: {"k": row["id"] % 25, "v": row["id"]}
    )
    out = ds.repartition(5, keys=["k"])

    _assert_keys_colocated(_keys_per_block(out, ["k"]))
    assert out.count() == 500


def test_multi_column_keys(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """Composite keys hash on all columns: every distinct (a, b) tuple lands in
    exactly one block."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    ds = ray.data.range(500, override_num_blocks=10).map(
        lambda row: {"a": row["id"] % 5, "b": row["id"] % 7, "v": row["id"]}
    )
    out = ds.repartition(4, keys=["a", "b"])

    _assert_keys_colocated(_keys_per_block(out, ["a", "b"]))
    assert out.count() == 500


def test_more_partitions_than_keys_emits_empty_blocks(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """Requesting more partitions than there are distinct keys emits the extra
    partitions as empty (0-row) blocks that still carry the dataset schema."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    # 3 distinct keys into 50 partitions => at most 3 non-empty, >=47 empty.
    ds = ray.data.range(600, override_num_blocks=10).map(
        lambda row: {"k": row["id"] % 3, "v": row["id"]}
    )
    out = ds.repartition(50, keys=["k"]).materialize()

    assert out.count() == 600
    assert out.num_blocks() == 50

    rows_per_block = []
    schemas = []
    for ref_bundle in out.iter_internal_ref_bundles():
        for block_ref in ref_bundle.block_refs:
            block = ray.get(block_ref)
            rows_per_block.append(block.num_rows)
            schemas.append(block.schema)

    assert rows_per_block.count(0) >= 47
    assert all(schema.equals(schemas[0]) for schema in schemas)

    _assert_keys_colocated(_keys_per_block(out, ["k"]))


def test_repartition_empty_dataset(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """Empty dataset should still output N blocks"""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    ds = ray.data.range(100, override_num_blocks=4).filter(lambda row: False)
    out = ds.repartition(4, keys=["id"]).materialize()
    assert out.count() == 0
    assert out.num_blocks() == 4
    rows_per_block = [
        ray.get(block_ref).num_rows
        for ref_bundle in out.iter_internal_ref_bundles()
        for block_ref in ref_bundle.block_refs
    ]
    assert rows_per_block == [0, 0, 0, 0]


def test_repartition_with_sort_produces_sorted_partitions(
    ray_start_regular_shared_2_cpus,
    restore_data_context,
    disable_fallback_to_object_extension,
):
    """Check that rows are sorted in every partition."""
    ctx = DataContext.get_current()
    ctx.shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE

    ds = ray.data.range(200, override_num_blocks=4)
    out = ds.repartition(4, keys=["id"], sort=True)

    for ref_bundle in out.iter_internal_ref_bundles():
        for block_ref in ref_bundle.block_refs:
            ids = ray.get(block_ref)["id"].to_pylist()
            assert ids == sorted(ids)


def test_get_shard_batch_no_timeout(ray_start_regular_shared_2_cpus):
    """timeout_s <= 0 fetches in a single blocking ray.get."""
    refs = [ray.put(i) for i in range(4)]
    out = _get_shard_batch(
        refs,
        partition_id=0,
        batch_index=0,
        num_batches=1,
        timeout_s=0,
    )
    assert out == [0, 1, 2, 3]


def test_get_shard_batch_returns_ready_values(ray_start_regular_shared_2_cpus):
    """A timeout that is never hit returns the values unchanged."""
    refs = [ray.put(i) for i in range(3)]
    out = _get_shard_batch(
        refs,
        partition_id=1,
        batch_index=0,
        num_batches=1,
        timeout_s=30.0,
    )
    assert out == [0, 1, 2]


def test_get_shard_batch_warns_then_raises_on_stall(
    ray_start_regular_shared_2_cpus, propagate_logs, caplog
):
    """A stalled fetch warns partway through, then raises at the timeout."""

    @ray.remote
    def _never_ready():
        import time

        time.sleep(1000)

    ref = _never_ready.remote()
    with caplog.at_level(
        "WARNING",
        logger="ray.data._internal.execution.operators.shuffle_operators.shuffle_tasks",
    ):
        with pytest.raises(GetTimeoutError):
            _get_shard_batch(
                [ref],
                partition_id=7,
                batch_index=0,
                num_batches=1,
                timeout_s=0.3,
            )
    assert [r.levelname for r in caplog.records].count("WARNING") == 1
    assert [r.levelname for r in caplog.records].count("ERROR") == 1
    assert "partition 7" in caplog.records[0].message
    ray.cancel(ref, force=True)


# --- Multi-input reduce -------------------------------------------------------
# TODO: move these multi-input ShuffleReduceOp tests (and the _get_shard_batch
# shuffle_tasks tests above) into a dedicated operator/task-level test file --
# they aren't specific to hash-shuffle-v2.
def _ipc_shard_bundle(partition_id, table):
    """One partition's shard as a ShuffleMapOp emits it: an IPC-encoded buffer
    stamped with the partition id."""
    from ray.data._internal.execution.interfaces import BlockEntry, RefBundle

    buf = _encode_partition_ipc(table, _ipc_write_options("none"))
    meta = BlockMetadata(
        num_rows=table.num_rows,
        size_bytes=table.nbytes,
        exec_stats=None,
        input_files=make_partition_sentinel(partition_id),
    )
    return RefBundle(
        (
            BlockEntry(
                ref=ray.put(buf),  # pyrefly: ignore[bad-argument-type]
                metadata=meta,
            ),
        ),
        schema=table.schema,
        owns_blocks=True,
    )


def _make_multi_input_reduce_op(reduce_fn, num_inputs=2, num_partitions=2):
    ctx = DataContext.get_current()
    maps = [
        ShuffleMapOp(
            InputDataBuffer(ctx, make_ref_bundles([[0]])),
            ctx,
            num_partitions=num_partitions,
            partition_fn=lambda t: {},
        )
        for _ in range(num_inputs)
    ]
    return ShuffleReduceOp(
        maps,
        ctx,
        num_partitions=num_partitions,
        reduce_fn=reduce_fn,
        disallow_block_splitting=True,
    )


def _drain_reduce_op(op, feed):
    """Run `op` over `feed` (bundle, input_index) pairs and return output tables."""
    op.start(ExecutionOptions(), noop_counter())
    for bundle, input_index in feed:
        op.add_input(bundle, input_index)
    op.all_inputs_done()
    run_op_tasks_sync(op)
    tables = []
    while op.has_next():
        for ref in op.get_next().block_refs:
            tables.append(ray.get(ref))
    return tables


def _concat_inputs_reduce_fn():
    def _reduce(partition_id, tables_by_input):
        tables = [t for shards in tables_by_input for t in shards]
        if tables:
            yield pa.concat_tables(tables)

    return _reduce


def test_reduce_op_combines_all_inputs(ray_start_regular_shared_2_cpus):
    """Both inputs' shards for a partition reach the reducer, in input order."""
    op = _make_multi_input_reduce_op(_concat_inputs_reduce_fn(), num_inputs=2)
    feed = [
        (_ipc_shard_bundle(0, pa.table({"src": ["L"], "v": [1]})), 0),
        (_ipc_shard_bundle(0, pa.table({"src": ["R"], "v": [2]})), 1),
    ]
    out = pa.concat_tables(_drain_reduce_op(op, feed))
    assert sorted(out.column("src").to_pylist()) == ["L", "R"]
    assert sorted(out.column("v").to_pylist()) == [1, 2]


def test_reduce_op_runs_when_an_input_is_missing(ray_start_regular_shared_2_cpus):
    """A partition that never receives one input (a block-less side) is still
    reduced -- the reducer sees an empty shard list for the missing input rather
    than the op hanging on a never-paired partition."""
    op = _make_multi_input_reduce_op(_concat_inputs_reduce_fn(), num_inputs=2)
    # Only input 0 delivers partition 0; input 1 never does.
    feed = [(_ipc_shard_bundle(0, pa.table({"src": ["L"], "v": [1]})), 0)]
    out = pa.concat_tables(_drain_reduce_op(op, feed))
    assert out.column("src").to_pylist() == ["L"]
    assert op.has_completed()


if __name__ == "__main__":
    import sys

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