"""Implementation of Google Artifact Registry for wandb launch."""

import logging
from typing import Optional, Tuple

import google.auth  # type: ignore
import google.cloud.artifactregistry  # type: ignore

from wandb.sdk.launch.errors import LaunchError
from wandb.sdk.launch.utils import (
    GCP_ARTIFACT_REGISTRY_URI_REGEX,
    event_loop_thread_exec,
)
from wandb.util import get_module

from .abstract import AbstractRegistry

_logger = logging.getLogger(__name__)

google = get_module(  # noqa: F811
    "google",
    required="The google package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
)
google.auth = get_module(  # noqa: F811
    "google.auth",
    required="The google-auth package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
)

google.cloud.artifactregistry = get_module(  # noqa: F811
    "google.cloud.artifactregistry",
    required="The google-cloud-artifactregistry package is required to use launch with Google. Please install it with `pip install wandb[launch]`.",
)


class GoogleArtifactRegistry(AbstractRegistry):
    """Google Artifact Registry helper for interacting with the registry.

    This helper should be constructed from either a uri or a repository,
    project, and optional image-name. If constructed from a uri, the uri
    must be of the form REGION-docker.pkg.dev/PROJECT/REPOSITORY/[IMAGE_NAME],
    with an optional https:// preceding.
    """

    def __init__(
        self,
        uri: Optional[str] = None,
        repository: Optional[str] = None,
        image_name: Optional[str] = None,
        project: Optional[str] = None,
        region: Optional[str] = None,
    ) -> None:
        """Initialize the Google Artifact Registry.

        Either uri or repository and image_name must be provided. Project and
        region are optional, and will be inferred from the uri if provided, or
        from the default credentials if not.

        Arguments:
            uri (optional): The uri of the repository.
            repository (optional): The repository name.
            image_name (optional): The image name.
            project (optional): The GCP project name.
            region (optional): The GCP region name.

        Raises:
            LaunchError: If verify is True and the container registry or its
                environment have not been properly configured. Or if the environment
                is not an instance of GcpEnvironment.
        """
        _logger.info(
            f"Initializing Google Artifact Registry with repository {repository} "
            f"and image name {image_name}"
        )

        if uri is not None:
            self.uri = uri
            # Raise an error if any other kwargs were provided in addition to uri.
            if any([repository, image_name, project, region]):
                raise LaunchError(
                    "The Google Artifact Registry must be specified with either "
                    "the uri key or the repository, image-name, project and region "
                    "keys, but not both."
                )
            match = GCP_ARTIFACT_REGISTRY_URI_REGEX.match(self.uri)
            if not match:
                raise LaunchError(
                    f"The Google Artifact Registry uri {self.uri} is invalid. "
                    "Please provide a uri of the form "
                    "REGION-docker.pkg.dev/PROJECT/REPOSITORY/IMAGE_NAME."
                )
            self.project = match.group("project")
            self.region = match.group("region")
            self.repository = match.group("repository")
            self.image_name = match.group("image_name")
        else:
            if any(x is None for x in (repository, region, image_name)):
                raise LaunchError(
                    "The Google Artifact Registry must be specified with either "
                    "the uri key or the repository, image-name, project and region "
                    "keys."
                )
            self.project = project
            self.region = region
            self.repository = repository
            self.image_name = image_name
            self.uri = f"{self.region}-docker.pkg.dev/{self.project}/{self.repository}/{self.image_name}"

        _missing_kwarg_msg = (
            "The Google Artifact Registry is missing the {} kwarg. "
            "Please specify it by name or as part of the uri argument."
        )
        if not self.region:
            raise LaunchError(_missing_kwarg_msg.format("region"))
        if not self.repository:
            raise LaunchError(_missing_kwarg_msg.format("repository"))
        if not self.image_name:
            raise LaunchError(_missing_kwarg_msg.format("image-name"))
        # Try to load default project from the default credentials.
        self.credentials, project = google.auth.default()
        self.project = self.project or project
        self.credentials.refresh(google.auth.transport.requests.Request())

    @classmethod
    def from_config(
        cls,
        config: dict,
    ) -> "GoogleArtifactRegistry":
        """Create a Google Artifact Registry from a config.

        Arguments:
            config: A dictionary containing the following keys:
                repository: The repository name.
                image-name: The image name.
            environment: A GcpEnvironment configured for access to this registry.

        Returns:
            A GoogleArtifactRegistry.
        """
        # TODO: Replace this with pydantic.
        acceptable_keys = {
            "uri",
            "type",
            "repository",
            "image-name",
            "region",
            "project",
        }
        unacceptable_keys = set(config.keys()) - acceptable_keys
        if unacceptable_keys:
            raise LaunchError(
                f"The Google Artifact Registry config contains unacceptable keys: "
                f"{unacceptable_keys}. Please remove these keys. The acceptable "
                f"keys are: {acceptable_keys}."
            )
        return cls(
            uri=config.get("uri"),
            repository=config.get("repository"),
            image_name=config.get("image-name"),
            project=config.get("project"),
            region=config.get("region"),
        )

    async def get_username_password(self) -> Tuple[str, str]:
        """Get the username and password for the registry.

        Returns:
            A tuple of the username and password.
        """
        if not self.credentials.token:
            self.credentials.refresh(google.auth.transport.requests.Request())
        return "oauth2accesstoken", self.credentials.token

    async def get_repo_uri(self) -> str:
        """Get the URI for the given repository.

        Arguments:
            repo_name: The repository name.

        Returns:
            The repository URI.
        """
        return (
            f"{self.region}-docker.pkg.dev/"
            f"{self.project}/{self.repository}/{self.image_name}"
        )

    async def check_image_exists(self, image_uri: str) -> bool:
        """Check if the image exists.

        Arguments:
            image_uri: The image URI.

        Returns:
            True if the image exists, False otherwise.
        """
        _logger.info(f"Checking if image {image_uri} exists")
        repo_uri, tag = image_uri.split(":")
        self_repo_uri = await self.get_repo_uri()
        if repo_uri != self_repo_uri:
            raise LaunchError(
                f"The image {image_uri} does not match to the image uri "
                f"repository {self.uri}."
            )
        parent = f"projects/{self.project}/locations/{self.region}/repositories/{self.repository}"
        artifact_registry_client = event_loop_thread_exec(
            google.cloud.artifactregistry.ArtifactRegistryClient
        )
        client = await artifact_registry_client(credentials=self.credentials)
        list_images = event_loop_thread_exec(client.list_docker_images)
        try:
            for image in await list_images(request={"parent": parent}):
                if tag in image.tags:
                    return True
        except google.api_core.exceptions.NotFound as e:  # type: ignore[attr-defined]
            raise LaunchError(
                f"The Google Artifact Registry repository {self.repository} "
                f"does not exist. Please create it or modify your registry configuration."
            ) from e
        return False
