import json
import logging
import os
import shutil
import uuid
from datetime import datetime
from pathlib import Path
from pprint import pformat
from typing import Any, Callable, Literal, NamedTuple, Optional, Union
from urllib.parse import urlparse

import yaml
from packaging.requirements import InvalidRequirement, Requirement

import mlflow
from mlflow.entities import LoggedModel, LoggedModelOutput, Metric
from mlflow.entities.model_registry.prompt_version import PromptVersion
from mlflow.environment_variables import (
    MLFLOW_PRINT_MODEL_URLS_ON_CREATION,
    MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING,
)
from mlflow.exceptions import MlflowException
from mlflow.models.auth_policy import AuthPolicy
from mlflow.models.resources import Resource, ResourceType, _ResourceBuilder
from mlflow.protos.databricks_pb2 import (
    INVALID_PARAMETER_VALUE,
    RESOURCE_DOES_NOT_EXIST,
)
from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
from mlflow.tracking._tracking_service.utils import _resolve_tracking_uri
from mlflow.tracking.artifact_utils import _download_artifact_from_uri, _upload_artifact_to_uri
from mlflow.tracking.fluent import (
    _get_active_model_context,
    _set_active_model_id,
    _use_logged_model,
)
from mlflow.utils.annotations import experimental
from mlflow.utils.databricks_utils import (
    _construct_databricks_logged_model_url,
    get_databricks_runtime_version,
    get_workspace_id,
    get_workspace_url,
    is_in_databricks_runtime,
)
from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
from mlflow.utils.environment import (
    _CONDA_ENV_FILE_NAME,
    _PYTHON_ENV_FILE_NAME,
    _REQUIREMENTS_FILE_NAME,
    _add_or_overwrite_requirements,
    _get_requirements_from_file,
    _remove_requirements,
    _write_requirements_to_file,
)
from mlflow.utils.file_utils import TempDir
from mlflow.utils.logging_utils import eprint
from mlflow.utils.mlflow_tags import MLFLOW_MODEL_IS_EXTERNAL
from mlflow.utils.uri import (
    append_to_uri_path,
    get_uri_scheme,
    is_databricks_uri,
)

_logger = logging.getLogger(__name__)

# NOTE: The MLMODEL_FILE_NAME constant is considered @developer_stable
MLMODEL_FILE_NAME = "MLmodel"
_DATABRICKS_FS_LOADER_MODULE = "databricks.feature_store.mlflow_model"
_LOG_MODEL_METADATA_WARNING_TEMPLATE = (
    "Logging model metadata to the tracking server has failed. The model artifacts "
    "have been logged successfully under %s. Set logging level to DEBUG via "
    '`logging.getLogger("mlflow").setLevel(logging.DEBUG)` to see the full traceback.'
)
_LOG_MODEL_MISSING_SIGNATURE_WARNING = (
    "Model logged without a signature. Signatures are required for Databricks UC model registry "
    "as they validate model inputs and denote the expected schema of model outputs. "
    f"Please visit https://www.mlflow.org/docs/{mlflow.__version__.replace('.dev0', '')}/"
    "model/signatures.html#how-to-set-signatures-on-models for instructions on setting "
    "signature on models."
)
_LOG_MODEL_MISSING_INPUT_EXAMPLE_WARNING = (
    "Model logged without a signature and input example. Please set `input_example` parameter "
    "when logging the model to auto infer the model signature."
)
# NOTE: The _MLFLOW_VERSION_KEY constant is considered @developer_stable
_MLFLOW_VERSION_KEY = "mlflow_version"
METADATA_FILES = [
    MLMODEL_FILE_NAME,
    _CONDA_ENV_FILE_NAME,
    _REQUIREMENTS_FILE_NAME,
    _PYTHON_ENV_FILE_NAME,
]
MODEL_CONFIG = "config"
MODEL_CODE_PATH = "model_code_path"
SET_MODEL_ERROR = (
    "Model should either be an instance of PyFuncModel, Langchain type, or LlamaIndex index."
)
ENV_VAR_FILE_NAME = "environment_variables.txt"
ENV_VAR_FILE_HEADER = (
    "# This file records environment variable names that are used during model inference.\n"
    "# They might need to be set when creating a serving endpoint from this model.\n"
    "# Note: it is not guaranteed that all environment variables listed here are required\n"
)


