import os
import shutil
from typing import TYPE_CHECKING, Callable, Optional, Union

from enum import Enum
import tempfile
import tensorflow as tf
from tensorflow import keras
import warnings

from ray.air._internal.checkpointing import save_preprocessor_to_dir
from ray.air.checkpoint import Checkpoint
from ray.air.constants import MODEL_KEY, PREPROCESSOR_KEY
from ray.train.data_parallel_trainer import _load_checkpoint_dict
from ray.train._internal.framework_checkpoint import FrameworkCheckpoint
from ray.util.annotations import PublicAPI

if TYPE_CHECKING:
    from ray.data.preprocessor import Preprocessor


@PublicAPI(stability="beta")
class TensorflowCheckpoint(FrameworkCheckpoint):
    """A :py:class:`~ray.train.Checkpoint` with TensorFlow-specific functionality."""

    MODEL_FILENAME_KEY = "_model_filename"

    @classmethod
    def from_model(
        cls,
        model: keras.Model,
        *,
        preprocessor: Optional["Preprocessor"] = None,
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras model.

        The checkpoint created with this method needs to be paired with
        `model` when used.

        Args:
            model: The Keras model, whose weights are stored in the checkpoint.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` containing the specified model.

        Examples:

            .. testcode::

                from ray.train.tensorflow import TensorflowCheckpoint
                import tensorflow as tf

                model = tf.keras.applications.resnet.ResNet101()
                checkpoint = TensorflowCheckpoint.from_model(model)

            .. testoutput::
                :options: +MOCK
                :hide:

                ...  # Model may or may not be downloaded

        """
        tempdir = tempfile.mkdtemp()
        filename = "model.keras"
        model.save(os.path.join(tempdir, filename))

        checkpoint = cls.from_directory(tempdir)
        if preprocessor:
            checkpoint.set_preprocessor(preprocessor)
        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
        return checkpoint

    @classmethod
    def from_h5(
        cls, file_path: str, *, preprocessor: Optional["Preprocessor"] = None
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
        model from H5 format.

        The checkpoint generated by this method contains all the information needed.
        Thus no `model` is needed to be supplied when using this checkpoint.

        Args:
            file_path: The path to the .h5 file to load model from. This is the
                same path that is used for ``model.save(path)``.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` converted from h5 format.

        """
        if not os.path.isfile(file_path) or not file_path.endswith(".h5"):
            raise ValueError(
                "Please supply a h5 file path to `TensorflowCheckpoint.from_h5()`."
            )
        tempdir = tempfile.mkdtemp()
        filename = os.path.basename(file_path)
        new_checkpoint_file = os.path.join(tempdir, filename)
        shutil.copy(file_path, new_checkpoint_file)

        checkpoint = cls.from_directory(tempdir)
        if preprocessor:
            checkpoint.set_preprocessor(preprocessor)
        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: filename})
        return checkpoint

    @classmethod
    def from_saved_model(
        cls, dir_path: str, *, preprocessor: Optional["Preprocessor"] = None
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.train.Checkpoint` that stores a Keras
        model from SavedModel format.

        The checkpoint generated by this method contains all the information needed.
        Thus no `model` is needed to be supplied when using this checkpoint.

        Args:
            dir_path: The directory containing the saved model. This is the same
                directory as used by ``model.save(dir_path)``.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` converted from SavedModel format.

        """
        if not os.path.isdir(dir_path):
            raise ValueError(
                "Please supply a directory to `TensorflowCheckpoint.from_saved_model`"
            )
        tempdir = tempfile.mkdtemp()
        shutil.copytree(dir_path, tempdir, dirs_exist_ok=True)

        checkpoint = cls.from_directory(tempdir)
        if preprocessor:
            checkpoint.set_preprocessor(preprocessor)
        # NOTE: The entire directory is the checkpoint.
        checkpoint.update_metadata({cls.MODEL_FILENAME_KEY: "."})
        return checkpoint

    def get_model(
        self,
    ) -> tf.keras.Model:
        """Retrieve the model stored in this checkpoint.

        Returns:
            The Tensorflow Keras model stored in the checkpoint.
        """
        metadata = self.get_metadata()
        if self.MODEL_FILENAME_KEY not in metadata:
            raise ValueError(
                "`TensorflowCheckpoint` cannot retrieve the model if you override the "
                "checkpoint metadata. Please use `Checkpoint.update_metadata` instead."
            )
        model_filename = metadata[self.MODEL_FILENAME_KEY]
        with self.as_directory() as checkpoint_dir:
            model_path = os.path.join(checkpoint_dir, model_filename)
            return keras.models.load_model(model_path)


@PublicAPI(stability="beta")
class LegacyTensorflowCheckpoint(Checkpoint):
    """A :py:class:`~ray.air.checkpoint.Checkpoint` with TensorFlow-specific
    functionality.

    Create this from a generic :py:class:`~ray.air.checkpoint.Checkpoint` by calling
    ``TensorflowCheckpoint.from_checkpoint(ckpt)``.
    """

    _SERIALIZED_ATTRS = Checkpoint._SERIALIZED_ATTRS + ("_flavor", "_h5_file_path")

    class Flavor(Enum):
        # Various flavors with which TensorflowCheckpoint is generated.
        # This is necessary metadata to decide how to load model from a checkpoint.
        MODEL_WEIGHTS = 1
        SAVED_MODEL = 2
        H5 = 3

    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self._flavor = None
        # Will only be set when `self._flavor` is `H5`.
        self._h5_file_path = None

    @classmethod
    def from_model(
        cls,
        model: keras.Model,
        *,
        preprocessor: Optional["Preprocessor"] = None,
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.air.checkpoint.Checkpoint` that stores a Keras
        model.

        The checkpoint created with this method needs to be paired with
        `model` when used.

        Args:
            model: The Keras model, whose weights are stored in the checkpoint.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` containing the specified model.

        Examples:

            .. testcode::

                from ray.train.tensorflow import TensorflowCheckpoint
                import tensorflow as tf

                model = tf.keras.applications.resnet.ResNet101()
                checkpoint = TensorflowCheckpoint.from_model(model)

            .. testoutput::
                :options: +MOCK
                :hide:

                ...  # Model may or may not be downloaded

        """
        checkpoint = cls.from_dict(
            {PREPROCESSOR_KEY: preprocessor, MODEL_KEY: model.get_weights()}
        )
        checkpoint._flavor = cls.Flavor.MODEL_WEIGHTS
        return checkpoint

    @classmethod
    def from_h5(
        cls, file_path: str, *, preprocessor: Optional["Preprocessor"] = None
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.air.checkpoint.Checkpoint` that stores a Keras
        model from H5 format.

        The checkpoint generated by this method contains all the information needed.
        Thus no `model` is needed to be supplied when using this checkpoint.

        `file_path` must maintain validity even after this function returns.
        Some new files/directories may be added to the parent directory of `file_path`,
        as a side effect of this method.

        Args:
            file_path: The path to the .h5 file to load model from. This is the
                same path that is used for ``model.save(path)``.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` converted from h5 format.

        """
        if not os.path.isfile(file_path) or not file_path.endswith(".h5"):
            raise ValueError(
                "Please supply a h5 file path to `TensorflowCheckpoint.from_h5()`."
            )
        dir_path = os.path.dirname(os.path.abspath(file_path))
        if preprocessor:
            save_preprocessor_to_dir(preprocessor, dir_path)
        checkpoint = cls.from_directory(dir_path)
        checkpoint._flavor = cls.Flavor.H5
        checkpoint._h5_file_path = os.path.basename(file_path)
        return checkpoint

    @classmethod
    def from_saved_model(
        cls, dir_path: str, *, preprocessor: Optional["Preprocessor"] = None
    ) -> "TensorflowCheckpoint":
        """Create a :py:class:`~ray.air.checkpoint.Checkpoint` that stores a Keras
        model from SavedModel format.

        The checkpoint generated by this method contains all the information needed.
        Thus no `model` is needed to be supplied when using this checkpoint.

        `dir_path` must maintain validity even after this function returns.
        Some new files/directories may be added to `dir_path`, as a side effect
        of this method.

        Args:
            dir_path: The directory containing the saved model. This is the same
                directory as used by ``model.save(dir_path)``.
            preprocessor: A fitted preprocessor to be applied before inference.

        Returns:
            A :py:class:`TensorflowCheckpoint` converted from SavedModel format.

        """
        if preprocessor:
            save_preprocessor_to_dir(preprocessor, dir_path)
        checkpoint = cls.from_directory(dir_path)
        checkpoint._flavor = cls.Flavor.SAVED_MODEL
        return checkpoint

    def get_model(
        self,
        model: Optional[Union[tf.keras.Model, Callable[[], tf.keras.Model]]] = None,
    ) -> tf.keras.Model:
        """Retrieve the model stored in this checkpoint.

        Args:
            model: This arg is expected only if the original checkpoint
                was created via `TensorflowCheckpoint.from_model`.
            model_definition: This parameter is deprecated. Use `model` instead.

        Returns:
            The Tensorflow Keras model stored in the checkpoint.
        """
        if model is not None and self._flavor is not self.Flavor.MODEL_WEIGHTS:
            warnings.warn(
                "TensorflowCheckpoint was created from "
                "TensorflowCheckpoint.from_saved_model` or "
                "`TensorflowCheckpoint.from_h5`, which already contains all the "
                "information needed. This means "
                "if you are using TensorflowPredictor directly, you should do "
                "`TensorflowPredictor.from_checkpoint(checkpoint)` by "
                "removing kwargs `model=`."
            )
        if self._flavor is self.Flavor.MODEL_WEIGHTS and model is None:
            raise ValueError(
                "Expecting `model` argument when checkpoint is "
                "saved through `TensorflowCheckpoint.from_model()`."
            )

        if self._flavor is self.Flavor.MODEL_WEIGHTS:
            model = self._get_model_from_weights(model)
        elif self._flavor is self.Flavor.H5:
            model = self._get_model_from_h5()
        elif self._flavor == self.Flavor.SAVED_MODEL:
            model = self._get_model_from_saved_model()
        else:
            raise RuntimeError(
                "Avoid directly using `from_dict` or "
                "`from_directory` directly. Make sure "
                "that the checkpoint was generated by "
                "`TensorflowCheckpoint.from_model`, "
                "`TensorflowCheckpoint.from_saved_model` or "
                "`TensorflowCheckpoint.from_h5`."
            )

        return model

    def _get_model_from_weights(
        self, model: Union[tf.keras.Model, Callable[[], tf.keras.Model]]
    ) -> tf.keras.Model:
        assert self._flavor is self.Flavor.MODEL_WEIGHTS
        assert model is not None

        if callable(model):
            model = model()

        model_weights, _ = _load_checkpoint_dict(self, "TensorflowTrainer")
        model.set_weights(model_weights)

        return model

    def _get_model_from_h5(self) -> tf.keras.Model:
        assert self._flavor is self.Flavor.H5

        with self.as_directory() as checkpoint_dir:
            return keras.models.load_model(
                os.path.join(checkpoint_dir, self._h5_file_path)
            )

    def _get_model_from_saved_model(self) -> tf.keras.Model:
        assert self._flavor is self.Flavor.SAVED_MODEL

        with self.as_directory() as checkpoint_dir:
            return keras.models.load_model(checkpoint_dir)
