import threading
from typing import Callable, Dict

import pytest

import ray
from ray.data._internal.execution.block_ref_counter import BlockRefCounter


def _ref(uid: int) -> "ray.ObjectRef":
    """Real ObjectRef with a deterministic distinct 28-byte ID, no Ray cluster needed."""
    return ray.ObjectRef(uid.to_bytes(28, "big"))


class FakeAddObjectOutOfScopeCallback:
    """Test double for CoreWorker.add_object_out_of_scope_callback.

    Records each registered callback keyed by the block's object-ID bytes so a
    test can fire it explicitly. Set should_registration_fail=True to simulate an object
    that is already out of scope at registration time.
    """

    def __init__(self, should_registration_fail: bool = False):
        self._should_fail = should_registration_fail
        self._callbacks: Dict[bytes, Callable[[bytes], None]] = {}

    def __call__(
        self, object_ref: "ray.ObjectRef", callback: Callable[[bytes], None]
    ) -> bool:
        if not self._should_fail:
            self._callbacks[object_ref.binary()] = callback
        return not self._should_fail

    def fire(self, object_ref: "ray.ObjectRef") -> None:
        id_binary = object_ref.binary()
        self._callbacks[id_binary](id_binary)


class TestBlockRefCounterAccounting:
    def test_single_block_produced_and_released(self):
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        ref = _ref(1)

        counter.on_block_produced(ref, 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 1

        add_cb.fire(ref)
        assert counter.get_object_store_memory_usage("op_a") == 0

    def test_multiple_blocks_same_producer(self):
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        ref1, ref2 = _ref(1), _ref(2)

        counter.on_block_produced(ref1, 1, "op_a")
        counter.on_block_produced(ref2, 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 2

        add_cb.fire(ref1)
        assert counter.get_object_store_memory_usage("op_a") == 1
        add_cb.fire(ref2)
        assert counter.get_object_store_memory_usage("op_a") == 0

    def test_multiple_producers_isolated(self):
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        ref1, ref2 = _ref(1), _ref(2)

        counter.on_block_produced(ref1, 1, "op_a")
        counter.on_block_produced(ref2, 1, "op_b")
        assert counter.get_object_store_memory_usage("op_a") == 1
        assert counter.get_object_store_memory_usage("op_b") == 1

        add_cb.fire(ref1)
        assert counter.get_object_store_memory_usage("op_a") == 0
        assert counter.get_object_store_memory_usage("op_b") == 1

    def test_duplicate_registration_is_noop(self):
        """on_block_produced is idempotent: a duplicate ref is silently ignored.

        This matters when an AllToAllOperator forwards an input ref unchanged;
        the ref was already registered by the upstream producer."""
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        ref = _ref(1)

        counter.on_block_produced(ref, 1, "op_a")
        counter.on_block_produced(ref, 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 1

        add_cb.fire(ref)
        assert counter.get_object_store_memory_usage("op_a") == 0

    def test_register_when_object_already_out_of_scope(self):
        """If registration reports the object is already gone, the increment is
        undone immediately so the producer nets to zero."""
        add_cb = FakeAddObjectOutOfScopeCallback(should_registration_fail=True)
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)

        counter.on_block_produced(_ref(1), 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 0


class TestBlockRefCounterClear:
    def test_clear_resets_usage(self):
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        counter.on_block_produced(_ref(1), 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 1

        counter.clear()
        assert counter.get_object_store_memory_usage("op_a") == 0

    def test_stale_callback_after_clear_is_noop(self):
        """A stale callback firing after clear() must not touch accounting
        recorded after the reset."""
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        stale_ref = _ref(1)
        counter.on_block_produced(stale_ref, 1, "op_a")

        counter.clear()

        counter.on_block_produced(_ref(2), 1, "op_a")
        assert counter.get_object_store_memory_usage("op_a") == 1

        add_cb.fire(stale_ref)
        assert counter.get_object_store_memory_usage("op_a") == 1


class TestBlockRefCounterThreadSafety:
    def test_concurrent_callbacks_dont_corrupt_state(self):
        """Many threads firing callbacks at once must not corrupt the count."""
        add_cb = FakeAddObjectOutOfScopeCallback()
        counter = BlockRefCounter(add_object_out_of_scope_callback=add_cb)
        producer_id = "op_concurrent"
        n = 50
        refs = [_ref(i) for i in range(n)]
        for ref in refs:
            counter.on_block_produced(ref, 1, producer_id)
        assert counter.get_object_store_memory_usage(producer_id) == n

        threads = [threading.Thread(target=add_cb.fire, args=(ref,)) for ref in refs]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert counter.get_object_store_memory_usage(producer_id) == 0


if __name__ == "__main__":
    import sys

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