class ModelInfo:
    """
    The metadata of a logged MLflow Model.
    """

    def __init__(
        self,
        artifact_path: str,
        flavors: dict[str, Any],
        model_uri: str,
        model_uuid: str,
        run_id: str,
        saved_input_example_info: Optional[dict[str, Any]],
        signature,  # Optional[ModelSignature]
        utc_time_created: str,
        mlflow_version: str,
        metadata: Optional[dict[str, Any]] = None,
        registered_model_version: Optional[int] = None,
        env_vars: Optional[list[str]] = None,
        prompts: Optional[list[str]] = None,
        logged_model: Optional[LoggedModel] = None,
    ):
        self._artifact_path = artifact_path
        self._flavors = flavors
        self._model_uri = model_uri
        self._model_uuid = model_uuid
        self._run_id = run_id
        self._saved_input_example_info = saved_input_example_info
        self._signature = signature
        self._utc_time_created = utc_time_created
        self._mlflow_version = mlflow_version
        self._metadata = metadata
        self._prompts = prompts
        self._registered_model_version = registered_model_version
        self._env_vars = env_vars
        self._logged_model = logged_model

    @property
    def artifact_path(self) -> str:
        """
        Run relative path identifying the logged model.

        :getter: Retrieves the relative path of the logged model.
        :type: str
        """
        return self._artifact_path

    @property
    def flavors(self) -> dict[str, Any]:
        """
        A dictionary mapping the flavor name to how to serve
        the model as that flavor.

        :getter: Gets the mapping for the logged model's flavor that defines parameters used in
            serving of the model
        :type: Dict[str, str]

        .. code-block:: python
            :caption: Example flavor mapping for a scikit-learn logged model

            {
                "python_function": {
                    "model_path": "model.pkl",
                    "loader_module": "mlflow.sklearn",
                    "python_version": "3.8.10",
                    "env": "conda.yaml",
                },
                "sklearn": {
                    "pickled_model": "model.pkl",
                    "sklearn_version": "0.24.1",
                    "serialization_format": "cloudpickle",
                },
            }

        """
        return self._flavors

    @property
    def model_uri(self) -> str:
        """
        The ``model_uri`` of the logged model in the format ``'runs:/<run_id>/<artifact_path>'``.

        :getter: Gets the uri path of the logged model from the uri `runs:/<run_id>` path
            encapsulation
        :type: str
        """
        return self._model_uri

    @property
    def model_uuid(self) -> str:
        """
        The ``model_uuid`` of the logged model,
        e.g., ``'39ca11813cfc46b09ab83972740b80ca'``.

        :getter: [Legacy] Gets the model_uuid (run_id) of a logged model
        :type: str
        """
        return self._model_uuid

    @property
    def run_id(self) -> str:
        """
        The ``run_id`` associated with the logged model,
        e.g., ``'8ede7df408dd42ed9fc39019ef7df309'``

        :getter: Gets the run_id identifier for the logged model
        :type: str
        """
        return self._run_id

    @property
    def saved_input_example_info(self) -> Optional[dict[str, Any]]:
        """
        A dictionary that contains the metadata of the saved input example, e.g.,
        ``{"artifact_path": "input_example.json", "type": "dataframe", "pandas_orient": "split"}``.

        :getter: Gets the input example if specified during model logging
        :type: Optional[Dict[str, str]]
        """
        return self._saved_input_example_info

    @property
    def signature(self):  # -> Optional[ModelSignature]
        """
        A :py:class:`ModelSignature <mlflow.models.ModelSignature>` that describes the
        model input and output.

        :getter: Gets the model signature if it is defined
        :type: Optional[ModelSignature]
        """
        return self._signature

    @property
    def utc_time_created(self) -> str:
        """
        The UTC time that the logged model is created, e.g., ``'2022-01-12 05:17:31.634689'``.

        :getter: Gets the UTC formatted timestamp for when the model was logged
        :type: str
        """
        return self._utc_time_created

    @property
    def mlflow_version(self) -> str:
        """
        Version of MLflow used to log the model

        :getter: Gets the version of MLflow that was installed when a model was logged
        :type: str
        """
        return self._mlflow_version

    @property
    def env_vars(self) -> Optional[list[str]]:
        """
        Environment variables used during the model logging process.

        :getter: Gets the environment variables used during the model logging process.
        :type: Optional[List[str]]
        """
        return self._env_vars

    @env_vars.setter
    def env_vars(self, value: Optional[list[str]]) -> None:
        if value and not (isinstance(value, list) and all(isinstance(x, str) for x in value)):
            raise TypeError(f"env_vars must be a list of strings. Got: {value}")
        self._env_vars = value

    @property
    def metadata(self) -> Optional[dict[str, Any]]:
        """
        User defined metadata added to the model.

        :getter: Gets the user-defined metadata about a model
        :type: Optional[Dict[str, Any]]

        .. code-block:: python
            :caption: Example usage of Model Metadata

            # Create and log a model with metadata to the Model Registry

            from sklearn import datasets
            from sklearn.ensemble import RandomForestClassifier
            import mlflow
            from mlflow.models import infer_signature

            with mlflow.start_run():
                iris = datasets.load_iris()
                clf = RandomForestClassifier()
                clf.fit(iris.data, iris.target)
                signature = infer_signature(iris.data, iris.target)
                mlflow.sklearn.log_model(
                    clf,
                    name="iris_rf",
                    signature=signature,
                    registered_model_name="model-with-metadata",
                    metadata={"metadata_key": "metadata_value"},
                )

            # model uri for the above model
            model_uri = "models:/model-with-metadata/1"

            # Load the model and access the custom metadata from its ModelInfo object
            model = mlflow.pyfunc.load_model(model_uri=model_uri)
            assert model.metadata.get_model_info().metadata["metadata_key"] == "metadata_value"

            # Load the ModelInfo and access the custom metadata
            model_info = mlflow.models.get_model_info(model_uri=model_uri)
            assert model_info.metadata["metadata_key"] == "metadata_value"
        """
        return self._metadata

    @property
    def prompts(self) -> Optional[list[str]]:
        """A list of prompt URIs associated with the model."""
        return self._prompts

    @property
    def registered_model_version(self) -> Optional[int]:
        """
        The registered model version, if the model is registered.

        :getter: Gets the registered model version, if the model is registered in Model Registry.
        :setter: Sets the registered model version.
        :type: Optional[int]
        """
        return self._registered_model_version

    @registered_model_version.setter
    def registered_model_version(self, value) -> None:
        self._registered_model_version = value

    @property
    def model_id(self) -> str:
        """
        The model ID of the logged model.

        :getter: Gets the model ID of the logged model
        """
        return self._logged_model.model_id

    @property
    def metrics(self) -> Optional[list[Metric]]:
        """
        Returns the metrics of the logged model.

        :getter: Retrieves the metrics of the logged model
        """
        return self._logged_model.metrics

    @property
    def params(self) -> dict[str, str]:
        """
        Returns the parameters of the logged model.

        :getter: Retrieves the parameters of the logged model
        """
        return self._logged_model.params

    @property
    def tags(self) -> dict[str, str]:
        """
        Returns the tags of the logged model.

        :getter: Retrieves the tags of the logged model
        """
        return self._logged_model.tags

    @property
    def creation_timestamp(self) -> int:
        """
        Returns the creation timestamp of the logged model.

        :getter:  the creation timestamp of the logged model
        """
        return self._logged_model.creation_timestamp

    @property
    def name(self) -> str:
        """
        Returns the name of the logged model.
        """
        return self._logged_model.name


