import threading
from collections import defaultdict
from typing import Callable, Dict, Optional

import ray
from ray._private.worker import global_worker


class BlockRefCounter:
    """Tracks object-store memory usage per operator via Ray Core callbacks.

    The callback fires when:
    - All Python ObjectRefs wrapping the block's ObjectID are garbage-collected, AND
    - All Ray tasks that received the block as an argument have completed.
    """

    def __init__(
        self,
        add_object_out_of_scope_callback: Optional[
            Callable[["ray.ObjectRef", Callable[[bytes], None]], bool]
        ] = None,
    ):
        if add_object_out_of_scope_callback is None:
            add_object_out_of_scope_callback = (
                global_worker.core_worker.add_object_out_of_scope_callback  # pyrefly: ignore[missing-attribute]
            )
        self._add_callback_fn = add_object_out_of_scope_callback
        # IDs of live blocks. Stale callbacks (fired after clear()) check
        # membership here and no-op, preventing negative _bytes_by_producer.
        self._registered_ids: set[bytes] = set()
        # (producer_id -> total live bytes); maintained incrementally for O(1) reads.
        self._bytes_by_producer: Dict[str, int] = defaultdict(int)
        self._lock = threading.Lock()

    def on_block_produced(
        self,
        block_ref: "ray.ObjectRef",
        size_bytes: int,
        producer_id: str,
    ) -> None:
        """Register a block and attribute its memory to producer_id.

        Registers a Ray Core out-of-scope callback so that when all references
        to block_ref are gone the bytes are automatically removed from the
        producer's usage.

        Idempotent: calling twice with the same block_ref is a no-op.
        """
        id_binary = block_ref.binary()
        with self._lock:
            if id_binary in self._registered_ids:
                return
            self._registered_ids.add(id_binary)
            self._bytes_by_producer[producer_id] += size_bytes

        def _on_object_freed(id_bytes: bytes) -> None:
            with self._lock:
                if id_bytes not in self._registered_ids:
                    # Already cleared (e.g. by clear()), nothing to do.
                    return
                self._registered_ids.discard(id_bytes)
                self._bytes_by_producer[producer_id] -= size_bytes

        # TODO(srayhome): This raises ValueError for blocks not owned by this
        # worker (e.g., materialized dataset passed to streaming_split). We may
        # need to guard this with RefBundle.owns_blocks or skip registration at
        # the InputDataBuffer level.
        registered = self._add_callback_fn(block_ref, _on_object_freed)
        if not registered:
            _on_object_freed(id_binary)

    def get_object_store_memory_usage(self, producer_id: str) -> int:
        """Total bytes of live blocks attributed to producer_id."""
        with self._lock:
            return self._bytes_by_producer.get(producer_id, 0)

    def clear(self) -> None:
        """Reset all accounting, e.g. on executor shutdown.

        Any previously registered Ray Core callbacks firing after clear()
        will be silently ignored because _registered_ids is empty.
        """
        with self._lock:
            self._registered_ids.clear()
            self._bytes_by_producer.clear()