#  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 a base container registry class."""

import re
from typing import TYPE_CHECKING, Optional, Tuple, Type, cast

from pydantic import Field, field_validator

from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
from zenml.enums import StackComponentType
from zenml.models import ServiceConnectorRequirements
from zenml.secret.schemas import BasicAuthSecretSchema
from zenml.stack.authentication_mixin import (
    AuthenticationConfigMixin,
    AuthenticationMixin,
)
from zenml.stack.flavor import Flavor
from zenml.utils import docker_utils

if TYPE_CHECKING:
    from docker.client import DockerClient


class BaseContainerRegistryConfig(AuthenticationConfigMixin):
    """Base config for a container registry.

    Configuration for connecting to container image registries.
    Field descriptions are defined inline using Field() descriptors.
    """

    uri: str = Field(
        description="Container registry URI (e.g., 'gcr.io' for Google Container "
        "Registry, 'docker.io' for Docker Hub, 'registry.gitlab.com' for GitLab "
        "Container Registry, 'ghcr.io' for GitHub Container Registry). This is "
        "the base URL where container images will be pushed to and pulled from."
    )
    default_repository: Optional[str] = Field(
        default=None,
        description="Default repository namespace for image storage (e.g., "
        "'username' for Docker Hub, 'project-id' for GCR, 'organization' for "
        "GitHub Container Registry). If not specified, images will be stored at "
        "the registry root. For Docker Hub this would mean only official images "
        "can be pushed.",
    )

    @field_validator("uri")
    @classmethod
    def strip_trailing_slash(cls, uri: str) -> str:
        """Removes trailing slashes from the URI.

        Args:
            uri: The URI to be stripped.

        Returns:
            The URI without trailing slashes.
        """
        return uri.rstrip("/")

    @property
    def is_local(self) -> bool:
        """Checks if this stack component is running locally.

        Returns:
            True if this config is for a local component, False otherwise.
        """
        return bool(re.fullmatch(r"localhost:[0-9]{4,5}", self.uri))


class BaseContainerRegistry(AuthenticationMixin):
    """Base class for all ZenML container registries."""

    _docker_client: Optional["DockerClient"] = None

    @property
    def config(self) -> BaseContainerRegistryConfig:
        """Returns the `BaseContainerRegistryConfig` config.

        Returns:
            The configuration.
        """
        return cast(BaseContainerRegistryConfig, self._config)

    @property
    def requires_authentication(self) -> bool:
        """Returns whether the container registry requires authentication.

        Returns:
            `True` if the container registry requires authentication,
            `False` otherwise.
        """
        return bool(self.config.authentication_secret)

    @property
    def credentials(self) -> Optional[Tuple[str, str]]:
        """Username and password to authenticate with this container registry.

        Returns:
            Tuple with username and password if this container registry
            requires authentication, `None` otherwise.
        """
        secret = self.get_typed_authentication_secret(
            expected_schema_type=BasicAuthSecretSchema
        )
        if secret:
            return secret.username, secret.password

        connector = self.get_connector()
        if connector:
            from zenml.service_connectors.docker_service_connector import (
                DockerServiceConnector,
            )

            if isinstance(connector, DockerServiceConnector):
                return (
                    connector.config.username.get_secret_value(),
                    connector.config.password.get_secret_value(),
                )

        return None

    @property
    def docker_client(self) -> "DockerClient":
        """Returns a Docker client for this container registry.

        Returns:
            The Docker client.

        Raises:
            RuntimeError: If the connector does not return a Docker client.
        """
        from docker.client import DockerClient

        # Refresh the client also if the connector has expired
        if self._docker_client and not self.connector_has_expired():
            return self._docker_client

        connector = self.get_connector()
        if connector:
            client = connector.connect()
            if not isinstance(client, DockerClient):
                raise RuntimeError(
                    f"Expected a DockerClient while trying to use the "
                    f"linked connector, but got {type(client)}."
                )
            self._docker_client = client
        else:
            self._docker_client = (
                docker_utils._try_get_docker_client_from_env()
            )

            credentials = self.credentials
            if credentials:
                username, password = credentials
                self._docker_client.login(
                    username=username,
                    password=password,
                    registry=self.config.uri,
                    reauth=True,
                )

        return self._docker_client

    def is_valid_image_name_for_registry(self, image_name: str) -> bool:
        """Check if the image name is valid for the container registry.

        Args:
            image_name: The name of the image.

        Returns:
            `True` if the image name is valid for the container registry,
            `False` otherwise.
        """
        # Remove prefixes to make sure this logic also works for DockerHub
        image_name = image_name.removeprefix("index.docker.io/")
        image_name = image_name.removeprefix("docker.io/")

        registry_uri = self.config.uri.removeprefix("index.docker.io/")
        registry_uri = registry_uri.removeprefix("docker.io/")

        return image_name.startswith(registry_uri)

    def prepare_image_push(self, image_name: str) -> None:
        """Preparation before an image gets pushed.

        Subclasses can overwrite this to do any necessary checks or
        preparations before an image gets pushed.

        Args:
            image_name: Name of the docker image that will be pushed.
        """

    def push_image(self, image_name: str) -> str:
        """Pushes a docker image.

        Args:
            image_name: Name of the docker image that will be pushed.

        Returns:
            The Docker repository digest of the pushed image.

        Raises:
            ValueError: If the image name is not associated with this
                container registry.
        """
        if not self.is_valid_image_name_for_registry(image_name):
            raise ValueError(
                f"Docker image `{image_name}` does not belong to container "
                f"registry `{self.config.uri}`."
            )

        self.prepare_image_push(image_name)
        return docker_utils.push_image(
            image_name, docker_client=self.docker_client
        )

    def get_image_repo_digest(self, image_name: str) -> Optional[str]:
        """Get the repository digest of an image.

        Args:
            image_name: The name of the image.

        Returns:
            The repository digest of the image.
        """
        if not self.is_valid_image_name_for_registry(image_name):
            return None

        try:
            metadata = self.docker_client.images.get_registry_data(image_name)
        except Exception:
            return None

        return cast(str, metadata.id.split(":")[-1])


class BaseContainerRegistryFlavor(Flavor):
    """Base flavor for container registries."""

    @property
    def type(self) -> StackComponentType:
        """Returns the flavor type.

        Returns:
            The flavor type.
        """
        return StackComponentType.CONTAINER_REGISTRY

    @property
    def service_connector_requirements(
        self,
    ) -> Optional[ServiceConnectorRequirements]:
        """Service connector resource requirements for service connectors.

        Specifies resource requirements that are used to filter the available
        service connector types that are compatible with this flavor.

        Returns:
            Requirements for compatible service connectors, if a service
            connector is required for this flavor.
        """
        return ServiceConnectorRequirements(
            resource_type=DOCKER_REGISTRY_RESOURCE_TYPE,
            resource_id_attr="uri",
        )

    @property
    def config_class(self) -> Type[BaseContainerRegistryConfig]:
        """Config class for this flavor.

        Returns:
            The config class.
        """
        return BaseContainerRegistryConfig

    @property
    def implementation_class(self) -> Type[BaseContainerRegistry]:
        """Implementation class.

        Returns:
            The implementation class.
        """
        return BaseContainerRegistry