class Model:
    """
    An MLflow Model that can support multiple model flavors. Provides APIs for implementing
    new Model flavors.
    """

    def __init__(
        self,
        artifact_path=None,
        run_id=None,
        utc_time_created=None,
        flavors=None,
        signature=None,  # ModelSignature
        saved_input_example_info: Optional[dict[str, Any]] = None,
        model_uuid: Union[str, Callable[[], str], None] = lambda: uuid.uuid4().hex,
        mlflow_version: Union[str, None] = mlflow.version.VERSION,
        metadata: Optional[dict[str, Any]] = None,
        model_size_bytes: Optional[int] = None,
        resources: Optional[Union[str, list[Resource]]] = None,
        env_vars: Optional[list[str]] = None,
        auth_policy: Optional[AuthPolicy] = None,
        model_id: Optional[str] = None,
        prompts: Optional[list[str]] = None,
        **kwargs,
    ):
        # store model id instead of run_id and path to avoid confusion when model gets exported
        self.run_id = run_id
        self.artifact_path = artifact_path
        self.utc_time_created = str(utc_time_created or datetime.now())
        self.flavors = flavors if flavors is not None else {}
        self.signature = signature
        self.saved_input_example_info = saved_input_example_info
        self.model_uuid = model_uuid() if callable(model_uuid) else model_uuid
        self.mlflow_version = mlflow_version
        self.metadata = metadata
        self.prompts = prompts
        self.model_size_bytes = model_size_bytes
        self.resources = resources
        self.env_vars = env_vars
        self.auth_policy = auth_policy
        self.model_id = model_id
        self.__dict__.update(kwargs)

    def __eq__(self, other):
        if not isinstance(other, Model):
            return False
        return self.__dict__ == other.__dict__

    def get_input_schema(self):
        """
        Retrieves the input schema of the Model iff the model was saved with a schema definition.
        """
        return self.signature.inputs if self.signature is not None else None

    def get_output_schema(self):
        """
        Retrieves the output schema of the Model iff the model was saved with a schema definition.
        """
        return self.signature.outputs if self.signature is not None else None

    def get_params_schema(self):
        """
        Retrieves the parameters schema of the Model iff the model was saved with a schema
        definition.
        """
        return getattr(self.signature, "params", None)

    def get_serving_input(self, path: str) -> Optional[str]:
        """
        Load serving input example from a model directory. Returns None if there is no serving input
        example.

        Args:
            path: Path to the model directory.

        Returns:
            Serving input example or None if the model has no serving input example.
        """
        from mlflow.models.utils import _load_serving_input_example

        return _load_serving_input_example(self, path)

    def load_input_example(self, path: Optional[str] = None) -> Optional[str]:
        """
        Load the input example saved along a model. Returns None if there is no example metadata
        (i.e. the model was saved without example). Raises FileNotFoundError if there is model
        metadata but the example file is missing.

        Args:
            path: Model or run URI, or path to the `model` directory.
                e.g. models://<model_name>/<model_version>, runs:/<run_id>/<artifact_path>
                or /path/to/model

        Returns:
            Input example (NumPy ndarray, SciPy csc_matrix, SciPy csr_matrix,
            pandas DataFrame, dict) or None if the model has no example.
        """

        # Just-in-time import to only load example-parsing libraries (e.g. numpy, pandas, etc.) if
        # example is requested.
        from mlflow.models.utils import _read_example

        if path is None:
            path = (
                f"runs:/{self.run_id}/{self.artifact_path}"
                if self.model_id is None
                else self.artifact_path
            )

        return _read_example(self, str(path))

    def load_input_example_params(self, path: str):
        """
        Load the params of input example saved along a model. Returns None if there are no params in
        the input_example.

        Args:
            path: Path to the model directory.

        Returns:
            params (dict) or None if the model has no params.
        """
        from mlflow.models.utils import _read_example_params

        return _read_example_params(self, path)

    def add_flavor(self, name, **params) -> "Model":
        """Add an entry for how to serve the model in a given format."""
        self.flavors[name] = params
        return self

    @property
    def metadata(self) -> Optional[dict[str, Any]]:
        """
        Custom metadata dictionary passed to the model and stored in the MLmodel file.

        :getter: Retrieves custom metadata that have been applied to a model instance.
        :setter: Sets a dictionary of custom keys and values to be included with the model instance
        :type: Optional[Dict[str, Any]]

        Returns:
            A Dictionary of user-defined metadata iff defined.

        .. code-block:: python
            :caption: Example

            # Create and log a model with metadata to the Model Registry
            from sklearn import datasets
            from sklearn.ensemble import RandomForestClassifier
            import mlflow
            from mlflow.models import infer_signature

            with mlflow.start_run():
                iris = datasets.load_iris()
                clf = RandomForestClassifier()
                clf.fit(iris.data, iris.target)
                signature = infer_signature(iris.data, iris.target)
                mlflow.sklearn.log_model(
                    clf,
                    name="iris_rf",
                    signature=signature,
                    registered_model_name="model-with-metadata",
                    metadata={"metadata_key": "metadata_value"},
                )

            # model uri for the above model
            model_uri = "models:/model-with-metadata/1"

            # Load the model and access the custom metadata
            model = mlflow.pyfunc.load_model(model_uri=model_uri)
            assert model.metadata.metadata["metadata_key"] == "metadata_value"
        """

        return self._metadata

    @metadata.setter
    def metadata(self, value: Optional[dict[str, Any]]) -> None:
        self._metadata = value

    @property
    def signature(self):  # -> Optional[ModelSignature]
        """
        An optional definition of the expected inputs to and outputs from a model object, defined
        with both field names and data types. Signatures support both column-based and tensor-based
        inputs and outputs.

        :getter: Retrieves the signature of a model instance iff the model was saved with a
            signature definition.
        :setter: Sets a signature to a model instance.
        :type: Optional[ModelSignature]
        """
        return self._signature

    @signature.setter
    def signature(self, value) -> None:
        # signature cannot be set to `False`, which is used in `log_model` and `save_model` calls
        # to disable automatic signature inference
        if value is not False:
            self._signature = value

    @property
    def saved_input_example_info(self) -> Optional[dict[str, Any]]:
        """
        A dictionary that contains the metadata of the saved input example, e.g.,
        ``{"artifact_path": "input_example.json", "type": "dataframe", "pandas_orient": "split"}``.
        """
        return self._saved_input_example_info

    @saved_input_example_info.setter
    def saved_input_example_info(self, value: dict[str, Any]) -> None:
        self._saved_input_example_info = value

    @property
    def model_size_bytes(self) -> Optional[int]:
        """
        An optional integer that represents the model size in bytes

        :getter: Retrieves the model size if it's calculated when the model is saved
        :setter: Sets the model size to a model instance
        :type: Optional[int]
        """
        return self._model_size_bytes

    @model_size_bytes.setter
    def model_size_bytes(self, value: Optional[int]) -> None:
        self._model_size_bytes = value

    @experimental(version="2.13.0")
    @property
    def resources(self) -> dict[str, dict[ResourceType, list[dict[str, Any]]]]:
        """
        An optional dictionary that contains the resources required to serve the model.

        :getter: Retrieves the resources required to serve the model
        :setter: Sets the resources required to serve the model
        :type: Dict[str, Dict[ResourceType, List[Dict]]]
        """
        return self._resources

    @experimental(version="2.13.0")
    @resources.setter
    def resources(self, value: Optional[Union[str, list[Resource]]]) -> None:
        if isinstance(value, (Path, str)):
            serialized_resource = _ResourceBuilder.from_yaml_file(value)
        elif isinstance(value, list) and all(isinstance(resource, Resource) for resource in value):
            serialized_resource = _ResourceBuilder.from_resources(value)
        else:
            serialized_resource = value
        self._resources = serialized_resource

    @experimental(version="2.21.0")
    @property
    def auth_policy(self) -> dict[str, dict[str, Any]]:
        """
        An optional dictionary that contains the auth policy required to serve the model.

        :getter: Retrieves the auth_policy required to serve the model
        :setter: Sets the auth_policy required to serve the model
        :type: Dict[str, dict]
        """
        return self._auth_policy

    @experimental(version="2.21.0")
    @auth_policy.setter
    def auth_policy(self, value: Optional[Union[dict[str, Any], AuthPolicy]]) -> None:
        self._auth_policy = value.to_dict() if isinstance(value, AuthPolicy) else value

    @property
    def env_vars(self) -> Optional[list[str]]:
        return self._env_vars

    @env_vars.setter
    def env_vars(self, value: Optional[list[str]]) -> None:
        if value and not (isinstance(value, list) and all(isinstance(x, str) for x in value)):
            raise TypeError(f"env_vars must be a list of strings. Got: {value}")
        self._env_vars = value

    def _is_signature_from_type_hint(self):
        return self.signature._is_signature_from_type_hint if self.signature is not None else False

    def _is_type_hint_from_example(self):
        return self.signature._is_type_hint_from_example if self.signature is not None else False

    def get_model_info(self, logged_model: Optional[LoggedModel] = None) -> ModelInfo:
        """
        Create a :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
        model metadata.
        """
        return ModelInfo(
            artifact_path=self.artifact_path,
            flavors=self.flavors,
            model_uri=(
                f"models:/{self.model_id}"
                if self.model_id
                else f"runs:/{self.run_id}/{self.artifact_path}"
            ),
            model_uuid=self.model_uuid,
            run_id=self.run_id,
            saved_input_example_info=self.saved_input_example_info,
            signature=self.signature,
            utc_time_created=self.utc_time_created,
            mlflow_version=self.mlflow_version,
            metadata=self.metadata,
            prompts=self.prompts,
            env_vars=self.env_vars,
            logged_model=logged_model,
        )

    def get_tags_dict(self) -> dict[str, Any]:
        result = self.to_dict()

        tags = {
            key: value
            for key, value in result.items()
            if key in ["run_id", "utc_time_created", "artifact_path", "model_uuid"]
        }

        tags["flavors"] = {
            flavor: (
                {k: v for k, v in config.items() if k != "config"}
                if isinstance(config, dict)
                else config
            )
            for flavor, config in result.get("flavors", {}).items()
        }

        return tags

    def to_dict(self) -> dict[str, Any]:
        """Serialize the model to a dictionary."""
        res = {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
        databricks_runtime = get_databricks_runtime_version()
        if databricks_runtime:
            res["databricks_runtime"] = databricks_runtime
        if self.signature is not None:
            res["signature"] = self.signature.to_dict()
            res["is_signature_from_type_hint"] = self.signature._is_signature_from_type_hint
            res["type_hint_from_example"] = self.signature._is_type_hint_from_example
        if self.saved_input_example_info is not None:
            res["saved_input_example_info"] = self.saved_input_example_info
        if self.mlflow_version is None and _MLFLOW_VERSION_KEY in res:
            res.pop(_MLFLOW_VERSION_KEY)
        if self.metadata is not None:
            res["metadata"] = self.metadata
        if self.prompts is not None:
            res["prompts"] = self.prompts
        if self.resources is not None:
            res["resources"] = self.resources
        if self.model_size_bytes is not None:
            res["model_size_bytes"] = self.model_size_bytes
        if self.auth_policy is not None:
            res["auth_policy"] = self.auth_policy
        # Exclude null fields in case MLmodel file consumers such as Model Serving may not
        # handle them correctly.
        if self.artifact_path is None:
            res.pop("artifact_path", None)
        if self.run_id is None:
            res.pop("run_id", None)
        if self.env_vars is not None:
            res["env_vars"] = self.env_vars
        return res

    def to_yaml(self, stream=None) -> str:
        """Write the model as yaml string."""
        return yaml.safe_dump(self.to_dict(), stream=stream, default_flow_style=False)

    def __str__(self):
        return self.to_yaml()

    def to_json(self) -> str:
        """Write the model as json."""
        return json.dumps(self.to_dict())

    def save(self, path) -> None:
        """Write the model as a local YAML file."""
        with open(path, "w") as out:
            self.to_yaml(out)

    @classmethod
    def load(cls, path) -> "Model":
        """
        Load a model from its YAML representation.

        Args:
            path: A local filesystem path or URI referring to the MLmodel YAML file
                representation of the Model object or to the directory containing
                the MLmodel YAML file representation.

        Returns:
            An instance of Model.

        .. code-block:: python
            :caption: example

            from mlflow.models import Model

            # Load the Model object from a local MLmodel file
            model1 = Model.load("~/path/to/my/MLmodel")

            # Load the Model object from a remote model directory
            model2 = Model.load("s3://mybucket/path/to/my/model")
        """

        # Check if the path is a local directory and not remote
        sep = os.path.sep
        path = str(path).rstrip(sep)
        path_scheme = urlparse(path).scheme
        if (not path_scheme or path_scheme == "file") and not os.path.exists(path):
            raise MlflowException(
                f'Could not find an "{MLMODEL_FILE_NAME}" configuration file at "{path}"',
                RESOURCE_DOES_NOT_EXIST,
            )

        if ModelsArtifactRepository._is_logged_model_uri(path):
            path = ModelsArtifactRepository.get_underlying_uri(path)

        is_model_dir = path.rsplit(sep, maxsplit=1)[-1] != MLMODEL_FILE_NAME
        mlmodel_file_path = f"{path}/{MLMODEL_FILE_NAME}" if is_model_dir else path
        mlmodel_local_path = _download_artifact_from_uri(artifact_uri=mlmodel_file_path)
        with open(mlmodel_local_path) as f:
            model_dict = yaml.safe_load(f)
        return cls.from_dict(model_dict)

    @classmethod
    def from_dict(cls, model_dict) -> "Model":
        """Load a model from its YAML representation."""

        from mlflow.models.signature import ModelSignature

        model_dict = model_dict.copy()
        if "signature" in model_dict and isinstance(model_dict["signature"], dict):
            signature = ModelSignature.from_dict(model_dict["signature"])
            if "is_signature_from_type_hint" in model_dict:
                signature._is_signature_from_type_hint = model_dict.pop(
                    "is_signature_from_type_hint"
                )
            if "type_hint_from_example" in model_dict:
                signature._is_type_hint_from_example = model_dict.pop("type_hint_from_example")
            model_dict["signature"] = signature

        if "model_uuid" not in model_dict:
            model_dict["model_uuid"] = None

        if _MLFLOW_VERSION_KEY not in model_dict:
            model_dict[_MLFLOW_VERSION_KEY] = None
        return cls(**model_dict)

    # MLflow 2.x log_model API. Only spark flavors uses this API.
    # https://github.com/mlflow/mlflow/blob/fd2d9861fa52eeca178825c871d5d29fbb3b95c4/mlflow/models/model.py#L773-L982
    @format_docstring(LOG_MODEL_PARAM_DOCS)
    @classmethod
    def _log_v2(
        cls,
        artifact_path,
        flavor,
        registered_model_name=None,
        await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
        metadata=None,
        run_id=None,
        resources=None,
        auth_policy=None,
        prompts=None,
        **kwargs,
    ) -> ModelInfo:
        """
        Log model using supplied flavor module. If no run is active, this method will create a new
        active run.

        Args:
            artifact_path: Run relative path identifying the model.
            flavor: Flavor module to save the model with. The module must have
                the ``save_model`` function that will persist the model as a valid
                MLflow model.
            registered_model_name: If given, create a model version under
                ``registered_model_name``, also creating a registered model if
                one with the given name does not exist.
            await_registration_for: Number of seconds to wait for the model version to finish
                being created and is in ``READY`` status. By default, the
                function waits for five minutes. Specify 0 or None to skip
                waiting.
            metadata: {{ metadata }}
            run_id: The run ID to associate with this model. If not provided,
                a new run will be started.
            resources: {{ resources }}
            auth_policy: {{ auth_policy }}
            prompts: {{ prompts }}
            kwargs: Extra args passed to the model flavor.

        Returns:
            A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
            metadata of the logged model.
        """

        # Only one of Auth policy and resources should be defined

        if resources is not None and auth_policy is not None:
            raise ValueError("Only one of `resources`, and `auth_policy` can be specified.")

        from mlflow.utils.model_utils import _validate_and_get_model_config_from_file

        registered_model = None
        with TempDir() as tmp:
            local_path = tmp.path("model")
            if run_id is None:
                run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
            if prompts is not None:
                # Convert to URIs for serialization
                prompts = [pr.uri if isinstance(pr, PromptVersion) else pr for pr in prompts]
            mlflow_model = cls(
                artifact_path=artifact_path,
                run_id=run_id,
                metadata=metadata,
                resources=resources,
                auth_policy=auth_policy,
                prompts=prompts,
            )
            flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
            # `save_model` calls `load_model` to infer the model requirements, which may result in
            # __pycache__ directories being created in the model directory.
            for pycache in Path(local_path).rglob("__pycache__"):
                shutil.rmtree(pycache, ignore_errors=True)

            if is_in_databricks_runtime():
                _copy_model_metadata_for_uc_sharing(local_path, flavor)

            tracking_uri = _resolve_tracking_uri()
            serving_input = mlflow_model.get_serving_input(local_path)
            # We check signature presence here as some flavors have a default signature as a
            # fallback when not provided by user, which is set during flavor's save_model() call.
            if mlflow_model.signature is None:
                if serving_input is None:
                    _logger.warning(
                        _LOG_MODEL_MISSING_INPUT_EXAMPLE_WARNING, extra={"color": "red"}
                    )
                elif tracking_uri == "databricks" or get_uri_scheme(tracking_uri) == "databricks":
                    _logger.warning(_LOG_MODEL_MISSING_SIGNATURE_WARNING, extra={"color": "red"})

            env_vars = None
            # validate input example works for serving when logging the model
            if serving_input and kwargs.get("validate_serving_input", True):
                from mlflow.models import validate_serving_input
                from mlflow.utils.model_utils import RECORD_ENV_VAR_ALLOWLIST, env_var_tracker

                with env_var_tracker() as tracked_env_names:
                    try:
                        validate_serving_input(
                            model_uri=local_path,
                            serving_input=serving_input,
                        )
                    except Exception as e:
                        serving_input_msg = (
                            serving_input[:50] + "..." if len(serving_input) > 50 else serving_input
                        )
                        _logger.warning(
                            f"Failed to validate serving input example {serving_input_msg}. "
                            "Alternatively, you can avoid passing input example and pass model "
                            "signature instead when logging the model. To ensure the input example "
                            "is valid prior to serving, please try calling "
                            "`mlflow.models.validate_serving_input` on the model uri and serving "
                            "input example. A serving input example can be generated from model "
                            "input example using "
                            "`mlflow.models.convert_input_example_to_serving_input` function.\n"
                            f"Got error: {e}",
                            exc_info=_logger.isEnabledFor(logging.DEBUG),
                        )
                    env_vars = (
                        sorted(
                            x
                            for x in tracked_env_names
                            if any(env_var in x for env_var in RECORD_ENV_VAR_ALLOWLIST)
                        )
                        or None
                    )
            if env_vars:
                # Keep the environment variable file as it serves as a check
                # for displaying tips in Databricks serving endpoint
                env_var_path = Path(local_path, ENV_VAR_FILE_NAME)
                env_var_path.write_text(ENV_VAR_FILE_HEADER + "\n".join(env_vars) + "\n")
                if len(env_vars) <= 3:
                    env_var_info = "[" + ", ".join(env_vars) + "]"
                else:
                    env_var_info = "[" + ", ".join(env_vars[:3]) + ", ... " + "]"
                    f"(check file {ENV_VAR_FILE_NAME} in the model's artifact folder for full list"
                    " of environment variable names)"
                _logger.info(
                    "Found the following environment variables used during model inference: "
                    f"{env_var_info}. Please check if you need to set them when deploying the "
                    "model. To disable this message, set environment variable "
                    f"`{MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING.name}` to `false`."
                )
                mlflow_model.env_vars = env_vars
                # mlflow_model is updated, rewrite the MLmodel file
                mlflow_model.save(os.path.join(local_path, MLMODEL_FILE_NAME))

            # Associate prompts to the model Run
            if prompts:
                client = mlflow.MlflowClient()
                for prompt in prompts:
                    client.link_prompt_version_to_run(run_id, prompt)

            mlflow.tracking.fluent.log_artifacts(local_path, mlflow_model.artifact_path, run_id)

            # if the model_config kwarg is passed in, then log the model config as an params
            if model_config := kwargs.get("model_config"):
                if isinstance(model_config, str):
                    try:
                        file_extension = os.path.splitext(model_config)[1].lower()
                        if file_extension == ".json":
                            with open(model_config) as f:
                                model_config = json.load(f)
                        elif file_extension in [".yaml", ".yml"]:
                            model_config = _validate_and_get_model_config_from_file(model_config)
                        else:
                            _logger.warning(
                                "Unsupported file format for model config: %s. "
                                "Failed to load model config.",
                                model_config,
                            )
                    except Exception as e:
                        _logger.warning("Failed to load model config from %s: %s", model_config, e)

                try:
                    from mlflow.models.utils import _flatten_nested_params

                    # We are using the `/` separator to flatten the nested params
                    # since we are using the same separator to log nested metrics.
                    params_to_log = _flatten_nested_params(model_config, sep="/")
                except Exception as e:
                    _logger.warning("Failed to flatten nested params: %s", str(e))
                    params_to_log = model_config

                try:
                    mlflow.tracking.fluent.log_params(params_to_log or {}, run_id=run_id)
                except Exception as e:
                    _logger.warning("Failed to log model config as params: %s", str(e))

            try:
                mlflow.tracking.fluent._record_logged_model(mlflow_model, run_id)
            except MlflowException:
                # We need to swallow all mlflow exceptions to maintain backwards compatibility with
                # older tracking servers. Only print out a warning for now.
                _logger.warning(_LOG_MODEL_METADATA_WARNING_TEMPLATE, mlflow.get_artifact_uri())
                _logger.debug("", exc_info=True)

            if registered_model_name is not None:
                registered_model = mlflow.tracking._model_registry.fluent._register_model(
                    f"runs:/{run_id}/{mlflow_model.artifact_path}",
                    registered_model_name,
                    await_registration_for=await_registration_for,
                    local_model_path=local_path,
                )

            model_info = mlflow_model.get_model_info()
            if registered_model is not None:
                model_info.registered_model_version = registered_model.version

        # If the model signature is Mosaic AI Agent compatible, render a recipe for evaluation.
        from mlflow.models.display_utils import maybe_render_agent_eval_recipe

        maybe_render_agent_eval_recipe(model_info)

        return model_info

    @format_docstring(LOG_MODEL_PARAM_DOCS)
    @classmethod
    def log(
        cls,
        artifact_path,
        flavor,
        registered_model_name=None,
        await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
        metadata=None,
        run_id=None,
        resources=None,
        auth_policy=None,
        prompts=None,
        name: Optional[str] = None,
        model_type: Optional[str] = None,
        params: Optional[dict[str, Any]] = None,
        tags: Optional[dict[str, Any]] = None,
        step: int = 0,
        model_id: Optional[str] = None,
        **kwargs,
    ) -> ModelInfo:
        """
        Log model using supplied flavor module. If no run is active, this method will create a new
        active run.

        Args:
            artifact_path: Deprecated. Use `name` instead.
            flavor: Flavor module to save the model with. The module must have
                the ``save_model`` function that will persist the model as a valid
                MLflow model.
            registered_model_name: If given, create a model version under
                ``registered_model_name``, also creating a registered model if
                one with the given name does not exist.
            await_registration_for: Number of seconds to wait for the model version to finish
                being created and is in ``READY`` status. By default, the
                function waits for five minutes. Specify 0 or None to skip
                waiting.
            metadata: {{ metadata }}
            run_id: The run ID to associate with this model.
            resources: {{ resources }}
            auth_policy: {{ auth_policy }}
            prompts: {{ prompts }}
            name: The name of the model.
            model_type: {{ model_type }}
            params: {{ params }}
            tags: {{ tags }}
            step: {{ step }}
            model_id: {{ model_id }}
            kwargs: Extra args passed to the model flavor.

        Returns:
            A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
            metadata of the logged model.
        """
        if name is not None and artifact_path is not None:
            raise MlflowException.invalid_parameter_value(
                "Both `artifact_path` (deprecated) and `name` parameters were specified. "
                "Please only specify `name`."
            )
        elif artifact_path is not None:
            _logger.warning("`artifact_path` is deprecated. Please use `name` instead.")

        name = name or artifact_path

        def log_model_metrics_for_step(client, model_id, run_id, step):
            metric_names = client.get_run(run_id).data.metrics.keys()
            metrics_for_step = []
            for metric_name in metric_names:
                history = client.get_metric_history(run_id, metric_name)
                metrics_for_step.extend(
                    [
                        Metric(
                            key=metric.key,
                            value=metric.value,
                            timestamp=metric.timestamp,
                            step=metric.step,
                            dataset_name=metric.dataset_name,
                            dataset_digest=metric.dataset_digest,
                            run_id=metric.run_id,
                            model_id=model_id,
                        )
                        for metric in history
                        if metric.step == step and metric.model_id is None
                    ]
                )
            client.log_batch(run_id=run_id, metrics=metrics_for_step)

        # Only one of Auth policy and resources should be defined

        if resources is not None and auth_policy is not None:
            raise ValueError("Only one of `resources`, and `auth_policy` can be specified.")

        registered_model = None
        with TempDir() as tmp:
            local_path = tmp.path("model")

            tracking_uri = _resolve_tracking_uri()
            client = mlflow.MlflowClient(tracking_uri)
            if not run_id:
                run_id = active_run.info.run_id if (active_run := mlflow.active_run()) else None

            if model_id is not None:
                model = client.get_logged_model(model_id)
            else:
                params = {
                    **(params or {}),
                    **(client.get_run(run_id).data.params if run_id else {}),
                }
                model = mlflow.initialize_logged_model(
                    # TODO: Update model name
                    name=name,
                    source_run_id=run_id,
                    model_type=model_type,
                    params={key: str(value) for key, value in params.items()},
                    tags={key: str(value) for key, value in tags.items()}
                    if tags is not None
                    else None,
                )
                if (
                    MLFLOW_PRINT_MODEL_URLS_ON_CREATION.get()
                    and is_databricks_uri(tracking_uri)
                    and (workspace_url := get_workspace_url())
                ):
                    logged_model_url = _construct_databricks_logged_model_url(
                        workspace_url,
                        model.experiment_id,
                        model.model_id,
                        get_workspace_id(),
                    )
                    eprint(f"🔗 View Logged Model at: {logged_model_url}")

            with _use_logged_model(model=model):
                if run_id is not None:
                    client.log_outputs(
                        run_id=run_id, models=[LoggedModelOutput(model.model_id, step=step)]
                    )
                    log_model_metrics_for_step(
                        client=client, model_id=model.model_id, run_id=run_id, step=step
                    )

                if prompts is not None:
                    # Convert to URIs for serialization
                    prompts = [pr.uri if isinstance(pr, PromptVersion) else pr for pr in prompts]

                mlflow_model = cls(
                    artifact_path=model.artifact_location,
                    model_uuid=model.model_id,
                    run_id=run_id,
                    metadata=metadata,
                    resources=resources,
                    auth_policy=auth_policy,
                    prompts=prompts,
                    model_id=model.model_id,
                )
                flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
                # `save_model` calls `load_model` to infer the model requirements, which may result
                # in __pycache__ directories being created in the model directory.
                for pycache in Path(local_path).rglob("__pycache__"):
                    shutil.rmtree(pycache, ignore_errors=True)

                if is_in_databricks_runtime():
                    _copy_model_metadata_for_uc_sharing(local_path, flavor)

                serving_input = mlflow_model.get_serving_input(local_path)
                # We check signature presence here as some flavors have a default signature as a
                # fallback when not provided by user, which is set during flavor's save_model()
                # call.
                if mlflow_model.signature is None:
                    if serving_input is None:
                        _logger.warning(
                            _LOG_MODEL_MISSING_INPUT_EXAMPLE_WARNING, extra={"color": "red"}
                        )
                    elif (
                        tracking_uri == "databricks" or get_uri_scheme(tracking_uri) == "databricks"
                    ):
                        _logger.warning(
                            _LOG_MODEL_MISSING_SIGNATURE_WARNING, extra={"color": "red"}
                        )

                env_vars = None
                # validate input example works for serving when logging the model
                if serving_input and kwargs.get("validate_serving_input", True):
                    from mlflow.models import validate_serving_input
                    from mlflow.utils.model_utils import RECORD_ENV_VAR_ALLOWLIST, env_var_tracker

                    with env_var_tracker() as tracked_env_names:
                        try:
                            validate_serving_input(
                                model_uri=local_path,
                                serving_input=serving_input,
                            )
                        except Exception as e:
                            serving_input_msg = (
                                serving_input[:50] + "..."
                                if len(serving_input) > 50
                                else serving_input
                            )
                            _logger.warning(
                                f"Failed to validate serving input example {serving_input_msg}. "
                                "Alternatively, you can avoid passing input example and pass model "
                                "signature instead when logging the model. To ensure the input "
                                "example is valid prior to serving, please try calling "
                                "`mlflow.models.validate_serving_input` on the model uri and "
                                "serving input example. A serving input example can be generated "
                                "from model input example using "
                                "`mlflow.models.convert_input_example_to_serving_input` function.\n"
                                f"Got error: {e}",
                                exc_info=_logger.isEnabledFor(logging.DEBUG),
                            )
                        env_vars = (
                            sorted(
                                x
                                for x in tracked_env_names
                                if any(env_var in x for env_var in RECORD_ENV_VAR_ALLOWLIST)
                            )
                            or None
                        )
                if env_vars:
                    # Keep the environment variable file as it serves as a check
                    # for displaying tips in Databricks serving endpoint
                    env_var_path = Path(local_path, ENV_VAR_FILE_NAME)
                    env_var_path.write_text(ENV_VAR_FILE_HEADER + "\n".join(env_vars) + "\n")
                    if len(env_vars) <= 3:
                        env_var_info = "[" + ", ".join(env_vars) + "]"
                    else:
                        env_var_info = "[" + ", ".join(env_vars[:3]) + ", ... " + "]"
                        f"(check file {ENV_VAR_FILE_NAME} in the model's artifact folder for full "
                        "list of environment variable names)"
                    _logger.info(
                        "Found the following environment variables used during model inference: "
                        f"{env_var_info}. Please check if you need to set them when deploying the "
                        "model. To disable this message, set environment variable "
                        f"`{MLFLOW_RECORD_ENV_VARS_IN_MODEL_LOGGING.name}` to `false`."
                    )
                    mlflow_model.env_vars = env_vars
                    # mlflow_model is updated, rewrite the MLmodel file
                    mlflow_model.save(os.path.join(local_path, MLMODEL_FILE_NAME))

                client.log_model_artifacts(model.model_id, local_path)
                # If the model was previously identified as external, delete the tag because
                # the model now has artifacts in MLflow Model format
                if model.tags.get(MLFLOW_MODEL_IS_EXTERNAL, "false").lower() == "true":
                    client.delete_logged_model_tag(model.model_id, MLFLOW_MODEL_IS_EXTERNAL)
                # client.finalize_logged_model(model.model_id, status=LoggedModelStatus.READY)

                # Associate prompts to the model Run
                if prompts and run_id:
                    client = mlflow.MlflowClient()
                    for prompt in prompts:
                        client.link_prompt_version_to_run(run_id, prompt)

                # if the model_config kwarg is passed in, then log the model config as an params
                if model_config := kwargs.get("model_config"):
                    if isinstance(model_config, str):
                        try:
                            file_extension = os.path.splitext(model_config)[1].lower()
                            if file_extension == ".json":
                                with open(model_config) as f:
                                    model_config = json.load(f)
                            elif file_extension in [".yaml", ".yml"]:
                                from mlflow.utils.model_utils import (
                                    _validate_and_get_model_config_from_file,
                                )

                                model_config = _validate_and_get_model_config_from_file(
                                    model_config
                                )
                            else:
                                _logger.warning(
                                    "Unsupported file format for model config: %s. "
                                    "Failed to load model config.",
                                    model_config,
                                )
                        except Exception as e:
                            _logger.warning(
                                "Failed to load model config from %s: %s", model_config, e
                            )

                    try:
                        from mlflow.models.utils import _flatten_nested_params

                        # We are using the `/` separator to flatten the nested params
                        # since we are using the same separator to log nested metrics.
                        params_to_log = _flatten_nested_params(model_config, sep="/")
                    except Exception as e:
                        _logger.warning("Failed to flatten nested params: %s", str(e))
                        params_to_log = model_config

                    try:
                        # do not log params to run if run_id is None, since that could trigger
                        # a new run to be created
                        if run_id:
                            mlflow.tracking.fluent.log_params(params_to_log or {}, run_id=run_id)
                    except Exception as e:
                        _logger.warning("Failed to log model config as params: %s", str(e))

            if registered_model_name is not None:
                registered_model = mlflow.tracking._model_registry.fluent._register_model(
                    f"models:/{model.model_id}",
                    registered_model_name,
                    await_registration_for=await_registration_for,
                    local_model_path=local_path,
                )
            model_info = mlflow_model.get_model_info(model)
            if registered_model is not None:
                model_info.registered_model_version = registered_model.version

        # If the model signature is Mosaic AI Agent compatible, render a recipe for evaluation.
        from mlflow.models.display_utils import maybe_render_agent_eval_recipe

        maybe_render_agent_eval_recipe(model_info)

        return model_info


def _copy_model_metadata_for_uc_sharing(local_path: str, flavor) -> None:
    """
    Copy model metadata files to a sub-directory 'metadata',
    For Databricks Unity Catalog sharing use-cases.

    Args:
        local_path: Local path to the model directory.
        flavor: Flavor module to save the model with.
    """
    from mlflow.models.wheeled_model import _ORIGINAL_REQ_FILE_NAME, WheeledModel

    metadata_path = os.path.join(local_path, "metadata")
    if isinstance(flavor, WheeledModel):
        # wheeled model updates several metadata files in original model directory
        # copy these updated metadata files to the 'metadata' subdirectory
        os.makedirs(metadata_path, exist_ok=True)
        for file_name in METADATA_FILES + [
            _ORIGINAL_REQ_FILE_NAME,
        ]:
            src_file_path = os.path.join(local_path, file_name)
            if os.path.exists(src_file_path):
                dest_file_path = os.path.join(metadata_path, file_name)
                shutil.copyfile(src_file_path, dest_file_path)
    else:
        os.makedirs(metadata_path, exist_ok=True)
        for file_name in METADATA_FILES:
            src_file_path = os.path.join(local_path, file_name)
            if os.path.exists(src_file_path):
                dest_file_path = os.path.join(metadata_path, file_name)
                shutil.copyfile(src_file_path, dest_file_path)


def get_model_info(model_uri: str) -> ModelInfo:
    """
    Get metadata for the specified model, such as its input/output signature.

    Args:
        model_uri: The location, in URI format, of the MLflow model. For example:

            - ``/Users/me/path/to/local/model``
            - ``relative/path/to/local/model``
            - ``s3://my_bucket/path/to/model``
            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
            - ``models:/<model_name>/<model_version>``
            - ``models:/<model_name>/<stage>``
            - ``mlflow-artifacts:/path/to/model``

            For more information about supported URI schemes, see
            `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
            artifact-locations>`_.

    Returns:
        A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
        metadata of the logged model.

    .. code-block:: python
        :caption: Example usage of get_model_info

        import mlflow.models
        import mlflow.sklearn
        from sklearn.ensemble import RandomForestRegressor

        with mlflow.start_run() as run:
            params = {"n_estimators": 3, "random_state": 42}
            X = [[0, 1]]
            y = [1]
            signature = mlflow.models.infer_signature(X, y)
            rfr = RandomForestRegressor(**params).fit(X, y)
            mlflow.log_params(params)
            mlflow.sklearn.log_model(rfr, name="sklearn-model", signature=signature)

        model_uri = f"runs:/{run.info.run_id}/sklearn-model"
        # Get model info with model_uri
        model_info = mlflow.models.get_model_info(model_uri)
        # Get model signature directly
        model_signature = model_info.signature
        assert model_signature == signature
    """
    return Model.load(model_uri).get_model_info()


class Files(NamedTuple):
    requirements: Path
    conda: Path


def get_model_requirements_files(resolved_uri: str) -> Files:
    requirements_txt_file = _download_artifact_from_uri(
        artifact_uri=append_to_uri_path(resolved_uri, _REQUIREMENTS_FILE_NAME)
    )
    conda_yaml_file = _download_artifact_from_uri(
        artifact_uri=append_to_uri_path(resolved_uri, _CONDA_ENV_FILE_NAME)
    )

    return Files(
        Path(requirements_txt_file),
        Path(conda_yaml_file),
    )


def update_model_requirements(
    model_uri: str,
    operation: Literal["add", "remove"],
    requirement_list: list[str],
) -> None:
    """
    Add or remove requirements from a model's conda.yaml and requirements.txt files.

    The process involves downloading these two files from the model artifacts
    (if they're non-local), updating them with the specified requirements,
    and then overwriting the existing files. Should the artifact repository
    associated with the model artifacts disallow overwriting, this function will
    fail.

    Note that model registry URIs (i.e. URIs in the form ``models:/``) are not
    supported, as artifacts in the model registry are intended to be read-only.

    If adding requirements, the function will overwrite any existing requirements
    that overlap, or else append the new requirements to the existing list.

    If removing requirements, the function will ignore any version specifiers,
    and remove all the specified package names. Any requirements that are not
    found in the existing files will be ignored.

    Args:
        model_uri: The location, in URI format, of the MLflow model. For example:

            - ``/Users/me/path/to/local/model``
            - ``relative/path/to/local/model``
            - ``s3://my_bucket/path/to/model``
            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
            - ``mlflow-artifacts:/path/to/model``

            For more information about supported URI schemes, see
            `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
            artifact-locations>`_.

        operation: The operation to perform. Must be one of "add" or "remove".

        requirement_list: A list of requirements to add or remove from the model.
            For example: ["numpy==1.20.3", "pandas>=1.3.3"]
    """
    resolved_uri = model_uri
    if ModelsArtifactRepository.is_models_uri(model_uri):
        if not ModelsArtifactRepository._is_logged_model_uri(model_uri):
            raise MlflowException(
                f'Failed to set requirements on "{model_uri}". '
                + "Model URIs with the `models:/` scheme are not supported.",
                INVALID_PARAMETER_VALUE,
            )
        resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)
    elif RunsArtifactRepository.is_runs_uri(model_uri):
        resolved_uri = RunsArtifactRepository.get_underlying_uri(model_uri)

    _logger.info(f"Retrieving model requirements files from {resolved_uri}...")
    local_paths = get_model_requirements_files(resolved_uri)
    conda_yaml_path = local_paths.conda
    requirements_txt_path = local_paths.requirements

    old_conda_reqs = _get_requirements_from_file(conda_yaml_path)
    old_requirements_reqs = _get_requirements_from_file(requirements_txt_path)

    requirements = []
    invalid_requirements = {}
    for s in requirement_list:
        try:
            requirements.append(Requirement(s.strip().lower()))
        except InvalidRequirement as e:
            invalid_requirements[s] = e
    if invalid_requirements:
        raise MlflowException.invalid_parameter_value(
            f"Found invalid requirements: {invalid_requirements}"
        )
    if operation == "add":
        updated_conda_reqs = _add_or_overwrite_requirements(requirements, old_conda_reqs)
        updated_requirements_reqs = _add_or_overwrite_requirements(
            requirements, old_requirements_reqs
        )
    else:
        updated_conda_reqs = _remove_requirements(requirements, old_conda_reqs)
        updated_requirements_reqs = _remove_requirements(requirements, old_requirements_reqs)

    _write_requirements_to_file(conda_yaml_path, updated_conda_reqs)
    _write_requirements_to_file(requirements_txt_path, updated_requirements_reqs)

    # just print conda reqs here to avoid log spam
    # it should be the same as requirements.txt anyway
    _logger.info(
        "Done updating requirements!\n\n"
        f"Old requirements:\n{pformat([str(req) for req in old_conda_reqs])}\n\n"
        f"Updated requirements:\n{pformat(updated_conda_reqs)}\n"
    )

    _logger.info(f"Uploading updated requirements files to {resolved_uri}...")
    _upload_artifact_to_uri(conda_yaml_path, resolved_uri)
    _upload_artifact_to_uri(requirements_txt_path, resolved_uri)


