#  Copyright (c) ZenML GmbH 2024. 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:
#
#       https://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 vLLM Inference Server Service."""

import argparse
import os
from typing import Any, List, Optional, Union

from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS
from zenml.logger import get_logger
from zenml.models.v2.misc.service import ServiceType
from zenml.services import (
    HTTPEndpointHealthMonitor,
    HTTPEndpointHealthMonitorConfig,
    LocalDaemonService,
    LocalDaemonServiceConfig,
    LocalDaemonServiceEndpoint,
    LocalDaemonServiceEndpointConfig,
    ServiceEndpointProtocol,
)
from zenml.services.service import BaseDeploymentService

logger = get_logger(__name__)


VLLM_PREDICTION_URL_PATH = "v1"
VLLM_HEALTHCHECK_URL_PATH = "health"


class VLLMDeploymentEndpointConfig(LocalDaemonServiceEndpointConfig):
    """vLLM deployment service configuration.

    Attributes:
        prediction_url_path: URI subpath for prediction requests
    """

    prediction_url_path: str


class VLLMDeploymentEndpoint(LocalDaemonServiceEndpoint):
    """A service endpoint exposed by the vLLM deployment daemon.

    Attributes:
        config: service endpoint configuration
    """

    config: VLLMDeploymentEndpointConfig
    monitor: HTTPEndpointHealthMonitor

    @property
    def prediction_url(self) -> Optional[str]:
        """Gets the prediction URL for the endpoint.

        Returns:
            the prediction URL for the endpoint
        """
        uri = self.status.uri
        if not uri:
            return None
        return os.path.join(uri, self.config.prediction_url_path)


class VLLMServiceConfig(LocalDaemonServiceConfig):
    """vLLM service configurations."""

    model: str
    port: int
    host: Optional[str] = None
    blocking: bool = True
    # If unspecified, model name or path will be used.
    tokenizer: Optional[str] = None
    served_model_name: Optional[Union[str, List[str]]] = None
    # Trust remote code from huggingface.
    trust_remote_code: Optional[bool] = False
    # ['auto', 'slow', 'mistral']
    tokenizer_mode: Optional[str] = "auto"
    # ['auto', 'half', 'float16', 'bfloat16', 'float', 'float32']
    dtype: Optional[str] = "auto"
    # The specific model version to use. It can be a branch name, a tag name, or a commit id.
    # If unspecified, will use the default version.
    revision: Optional[str] = None


class VLLMDeploymentService(LocalDaemonService, BaseDeploymentService):
    """vLLM Inference Server Deployment Service."""

    SERVICE_TYPE = ServiceType(
        name="vllm-deployment",
        type="model-serving",
        flavor="vllm",
        description="vLLM Inference prediction service",
    )
    config: VLLMServiceConfig
    endpoint: VLLMDeploymentEndpoint

    def __init__(self, config: VLLMServiceConfig, **attrs: Any):
        """Initialize the vLLM deployment service.

        Args:
            config: service configuration
            attrs: additional attributes to set on the service
        """
        if isinstance(config, VLLMServiceConfig) and "endpoint" not in attrs:
            endpoint = VLLMDeploymentEndpoint(
                config=VLLMDeploymentEndpointConfig(
                    protocol=ServiceEndpointProtocol.HTTP,
                    port=config.port,
                    ip_address=config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
                    prediction_url_path=VLLM_PREDICTION_URL_PATH,
                ),
                monitor=HTTPEndpointHealthMonitor(
                    config=HTTPEndpointHealthMonitorConfig(
                        healthcheck_uri_path=VLLM_HEALTHCHECK_URL_PATH,
                    )
                ),
            )
            attrs["endpoint"] = endpoint
        super().__init__(config=config, **attrs)

    def run(self) -> None:
        """Start the service."""
        logger.info(
            "Starting vLLM inference server service as blocking "
            "process... press CTRL+C once to stop it."
        )

        self.endpoint.prepare_for_start()

        import uvloop
        from vllm.entrypoints.openai.api_server import (
            run_server,
        )
        from vllm.entrypoints.openai.cli_args import (
            make_arg_parser,
        )
        from vllm.utils import (
            FlexibleArgumentParser,
        )

        try:
            parser: argparse.ArgumentParser = make_arg_parser(
                FlexibleArgumentParser()
            )
            # pass in empty list to get default args
            # otherwise it will try to get the args from sys.argv
            # and if there's a --config in there, it will want to use
            # that file for vLLM configuration, which is not what we want
            args: argparse.Namespace = parser.parse_args(args=[])
            # Override port with the available port
            self.config.port = self.endpoint.status.port or self.config.port

            # Update the arguments in place
            args.__dict__.update(self.config.model_dump())
            uvloop.run(run_server(args=args))
        except KeyboardInterrupt:
            logger.info("Stopping vLLM prediction service...")

    @property
    def prediction_url(self) -> Optional[str]:
        """Gets the prediction URL for the endpoint.

        Returns:
            the prediction URL for the endpoint
        """
        if not self.is_running:
            return None
        return self.endpoint.prediction_url

    def predict(self, data: "Any") -> "Any":
        """Make a prediction using the service.

        Args:
            data: data to make a prediction on

        Returns:
            The prediction result.

        Raises:
            Exception: if the service is not running
            ValueError: if the prediction endpoint is unknown.
        """
        if not self.is_running:
            raise Exception(
                "vLLM Inference service is not running. "
                "Please start the service before making predictions."
            )
        if self.endpoint.prediction_url is not None:
            from openai import OpenAI

            client = OpenAI(
                api_key="EMPTY",
                base_url=self.endpoint.prediction_url,
            )
            models = client.models.list()
            model = models.data[0].id
            result = client.completions.create(model=model, prompt=data)
            # TODO: We can add support for client.chat.completions.create
        else:
            raise ValueError("No endpoint known for prediction.")
        return result
