#  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 BentoML model deployer pipeline step."""

from typing import List, Optional, Tuple, cast

import bentoml
from bentoml._internal.bento import bento

from zenml import get_step_context, step
from zenml.integrations.bentoml.model_deployers.bentoml_model_deployer import (
    BentoMLModelDeployer,
)
from zenml.integrations.bentoml.services.bentoml_container_deployment import (
    BentoMLContainerDeploymentConfig,
    BentoMLContainerDeploymentService,
)
from zenml.integrations.bentoml.services.bentoml_local_deployment import (
    BentoMLLocalDeploymentConfig,
    BentoMLLocalDeploymentService,
    SSLBentoMLParametersConfig,
)
from zenml.integrations.bentoml.services.deployment_type import (
    BentoMLDeploymentType,
)
from zenml.logger import get_logger
from zenml.services.service import BaseService, ServiceConfig
from zenml.services.service_type import ServiceType
from zenml.utils import source_utils

logger = get_logger(__name__)


@step(enable_cache=True)
def bentoml_model_deployer_step(
    bento: bento.Bento,
    model_name: str,
    port: int,
    deployment_type: BentoMLDeploymentType = BentoMLDeploymentType.LOCAL,
    deploy_decision: bool = True,
    workers: Optional[int] = 1,
    backlog: Optional[int] = 2048,
    production: bool = False,
    working_dir: Optional[str] = None,
    host: Optional[str] = None,
    image: Optional[str] = None,
    image_tag: Optional[str] = None,
    platform: Optional[str] = None,
    ssl_certfile: Optional[str] = None,
    ssl_keyfile: Optional[str] = None,
    ssl_keyfile_password: Optional[str] = None,
    ssl_version: Optional[str] = None,
    ssl_cert_reqs: Optional[str] = None,
    ssl_ca_certs: Optional[str] = None,
    ssl_ciphers: Optional[str] = None,
    timeout: int = 30,
) -> BaseService:
    """Model deployer pipeline step for BentoML.

    This step deploys a given Bento to a local BentoML http prediction server.

    Args:
        bento: the bento artifact to deploy
        model_name: the name of the model to deploy.
        port: the port to use for the prediction service.
        deployment_type: the type of deployment to use. Either "local" or "container".
        deploy_decision: whether to deploy the model or not
        workers: number of workers to use for the prediction service
        backlog: the number of requests to queue up before rejecting requests.
        production: whether to deploy the service in production mode.
        working_dir: the working directory to use for the prediction service.
        host: the host to use for the prediction service.
        image: the image to use for the prediction service in the container deployment.
        image_tag: the image tag to use for the prediction service in the container deployment.
        platform: the platform to use to build the image for the container deployment.
        ssl_certfile: the path to the ssl cert file.
        ssl_keyfile: the path to the ssl key file.
        ssl_keyfile_password: the password for the ssl key file.
        ssl_version: the ssl version to use.
        ssl_cert_reqs: the ssl cert requirements.
        ssl_ca_certs: the path to the ssl ca certs.
        ssl_ciphers: the ssl ciphers to use.
        timeout: the number of seconds to wait for the service to start/stop.

    Returns:
        BentoML deployment service
    """
    # get the current active model deployer
    model_deployer = cast(
        BentoMLModelDeployer, BentoMLModelDeployer.get_active_model_deployer()
    )

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

    # Return the apis endpoint of the defined service to use in the predict.
    # This is a workaround to get the endpoints of the service defined as functions
    # from the user code in the BentoML service.
    def service_apis(bento_tag: str) -> List[str]:
        # Add working dir in the bentoml load
        service = bentoml.load(
            bento_identifier=bento_tag,
            working_dir=working_dir or source_utils.get_source_root(),
        )
        apis = service.apis
        apis_paths = list(apis.keys())
        return apis_paths

    def create_deployment_config(
        deployment_type: BentoMLDeploymentType,
    ) -> Tuple[ServiceConfig, ServiceType]:
        common_config = {
            "model_name": model_name,
            "bento_tag": str(bento.tag),
            "model_uri": bento.info.labels.get("model_uri"),
            "bento_uri": bento.info.labels.get("bento_uri"),
            "apis": service_apis(str(bento.tag)),
            "host": host,
            "port": port,
            "pipeline_name": pipeline_name,
            "pipeline_step_name": step_name,
            "workers": workers,
            "backlog": backlog,
        }

        if deployment_type == BentoMLDeploymentType.CONTAINER:
            return BentoMLContainerDeploymentConfig(
                **common_config,
                image=image,
                image_tag=image_tag,
                platform=platform,
            ), BentoMLContainerDeploymentService.SERVICE_TYPE
        else:
            return BentoMLLocalDeploymentConfig(
                **common_config,
                working_dir=working_dir or source_utils.get_source_root(),
                ssl_parameters=SSLBentoMLParametersConfig(
                    ssl_certfile=ssl_certfile,
                    ssl_keyfile=ssl_keyfile,
                    ssl_keyfile_password=ssl_keyfile_password,
                    ssl_version=ssl_version,
                    ssl_cert_reqs=ssl_cert_reqs,
                    ssl_ca_certs=ssl_ca_certs,
                    ssl_ciphers=ssl_ciphers,
                ),
                production=production,
            ), BentoMLLocalDeploymentService.SERVICE_TYPE

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

    # Creating a new service with inactive state and status by default
    service: Optional[BaseService] = None
    if existing_services:
        if deployment_type == BentoMLDeploymentType.CONTAINER:
            service = cast(
                BentoMLContainerDeploymentService, existing_services[0]
            )
        else:
            service = cast(BentoMLLocalDeploymentService, existing_services[0])

    if not deploy_decision and existing_services:
        logger.info(
            f"Skipping model deployment because the model quality does not "
            f"meet the criteria. Reusing last model server deployed by step "
            f"'{step_name}' and pipeline '{pipeline_name}' for model "
            f"'{model_name}'..."
        )
        assert service is not None
        if not service.is_running:
            service.start(timeout=timeout)
        return service

    # create a new model deployment and replace an old one if it exists
    new_service: BaseService
    if deployment_type == BentoMLDeploymentType.CONTAINER:
        new_service = cast(
            BentoMLContainerDeploymentService,
            model_deployer.deploy_model(
                replace=True,
                config=cast(BentoMLContainerDeploymentConfig, predictor_cfg),
                timeout=timeout,
                service_type=service_type,
            ),
        )
    else:
        new_service = cast(
            BentoMLLocalDeploymentService,
            model_deployer.deploy_model(
                replace=True,
                config=cast(BentoMLLocalDeploymentConfig, predictor_cfg),
                timeout=timeout,
                service_type=service_type,
            ),
        )

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

    return new_service
