#  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 flavor registry."""

from collections import defaultdict
from typing import DefaultDict, Dict, List, Type

from zenml.enums import StackComponentType
from zenml.integrations.registry import integration_registry
from zenml.logger import get_logger
from zenml.models import (
    FlavorFilterModel,
    FlavorResponseModel,
    FlavorUpdateModel,
)
from zenml.stack import Flavor
from zenml.zen_stores.base_zen_store import BaseZenStore

logger = get_logger(__name__)


class FlavorRegistry:
    """Registry for stack component flavors.

    The flavors defined by ZenML must be registered here.
    """

    def __init__(self) -> None:
        """Initialization of the flavors."""
        self._flavors: DefaultDict[
            StackComponentType, Dict[str, FlavorResponseModel]
        ] = defaultdict(dict)

    def register_flavors(self, store: BaseZenStore) -> None:
        """Register all flavors to the DB.

        Args:
            store: The instance of a store to use for persistence
        """
        self.register_builtin_flavors(store=store)
        self.register_integration_flavors(store=store)

    @property
    def builtin_flavors(self) -> List[Type[Flavor]]:
        """A list of all default in-built flavors.

        Returns:
            A list of builtin flavors.
        """
        from zenml.artifact_stores import LocalArtifactStoreFlavor
        from zenml.container_registries import (
            AzureContainerRegistryFlavor,
            DefaultContainerRegistryFlavor,
            DockerHubContainerRegistryFlavor,
            GCPContainerRegistryFlavor,
            GitHubContainerRegistryFlavor,
        )
        from zenml.image_builders import LocalImageBuilderFlavor
        from zenml.orchestrators import (
            LocalDockerOrchestratorFlavor,
            LocalOrchestratorFlavor,
        )
        from zenml.secrets_managers import LocalSecretsManagerFlavor

        flavors = [
            LocalArtifactStoreFlavor,
            LocalOrchestratorFlavor,
            LocalDockerOrchestratorFlavor,
            DefaultContainerRegistryFlavor,
            AzureContainerRegistryFlavor,
            DockerHubContainerRegistryFlavor,
            GCPContainerRegistryFlavor,
            GitHubContainerRegistryFlavor,
            LocalSecretsManagerFlavor,
            LocalImageBuilderFlavor,
        ]
        return flavors

    @property
    def integration_flavors(self) -> List[Type[Flavor]]:
        """A list of all default integration flavors.

        Returns:
            A list of integration flavors.
        """
        integrated_flavors = []
        for _, integration in integration_registry.integrations.items():
            for flavor in integration.flavors():
                integrated_flavors.append(flavor)

        return integrated_flavors

    def register_builtin_flavors(self, store: BaseZenStore) -> None:
        """Registers the default built-in flavors.

        Args:
            store: The instance of the zen_store to use
        """
        for flavor in self.builtin_flavors:
            flavor_request_model = flavor().to_model(
                integration="built-in",
                scoped_by_workspace=False,
                is_custom=False,
            )
            existing_flavor = store.list_flavors(
                FlavorFilterModel(
                    name=flavor_request_model.name,
                    type=flavor_request_model.type,
                )
            )
            if len(existing_flavor) == 0:
                store.create_flavor(flavor_request_model)
            else:
                flavor_update_model = FlavorUpdateModel.parse_obj(
                    flavor_request_model
                )
                store.update_flavor(existing_flavor[0].id, flavor_update_model)

    @staticmethod
    def register_integration_flavors(store: BaseZenStore) -> None:
        """Registers the flavors implemented by integrations.

        Args:
            store: The instance of the zen_store to use
        """
        for name, integration in integration_registry.integrations.items():
            integrated_flavors = integration.flavors()
            for flavor in integrated_flavors:
                flavor_request_model = flavor().to_model(
                    integration=name,
                    scoped_by_workspace=False,
                    is_custom=False,
                )
                existing_flavor = store.list_flavors(
                    FlavorFilterModel(
                        name=flavor_request_model.name,
                        type=flavor_request_model.type,
                    )
                )
                if len(existing_flavor) == 0:
                    store.create_flavor(flavor_request_model)
                else:
                    flavor_update_model = FlavorUpdateModel.parse_obj(
                        flavor_request_model
                    )
                    store.update_flavor(
                        existing_flavor[0].id, flavor_update_model
                    )