__mlflow_model__ = None


def _validate_langchain_model(model):
    from langchain_core.runnables.base import Runnable

    from mlflow.models.utils import _validate_and_get_model_code_path

    if isinstance(model, str):
        return _validate_and_get_model_code_path(model, None)

    if not isinstance(model, Runnable):
        raise MlflowException.invalid_parameter_value(
            "Model must be a Langchain Runnable type or path to a Langchain model, "
            f"got {type(model)}"
        )

    return model


def _validate_llama_index_model(model):
    from mlflow.llama_index.model import _validate_and_prepare_llama_index_model_or_path

    return _validate_and_prepare_llama_index_model_or_path(model, None)


@experimental(version="2.13.0")
def set_model(model) -> None:
    """
    When logging model as code, this function can be used to set the model object
    to be logged.

    Args:
        model: The model object to be logged. Supported model types are:

                - A Python function or callable object.
                - A Langchain model or path to a Langchain model.
                - A Llama Index model or path to a Llama Index model.
    """
    from mlflow.pyfunc import PythonModel

    if isinstance(model, str):
        raise mlflow.MlflowException(SET_MODEL_ERROR)

    if isinstance(model, PythonModel) or callable(model):
        globals()["__mlflow_model__"] = model
        return

    for validate_function in [_validate_langchain_model, _validate_llama_index_model]:
        try:
            globals()["__mlflow_model__"] = validate_function(model)
            return
        except Exception:
            pass

    raise mlflow.MlflowException(SET_MODEL_ERROR)


def _update_active_model_id_based_on_mlflow_model(mlflow_model: Model):
    """
    Update the current active model ID based on the provided MLflow model.
    Only set the active model ID if it is not already set by the user.
    This is useful for setting the active model ID when loading a model
    to ensure traces generated are associated with the loaded model.
    """
    if mlflow_model.model_id is None:
        return
    amc = _get_active_model_context()
    # only set the active model if the model is not set by the user
    if amc.model_id != mlflow_model.model_id and not amc.set_by_user:
        _set_active_model_id(model_id=mlflow_model.model_id)
