#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at:
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
#  or implied. See the License for the specific language governing
#  permissions and limitations under the License.
"""Implementation of the MLflow model deployer pipeline step."""

from typing import Optional, cast

from mlflow.tracking import MlflowClient, artifact_utils

from zenml import get_step_context, step
from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
from zenml.client import Client
from zenml.constants import (
    DEFAULT_SERVICE_START_STOP_TIMEOUT,
    MLFLOW_MODEL_FORMAT,
)
from zenml.integrations.mlflow.experiment_trackers.mlflow_experiment_tracker import (
    MLFlowExperimentTracker,
)
from zenml.integrations.mlflow.mlflow_utils import build_runs_uri, is_mlflow_3x
from zenml.integrations.mlflow.model_deployers.mlflow_model_deployer import (
    MLFlowModelDeployer,
)
from zenml.integrations.mlflow.model_registries.mlflow_model_registry import (
    MLFlowModelRegistry,
)
from zenml.integrations.mlflow.services.mlflow_deployment import (
    MLFlowDeploymentConfig,
    MLFlowDeploymentService,
)
from zenml.logger import get_logger
from zenml.model_registries.base_model_registry import (
    ModelRegistryModelMetadata,
    ModelVersionStage,
)

logger = get_logger(__name__)


@step(enable_cache=False)
def mlflow_model_deployer_step(
    model: UnmaterializedArtifact,
    deploy_decision: bool = True,
    experiment_name: Optional[str] = None,
    run_name: Optional[str] = None,
    model_name: str = "model",
    workers: int = 1,
    mlserver: bool = False,
    timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
) -> Optional[MLFlowDeploymentService]:
    """Model deployer pipeline step for MLflow.

    This step deploys a model logged in the MLflow artifact store to a
    deployment service. The user would typically use this step in a pipeline
    that deploys a model that was already registered in the MLflow model
    registr either manually or by using the `mlflow_model_registry_step`.

    Args:
        deploy_decision: whether to deploy the model or not
        model: the model artifact to deploy
        experiment_name: Name of the MLflow experiment in which the model was
            logged.
        run_name: Name of the MLflow run in which the model was logged.
        model_name: the name of the MLflow model logged in the MLflow artifact
            store for the current pipeline.
        workers: number of workers to use for the prediction service
        mlserver: set to True to use the MLflow MLServer backend (see
            https://github.com/SeldonIO/MLServer). If False, the
            MLflow built-in scoring server will be used.
        timeout: the number of seconds to wait for the service to start/stop.

    Returns:
        MLflow deployment service or None if no service was deployed or found

    Raises:
        ValueError: if the MLflow experiment tracker is not found
    """
    model_deployer = cast(
        MLFlowModelDeployer, MLFlowModelDeployer.get_active_model_deployer()
    )

    experiment_tracker = Client().active_stack.experiment_tracker

    if not isinstance(experiment_tracker, MLFlowExperimentTracker):
        raise ValueError(
            "MLflow model deployer step requires an MLflow experiment "
            "tracker. Please add an MLflow experiment tracker to your "
            "stack."
        )

    # get pipeline name, step name and run id
    step_context = get_step_context()
    pipeline_name = step_context.pipeline.name
    run_name = step_context.pipeline_run.name
    step_name = step_context.step_run.name

    # Configure Mlflow so the client points to the correct store
    experiment_tracker.configure_mlflow()
    client = MlflowClient()
    mlflow_run_id = experiment_tracker.get_run_id(
        experiment_name=experiment_name or pipeline_name,
        run_name=run_name or run_name,
    )

    # Fetch the model URI from the MLflow artifact store
    model_uri = ""
    if mlflow_run_id and client.list_artifacts(mlflow_run_id, model_name):
        # In MLflow 3.x, use runs:/ URI which correctly resolves to model location
        # In MLflow 2.x, use the artifact URI directly
        if is_mlflow_3x():
            model_uri = build_runs_uri(mlflow_run_id, model_name)
        else:
            model_uri = artifact_utils.get_artifact_uri(  # type: ignore[no-untyped-call, unused-ignore]
                run_id=mlflow_run_id, artifact_path=model_name
            )

    predictor_cfg = MLFlowDeploymentConfig(
        model_name=model_name or "",
        model_uri=model_uri,
        workers=workers,
        mlserver=mlserver,
        pipeline_name=pipeline_name,
        pipeline_step_name=step_name,
        timeout=timeout,
    )

    # Fetch existing services with same pipeline name, step name and model name
    existing_services = model_deployer.find_model_server(
        config=predictor_cfg.model_dump(),
    )

    # Check whether to deploy a new service
    if model_uri and deploy_decision:
        new_service = cast(
            MLFlowDeploymentService,
            model_deployer.deploy_model(
                replace=True,
                config=predictor_cfg,
                timeout=timeout,
                service_type=MLFlowDeploymentService.SERVICE_TYPE,
            ),
        )
        logger.info(
            f"MLflow deployment service started and reachable at:\n"
            f"    {new_service.prediction_url}\n"
        )
        logger.info("Stopping existing services...")
        for existing_service in existing_services:
            existing_service.stop(timeout)
        return new_service

    # Check whether an existing service can be reused
    elif existing_services:
        if not model_uri:
            logger.info(
                f"Skipping model deployment because an MLflow model with name "
                f" `{model_name}` was not trained in the current pipeline run."
            )
        elif not deploy_decision:
            logger.info(
                "Skipping model deployment because the deployment decision "
                "was negative."
            )
        logger.info(
            f"Reusing last model server deployed by step `{step_name}` of "
            f"pipeline '{pipeline_name}' for model '{model_name}'..."
        )
        service = cast(MLFlowDeploymentService, existing_services[0])
        if not service.is_running:
            service.start(timeout)
        return service

    # Log a warning if no service was deployed or found
    elif deploy_decision and not model_uri:
        logger.warning(
            f"An MLflow model with name `{model_name}` was not "
            f"logged in the current pipeline run and no running MLflow "
            f"model server was found. Please ensure that your pipeline "
            f"includes a step with an MLflow experiment configured that "
            "trains a model and logs it to MLflow. This could also happen "
            "if the current pipeline run did not log an MLflow model  "
            f"because the training step was cached."
        )
    return None


