#  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:
#
#       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 ZenML Service class."""

import json
import time
from abc import abstractmethod
from functools import wraps
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    ClassVar,
    Dict,
    Generator,
    Optional,
    Tuple,
    Type,
    TypeVar,
)
from uuid import UUID

from pydantic import ConfigDict

from zenml.console import console
from zenml.logger import get_logger
from zenml.services.service_endpoint import BaseServiceEndpoint
from zenml.services.service_monitor import HTTPEndpointHealthMonitor
from zenml.services.service_status import ServiceState, ServiceStatus
from zenml.services.service_type import ServiceType
from zenml.utils import source_utils
from zenml.utils.typed_model import BaseTypedModel

logger = get_logger(__name__)

T = TypeVar("T", bound=Callable[..., Any])

if TYPE_CHECKING:
    from zenml.models.v2.core.service import ServiceResponse

ZENM_ENDPOINT_PREFIX = "zenml-"


def update_service_status(
    pre_status: Optional[ServiceState] = None,
    post_status: Optional[ServiceState] = None,
    error_status: ServiceState = ServiceState.ERROR,
) -> Callable[[T], T]:
    """A decorator to update the service status before and after a method call.

    This decorator is used to wrap service methods and update the service status
    before and after the method call. If the method raises an exception, the
    service status is updated to reflect the error state.

    Args:
        pre_status: the status to update before the method call.
        post_status: the status to update after the method call.
        error_status: the status to update if the method raises an exception.

    Returns:
        Callable[..., Any]: The wrapped method with exception handling.
    """

    def decorator(func: T) -> T:
        @wraps(func)
        def wrapper(self: "BaseService", *args: Any, **kwargs: Any) -> Any:
            if pre_status:
                self.status.update_state(pre_status, "")
            try:
                logger.info(f"Calling {func.__name__} method...")
                result = func(self, *args, **kwargs)
                logger.info(f"{func.__name__} method executed successfully.")
                if post_status:
                    self.status.update_state(post_status, "")
                return result
            except Exception as e:
                logger.error(
                    f"Error occurred in {func.__name__} method: {str(e)}"
                )
                self.status.update_state(error_status, str(e))
                raise

        return wrapper  # type: ignore

    return decorator


class ServiceConfig(BaseTypedModel):
    """Generic service configuration.

    Concrete service classes should extend this class and add additional
    attributes that they want to see reflected and used in the service
    configuration.

    Attributes:
        name: name for the service instance
        description: description of the service
        pipeline_name: name of the pipeline that spun up the service
        pipeline_step_name: name of the pipeline step that spun up the service
        run_name: name of the pipeline run that spun up the service.
    """

    name: str = ""
    description: str = ""
    pipeline_name: str = ""
    pipeline_step_name: str = ""
    model_name: str = ""
    model_version: str = ""
    service_name: str = ""

    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
    #  fields defined under base models. If not handled, this raises a warning.
    #  It is possible to suppress this warning message with the following
    #  configuration, however the ultimate solution is to rename these fields.
    #  Even though they do not cause any problems right now, if we are not
    #  careful we might overwrite some fields protected by pydantic.
    model_config = ConfigDict(protected_namespaces=())

    def __init__(self, **data: Any):
        """Initialize the service configuration.

        Args:
            **data: keyword arguments.

        Raises:
            ValueError: if neither 'name' nor 'model_name' is set.
        """
        super().__init__(**data)
        if self.name or self.model_name:
            self.service_name = data.get(
                "service_name",
                f"{ZENM_ENDPOINT_PREFIX}{self.name or self.model_name}",
            )
        else:
            raise ValueError("Either 'name' or 'model_name' must be set.")

    def get_service_labels(self) -> Dict[str, str]:
        """Get the service labels.

        Returns:
            a dictionary of service labels.
        """
        labels = {}
        for k, v in self.model_dump().items():
            label = f"zenml_{k}".upper()
            labels[label] = str(v)
        return labels


