from typing import TYPE_CHECKING, Optional

import ray.cloudpickle as ray_pickle
from ray._common.utils import binary_to_hex, hex_to_binary
from ray.train._checkpoint import Checkpoint

if TYPE_CHECKING:
    from ray.data.preprocessor import Preprocessor

PREPROCESSOR_KEY = "preprocessor_pkl"


class FrameworkCheckpoint(Checkpoint):
    """A checkpoint to preserve the functionality of legacy
    framework-specific checkpoints.

    Example:

        >>> import tempfile
        >>> from ray.data.preprocessor import Preprocessor
        >>> checkpoint = FrameworkCheckpoint(tempfile.mkdtemp())
        >>> checkpoint.get_preprocessor() is None
        True
        >>> preprocessor = Preprocessor()
        >>> preprocessor._attr = 1234
        >>> checkpoint.set_preprocessor(preprocessor)
        >>> checkpoint.get_preprocessor()._attr
        1234
    """

    def get_preprocessor(self) -> Optional["Preprocessor"]:
        """Return the preprocessor stored in the checkpoint.

        .. warning::

            The checkpoint path must point to a **trusted** source.
            The preprocessor is stored as a pickle blob inside the checkpoint
            metadata. Loading a checkpoint from an untrusted path (shared
            storage, downloaded artifact, checkpoint produced by a different
            party) is equivalent to executing arbitrary Python code. Never
            call this method on a checkpoint you do not fully control.

        Returns:
            The preprocessor stored in the checkpoint, or ``None`` if no
            preprocessor was stored.
        """
        metadata = self.get_metadata()
        preprocessor_bytes = metadata.get(PREPROCESSOR_KEY)
        if preprocessor_bytes is None:
            return None
        return ray_pickle.loads(hex_to_binary(preprocessor_bytes))

    def set_preprocessor(self, preprocessor: "Preprocessor"):
        """Store a preprocessor with the checkpoint."""
        self.update_metadata(
            {PREPROCESSOR_KEY: binary_to_hex(ray_pickle.dumps(preprocessor))}
        )