#  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.
"""ZenML server deployer singleton implementation."""

from typing import ClassVar, Dict, Generator, Optional, Type

from zenml.config.global_config import GlobalConfiguration
from zenml.enums import ServerProviderType, StoreType
from zenml.logger import get_logger
from zenml.utils.singleton import SingletonMetaClass
from zenml.zen_server.deploy.base_provider import BaseServerProvider
from zenml.zen_server.deploy.deployment import (
    LocalServerDeployment,
    LocalServerDeploymentConfig,
)
from zenml.zen_server.deploy.exceptions import (
    ServerDeploymentError,
    ServerDeploymentNotFoundError,
    ServerProviderNotFoundError,
)
from zenml.zen_stores.base_zen_store import BaseZenStore
from zenml.zen_stores.rest_zen_store import RestZenStoreConfiguration

logger = get_logger(__name__)


class LocalServerDeployer(metaclass=SingletonMetaClass):
    """Local server deployer singleton.

    This class is responsible for managing the various server provider
    implementations and for directing server deployment lifecycle requests to
    the responsible provider. It acts as a facade built on top of the various
    server providers.
    """

    _providers: ClassVar[Dict[ServerProviderType, BaseServerProvider]] = {}

    @classmethod
    def register_provider(cls, provider: Type[BaseServerProvider]) -> None:
        """Register a server provider.

        Args:
            provider: The server provider to register.

        Raises:
            TypeError: If a provider with the same type is already registered.
        """
        if provider.TYPE in cls._providers:
            raise TypeError(
                f"Server provider '{provider.TYPE}' is already registered."
            )
        logger.debug(f"Registering server provider '{provider.TYPE}'.")
        cls._providers[provider.TYPE] = provider()

    @classmethod
    def get_provider(
        cls, provider_type: ServerProviderType
    ) -> BaseServerProvider:
        """Get the server provider associated with a provider type.

        Args:
            provider_type: The server provider type.

        Returns:
            The server provider associated with the provider type.

        Raises:
            ServerProviderNotFoundError: If no provider is registered for the
                given provider type.
        """
        if provider_type not in cls._providers:
            raise ServerProviderNotFoundError(
                f"Server provider '{provider_type}' is not registered."
            )
        return cls._providers[provider_type]

    def initialize_local_database(self) -> None:
        """Initialize the local ZenML database."""
        default_store_cfg = GlobalConfiguration().get_default_store()
        BaseZenStore.create_store(default_store_cfg)

    def deploy_server(
        self,
        config: LocalServerDeploymentConfig,
        timeout: Optional[int] = None,
        restart: bool = False,
    ) -> LocalServerDeployment:
        """Deploy the local ZenML server or update the existing deployment.

        Args:
            config: The server deployment configuration.
            timeout: The timeout in seconds to wait until the deployment is
                successful. If not supplied, the default timeout value specified
                by the provider is used.
            restart: If True, the existing server deployment will be torn down
                and a new server will be deployed.

        Returns:
            The local server deployment.
        """
        # Ensure that the local database is always initialized before any local
        # server is deployed or updated.
        self.initialize_local_database()

        try:
            self.get_server()
        except ServerDeploymentNotFoundError:
            pass
        else:
            return self.update_server(
                config=config, timeout=timeout, restart=restart
            )

        provider_name = config.provider.value
        provider = self.get_provider(config.provider)

        logger.info(f"Deploying a local {provider_name} ZenML server.")
        return provider.deploy_server(config, timeout=timeout)

    def update_server(
        self,
        config: LocalServerDeploymentConfig,
        timeout: Optional[int] = None,
        restart: bool = False,
    ) -> LocalServerDeployment:
        """Update an existing local ZenML server deployment.

        Args:
            config: The new server deployment configuration.
            timeout: The timeout in seconds to wait until the deployment is
                successful. If not supplied, a default timeout value of 30
                seconds is used.
            restart: If True, the existing server deployment will be torn down
                and a new server will be deployed.

        Returns:
            The updated server deployment.
        """
        # this will also raise ServerDeploymentNotFoundError if the server
        # does not exist
        existing_server = self.get_server()

        provider = self.get_provider(config.provider)
        if existing_server.config.provider != config.provider or restart:
            existing_provider = self.get_provider(
                existing_server.config.provider
            )

            # Tear down the existing server deployment
            existing_provider.remove_server(
                existing_server.config, timeout=timeout
            )

            # Deploy a new server with the new provider
            return provider.deploy_server(config, timeout=timeout)

        return provider.update_server(config, timeout=timeout)

    def remove_server(
        self,
        timeout: Optional[int] = None,
    ) -> None:
        """Tears down and removes all resources and files associated with the local ZenML server deployment.

        Args:
            timeout: The timeout in seconds to wait until the deployment is
                successfully torn down. If not supplied, a provider specific
                default timeout value is used.
        """
        # this will also raise ServerDeploymentNotFoundError if the server
        # does not exist
        try:
            server = self.get_server()
        except ServerDeploymentNotFoundError:
            return

        provider_name = server.config.provider.value
        provider = self.get_provider(server.config.provider)

        if self.is_connected_to_server():
            try:
                self.disconnect_from_server()
            except Exception as e:
                logger.warning(
                    f"Failed to disconnect from the local server: {e}"
                )

        logger.info(f"Tearing down the local {provider_name} ZenML server.")
        provider.remove_server(server.config, timeout=timeout)

    def is_connected_to_server(self) -> bool:
        """Check if the ZenML client is currently connected to the local ZenML server.

        Returns:
            True if the ZenML client is connected to the local ZenML server, False
            otherwise.
        """
        # this will also raise ServerDeploymentNotFoundError if the server
        # does not exist
        try:
            server = self.get_server()
        except ServerDeploymentNotFoundError:
            return False

        gc = GlobalConfiguration()
        return (
            server.status is not None
            and server.status.url is not None
            and gc.store_configuration.url == server.status.url
        )

    def connect_to_server(
        self,
    ) -> None:
        """Connect to the local ZenML server instance.

        Raises:
            ServerDeploymentError: If the local ZenML server is not running or
                is unreachable.
        """
        # this will also raise ServerDeploymentNotFoundError if the server
        # does not exist
        server = self.get_server()
        provider_name = server.config.provider.value

        gc = GlobalConfiguration()
        if not server.status or not server.status.url:
            raise ServerDeploymentError(
                f"The local {provider_name} ZenML server is not currently "
                "running or is unreachable."
            )

        store_config = RestZenStoreConfiguration(
            url=server.status.url,
        )

        if gc.store_configuration == store_config:
            logger.info(
                "Your client is already connected to the local "
                f"{provider_name} ZenML server."
            )
            return

        logger.info(
            f"Connecting to the local {provider_name} ZenML server "
            f"({store_config.url})."
        )

        gc.set_store(store_config)

        logger.info(
            f"Connected to the local {provider_name} ZenML server "
            f"({store_config.url})."
        )

    def disconnect_from_server(
        self,
    ) -> None:
        """Disconnect from the ZenML server instance."""
        gc = GlobalConfiguration()
        store_cfg = gc.store_configuration

        if store_cfg.type != StoreType.REST:
            logger.info(
                "Your client is not currently connected to a ZenML server."
            )
            return

        logger.info(
            f"Disconnecting from the local ({store_cfg.url}) ZenML server."
        )

        gc.set_default_store()

        logger.info("Disconnected from the local ZenML server.")

    def get_server(
        self,
    ) -> LocalServerDeployment:
        """Get the local server deployment.

        Returns:
            The local server deployment.

        Raises:
            ServerDeploymentNotFoundError: If no local server deployment is
                found.
        """
        for provider in self._providers.values():
            try:
                return provider.get_server(
                    LocalServerDeploymentConfig(provider=provider.TYPE)
                )
            except ServerDeploymentNotFoundError:
                pass

        raise ServerDeploymentNotFoundError(
            "No local server deployment was found."
        )

    def get_server_logs(
        self,
        follow: bool = False,
        tail: Optional[int] = None,
    ) -> Generator[str, bool, None]:
        """Retrieve the logs for the local ZenML server.

        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.
        """
        # this will also raise ServerDeploymentNotFoundError if the server
        # does not exist
        server = self.get_server()

        provider_name = server.config.provider.value
        provider = self.get_provider(server.config.provider)

        logger.info(
            f"Fetching logs from the local {provider_name} ZenML server..."
        )
        return provider.get_server_logs(
            server.config, follow=follow, tail=tail
        )
