class PipelineState(abc.ABC):
    """
    Object representing the state of a pipeline.

    This class defines the interface for accessing state variables like a dictionary,
    abstracting away whether the underlying storage is local, sharded, or global.
    """

    @abc.abstractmethod
    def __setitem__(self, key: str, value: Any):
        """
        Sets a state value for a given key.

        Args:
            key: The identifier for the state variable.
            value: The value to store.
        """

    @abc.abstractmethod
    def __getitem__(self, item: str) -> Any:
        """
        Retrieves a state value for a given key.

        Args:
            item: The identifier for the state variable.

        Returns:
            The value associated with the key.
        """

    @abc.abstractmethod
    def __contains__(self, item: str) -> bool:
        """
        Checks if a key exists in the state.

        Args:
            item: The identifier to check.

        Returns:
            True if the key exists, False otherwise.
        """

class PipelineStateHandler:
    """
    Manages the lifecycle and access patterns of pipeline states.

    This handler initializes the underlying storage and provides specific views
    (global or sharded) into that storage.
    """

    def __init__(self, sharding_spec: dict[str, ShardingSpecLeaf], num_shards: int):
        """
        Constructs a PipelineStateHandler object.

        Args:
            sharding_spec: A definition of how specific keys should be sharded.
            num_shards: The total number of shards in the pipeline.
        """

        self._storage = PipelineStateStorage(
            sharding_spec={(k,): v for k, v in sharding_spec.items()},
            num_shards=num_shards
        )

    def global_state(self) -> PipelineState:
        """
        Returns a view interface for accessing global state.

        Returns:
            A PipelineState interface that accesses the full, aggregated data.
        """

        return PipelineStateGlobal(self._storage)

    def sharded_state(self, shard_id: int) -> PipelineState:
        """
        Returns a view interface for accessing state specific to a shard ID.

        Args:
            shard_id: The index of the shard to access.

        Returns:
            A PipelineState interface that accesses partial data for the given shard.
        """

        return PipelineStateShard(self._storage, shard_id)

    def reset(self):
        """
        Resets the underlying storage, clearing all state.
        """

        self._storage.reset()