@step(enable_cache=False)
def mlflow_model_registry_deployer_step(
    registry_model_name: str,
    registry_model_version: Optional[str] = None,
    registry_model_stage: Optional[ModelVersionStage] = None,
    replace_existing: bool = True,
    model_name: str = "model",
    workers: int = 1,
    mlserver: bool = False,
    timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
) -> MLFlowDeploymentService:
    """Model deployer pipeline step for MLflow.

    Args:
        registry_model_name: the name of the model in the model registry
        registry_model_version: the version of the model in the model registry
        registry_model_stage: the stage of the model in the model registry
        replace_existing: Whether to create a new deployment service or not
        model_name: the name of the MLflow model logged in the MLflow artifact
            store for the current pipeline.
        workers: number of workers to use for the prediction service
        mlserver: set to True to use the MLflow MLServer backend (see
            https://github.com/SeldonIO/MLServer). If False, the
            MLflow built-in scoring server will be used.
        timeout: the number of seconds to wait for the service to start/stop.

    Returns:
        MLflow deployment service

    Raises:
        ValueError: if neither registry_model_version nor registry_model_stage
            is not provided
        ValueError: if No MLflow experiment tracker is found in the current
            active stack
        LookupError: if no model version is found in the MLflow model registry.
    """
    if not registry_model_version and not registry_model_stage:
        raise ValueError(
            "Either registry_model_version or registry_model_stage must "
            "be provided in addition to registry_model_name to the MLflow "
            "model registry deployer step. Since the "
            "mlflow_model_registry_deployer_step is used in conjunction with "
            "the mlflow_model_registry."
        )

    model_deployer = cast(
        MLFlowModelDeployer, MLFlowModelDeployer.get_active_model_deployer()
    )

    # fetch the MLflow model registry
    model_registry = Client().active_stack.model_registry
    if not isinstance(model_registry, MLFlowModelRegistry):
        raise ValueError(
            "The MLflow model registry step can only be used with an "
            "MLflow model registry."
        )

    # fetch the model version
    model_version = None
    if registry_model_version:
        try:
            model_version = model_registry.get_model_version(
                name=registry_model_name,
                version=registry_model_version,
            )
        except KeyError:
            model_version = None
    elif registry_model_stage:
        model_version = model_registry.get_latest_model_version(
            name=registry_model_name,
            stage=registry_model_stage,
        )
    if not model_version:
        raise LookupError(
            f"No Model Version found for model name "
            f"{registry_model_name} and version "
            f"{registry_model_version} or stage "
            f"{registry_model_stage}"
        )
    if model_version.model_format != MLFLOW_MODEL_FORMAT:
        raise ValueError(
            f"Model version {model_version.version} of model "
            f"{model_version.registered_model.name} is not an MLflow model."
            f"Only MLflow models can be deployed with the MLflow deployer "
            f"using this step."
        )
    # fetch existing services with same pipeline name, step name and model name

    existing_services = (
        model_deployer.find_model_server(
            model_name=registry_model_name,
            model_version=model_version.version,
        )
        if replace_existing
        else None
    )
    # create a config for the new model service
    metadata = model_version.metadata or ModelRegistryModelMetadata()
    predictor_cfg = MLFlowDeploymentConfig(
        name=model_name or None,
        model_name=registry_model_name,
        model_version=model_version.version,
        model_uri=model_version.model_source_uri,
        workers=workers,
        mlserver=mlserver,
        pipeline_name=metadata.zenml_pipeline_name or "",
        pipeline_step_name=metadata.zenml_step_name or "",
        timeout=timeout,
    )

    # create a new model deployment and replace an old one if it exists
    new_service = cast(
        MLFlowDeploymentService,
        model_deployer.deploy_model(
            replace=True,
            config=predictor_cfg,
            timeout=timeout,
            service_type=MLFlowDeploymentService.SERVICE_TYPE,
        ),
    )

    logger.info(
        f"MLflow deployment service started and reachable at:\n"
        f"    {new_service.prediction_url}\n"
    )

    if existing_services:
        logger.info("Stopping existing services...")
        for existing_service in existing_services:
            existing_service.stop(timeout)

    return new_service