class BaseService(BaseTypedModel):
    """Base service class.

    This class implements generic functionality concerning the life-cycle
    management and tracking of an external service (e.g. process, container,
    Kubernetes deployment etc.).

    Attributes:
        SERVICE_TYPE: a service type descriptor with information describing
            the service class. Every concrete service class must define this.
        admin_state: the administrative state of the service.
        uuid: unique UUID identifier for the service instance.
        config: service configuration
        status: service status
        endpoint: optional service endpoint
    """

    SERVICE_TYPE: ClassVar[ServiceType]

    uuid: UUID
    admin_state: ServiceState = ServiceState.INACTIVE
    config: ServiceConfig
    status: ServiceStatus
    # TODO [ENG-703]: allow multiple endpoints per service
    endpoint: Optional[BaseServiceEndpoint] = None

    def __init__(
        self,
        **attrs: Any,
    ) -> None:
        """Initialize the service instance.

        Args:
            **attrs: keyword arguments.
        """
        super().__init__(**attrs)
        self.config.name = self.config.name or self.__class__.__name__

    @classmethod
    def from_model(cls, model: "ServiceResponse") -> "BaseService":
        """Loads a service from a model.

        Args:
            model: The ServiceResponse to load from.

        Returns:
            The loaded service object.

        Raises:
            ValueError: if the service source is not found in the model.
        """
        if not model.service_source:
            raise ValueError("Service source not found in the model.")
        class_: Type[BaseService] = source_utils.load_and_validate_class(
            source=model.service_source, expected_class=BaseService
        )
        return class_(
            uuid=model.id,
            admin_state=model.admin_state,
            config=model.config,
            status=model.status,
            service_type=model.service_type.model_dump(),
            endpoint=model.endpoint,
        )

    @classmethod
    def from_json(cls, json_str: str) -> "BaseTypedModel":
        """Loads a service from a JSON string.

        Args:
            json_str: the JSON string to load from.

        Returns:
            The loaded service object.
        """
        service_dict = json.loads(json_str)
        class_: Type[BaseService] = source_utils.load_and_validate_class(
            source=service_dict["type"], expected_class=BaseService
        )
        return class_.from_dict(service_dict)

    @abstractmethod
    def check_status(self) -> Tuple[ServiceState, str]:
        """Check the the current operational state of the external service.

        This method should be overridden by subclasses that implement
        concrete service tracking functionality.

        Returns:
            The operational state of the external service and a message
            providing additional information about that state (e.g. a
            description of the error if one is encountered while checking the
            service status).
        """

    @abstractmethod
    def get_logs(
        self, follow: bool = False, tail: Optional[int] = None
    ) -> Generator[str, bool, None]:
        """Retrieve the service logs.

        This method should be overridden by subclasses that implement
        concrete service tracking functionality.

        Args:
            follow: if True, the logs will be streamed as they are written
            tail: only retrieve the last NUM lines of log output.

        Returns:
            A generator that can be accessed to get the service logs.
        """

    def update_status(self) -> None:
        """Update the status of the service.

        Check the current operational state of the external service
        and update the local operational status information to reflect it.

        This method should be overridden by subclasses that implement
        concrete service status tracking functionality.
        """
        logger.debug(
            "Running status check for service '%s' ...",
            self,
        )
        try:
            state, err = self.check_status()
            logger.debug(
                "Status check results for service '%s': %s [%s]",
                self,
                state.name,
                err,
            )
            self.status.update_state(state, err)

            # don't bother checking the endpoint state if the service is not active
            if self.status.state == ServiceState.INACTIVE:
                return

            if self.endpoint:
                self.endpoint.update_status()
        except Exception as e:
            logger.error(
                f"Failed to update status for service '{self}': {e}",
                exc_info=True,
            )
            self.status.update_state(ServiceState.ERROR, str(e))

    def get_service_status_message(self) -> str:
        """Get a service status message.

        Returns:
            A message providing information about the current operational
            state of the service.
        """
        return (
            f"  Administrative state: `{self.admin_state.value}`\n"
            f"  Operational state: `{self.status.state.value}`\n"
            f"  Last status message: '{self.status.last_error}'\n"
        )

    def poll_service_status(self, timeout: int = 0) -> bool:
        """Polls the external service status.

        It does this until the service operational state matches the
        administrative state, the service enters a failed state, or the timeout
        is reached.

        Args:
            timeout: maximum time to wait for the service operational state
                to match the administrative state, in seconds

        Returns:
            True if the service operational state matches the administrative
            state, False otherwise.
        """
        time_remaining = timeout
        while True:
            if self.admin_state == ServiceState.ACTIVE and self.is_running:
                return True
            if self.admin_state == ServiceState.INACTIVE and self.is_stopped:
                return True
            if self.is_failed:
                return False
            if time_remaining <= 0:
                break
            time.sleep(1)
            time_remaining -= 1

        if timeout > 0:
            logger.error(
                f"Timed out waiting for service {self} to become "
                f"{self.admin_state.value}:\n"
                + self.get_service_status_message()
            )

        return False

    @property
    def is_running(self) -> bool:
        """Check if the service is currently running.

        This method will actively poll the external service to get its status
        and will return the result.

        Returns:
            True if the service is running and active (i.e. the endpoints are
            responsive, if any are configured), otherwise False.
        """
        self.update_status()
        return self.status.state == ServiceState.ACTIVE and (
            not self.endpoint or self.endpoint.is_active()
        )

    @property
    def is_stopped(self) -> bool:
        """Check if the service is currently stopped.

        This method will actively poll the external service to get its status
        and will return the result.

        Returns:
            True if the service is stopped, otherwise False.
        """
        self.update_status()
        return self.status.state == ServiceState.INACTIVE

    @property
    def is_failed(self) -> bool:
        """Check if the service is currently failed.

        This method will actively poll the external service to get its status
        and will return the result.

        Returns:
            True if the service is in a failure state, otherwise False.
        """
        self.update_status()
        return self.status.state == ServiceState.ERROR

    def provision(self) -> None:
        """Provisions resources to run the service.

        Raises:
            NotImplementedError: if the service does not implement provisioning functionality
        """
        raise NotImplementedError(
            f"Provisioning resources not implemented for {self}."
        )

    def deprovision(self, force: bool = False) -> None:
        """Deprovisions all resources used by the service.

        Args:
            force: if True, the service will be deprovisioned even if it is
                in a failed state.

        Raises:
            NotImplementedError: if the service does not implement
                deprovisioning functionality.
        """
        raise NotImplementedError(
            f"Deprovisioning resources not implemented for {self}."
        )

    def update(self, config: ServiceConfig) -> None:
        """Update the service configuration.

        Args:
            config: the new service configuration.
        """
        self.config = config

    @update_service_status(
        pre_status=ServiceState.PENDING_STARTUP,
        post_status=ServiceState.ACTIVE,
    )
    def start(self, timeout: int = 0) -> None:
        """Start the service and optionally wait for it to become active.

        Args:
            timeout: amount of time to wait for the service to become active.
                If set to 0, the method will return immediately after checking
                the service status.
        """
        with console.status(f"Starting service '{self}'.\n"):
            self.admin_state = ServiceState.ACTIVE
            self.provision()
            if timeout > 0 and not self.poll_service_status(timeout):
                logger.error(
                    f"Failed to start service {self}\n"
                    + self.get_service_status_message()
                )

    @update_service_status(
        pre_status=ServiceState.PENDING_SHUTDOWN,
        post_status=ServiceState.INACTIVE,
    )
    def stop(self, timeout: int = 0, force: bool = False) -> None:
        """Stop the service and optionally wait for it to shutdown.

        Args:
            timeout: amount of time to wait for the service to shutdown.
                If set to 0, the method will return immediately after checking
                the service status.
            force: if True, the service will be stopped even if it is not
                currently running.
        """
        with console.status(f"Stopping service '{self}'.\n"):
            self.admin_state = ServiceState.INACTIVE
            self.deprovision(force)
            if timeout > 0:
                self.poll_service_status(timeout)
                if not self.is_stopped:
                    logger.error(
                        f"Failed to stop service {self}. Last state: "
                        f"'{self.status.state.value}'. Last error: "
                        f"'{self.status.last_error}'"
                    )

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

        Returns:
            the prediction URL for the endpoint
        """
        prediction_url = None
        if isinstance(self, BaseDeploymentService) and self.prediction_url:
            prediction_url = self.prediction_url
        elif self.endpoint:
            prediction_url = (
                self.endpoint.status.uri if self.endpoint.status else None
            )
        return prediction_url

    def get_healthcheck_url(self) -> Optional[str]:
        """Gets the healthcheck URL for the endpoint.

        Returns:
            the healthcheck URL for the endpoint
        """
        return (
            self.endpoint.monitor.get_healthcheck_uri(self.endpoint)
            if (self.endpoint and self.endpoint.monitor)
            and isinstance(self.endpoint.monitor, HTTPEndpointHealthMonitor)
            else None
        )

    def __repr__(self) -> str:
        """String representation of the service.

        Returns:
            A string representation of the service.
        """
        return f"{self.__class__.__qualname__}[{self.uuid}] (type: {self.SERVICE_TYPE.type}, flavor: {self.SERVICE_TYPE.flavor})"

    def __str__(self) -> str:
        """String representation of the service.

        Returns:
            A string representation of the service.
        """
        return self.__repr__()

    model_config = ConfigDict(
        # validate attribute assignments
        validate_assignment=True,
    )


class BaseDeploymentService(BaseService):
    """Base class for deployment services."""

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

        Returns:
            the prediction URL for the endpoint
        """
        return None

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

        Returns:
            the healthcheck URL for the endpoint
        """
        return None
