#  Copyright (c) ZenML GmbH 2023. 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 GCP Secrets Store."""

import json
import os
from typing import (
    Any,
    ClassVar,
    Dict,
    Optional,
    Type,
    cast,
)
from uuid import UUID

from google.api_core import exceptions as google_exceptions
from google.cloud.secretmanager import SecretManagerServiceClient
from pydantic import ConfigDict, model_validator

from zenml.enums import (
    SecretsStoreType,
)
from zenml.integrations.gcp import (
    GCP_CONNECTOR_TYPE,
    GCP_RESOURCE_TYPE,
)
from zenml.integrations.gcp.service_connectors.gcp_service_connector import (
    GCPAuthenticationMethods,
)
from zenml.logger import get_logger
from zenml.utils.pydantic_utils import before_validator_handler
from zenml.zen_stores.secrets_stores.service_connector_secrets_store import (
    ServiceConnectorSecretsStore,
    ServiceConnectorSecretsStoreConfiguration,
)

logger = get_logger(__name__)


GCP_ZENML_SECRET_NAME_PREFIX = "zenml"
ZENML_GROUP_KEY = "zenml-group-key"
ZENML_GCP_SECRET_SCOPE_PATH_SEPARATOR = "-"


class GCPSecretsStoreConfiguration(ServiceConnectorSecretsStoreConfiguration):
    """GCP secrets store configuration.

    Attributes:
        type: The type of the store.
    """

    type: SecretsStoreType = SecretsStoreType.GCP

    @property
    def project_id(self) -> str:
        """Get the GCP project ID.

        Returns:
            The GCP project ID.

        Raises:
            ValueError: If the project ID is not set.
        """
        project_id = self.auth_config.get("project_id")
        if project_id:
            return str(project_id)

        raise ValueError("GCP `project_id` must be specified in auth_config.")

    @model_validator(mode="before")
    @classmethod
    @before_validator_handler
    def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
        """Populate the connector configuration from legacy attributes.

        Args:
            data: Dict representing user-specified runtime settings.

        Returns:
            Validated settings.
        """
        # Search for legacy attributes and populate the connector configuration
        # from them, if they exist.
        if data.get("project_id"):
            if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
                logger.warning(
                    "The `project_id` GCP secrets store attribute is "
                    "deprecated and will be removed in a future version of ZenML. "
                    "Please use the `auth_method` and `auth_config` attributes "
                    "instead. Using an implicit GCP authentication to access "
                    "the GCP Secrets Manager API."
                )
                data["auth_method"] = GCPAuthenticationMethods.IMPLICIT
                data["auth_config"] = dict(
                    project_id=data.get("project_id"),
                )
            else:
                logger.warning(
                    "The `project_id` GCP secrets store attribute and the "
                    "`GOOGLE_APPLICATION_CREDENTIALS` environment variable are "
                    "deprecated and will be removed in a future version of ZenML. "
                    "Please use the `auth_method` and `auth_config` attributes "
                    "instead."
                )
                data["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT
                data["auth_config"] = dict(
                    project_id=data.get("project_id"),
                )
                # Load the service account credentials from the file
                with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f:
                    data["auth_config"]["service_account_json"] = f.read()

        return data

    model_config = ConfigDict(extra="allow")


class GCPSecretsStore(ServiceConnectorSecretsStore):
    """Secrets store implementation that uses the GCP Secrets Manager API."""

    config: GCPSecretsStoreConfiguration
    TYPE: ClassVar[SecretsStoreType] = SecretsStoreType.GCP
    CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]] = (
        GCPSecretsStoreConfiguration
    )
    SERVICE_CONNECTOR_TYPE: ClassVar[str] = GCP_CONNECTOR_TYPE
    SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str] = GCP_RESOURCE_TYPE

    _client: Optional[SecretManagerServiceClient] = None

    @property
    def client(self) -> SecretManagerServiceClient:
        """Initialize and return the GCP Secrets Manager client.

        Returns:
            The GCP Secrets Manager client instance.
        """
        return cast(SecretManagerServiceClient, super().client)

    # ====================================
    # Secrets Store interface implementation
    # ====================================

    # --------------------------------
    # Initialization and configuration
    # --------------------------------

    def _initialize_client_from_connector(self, client: Any) -> Any:
        """Initialize the GCP Secrets Manager client from the service connector client.

        Args:
            client: The authenticated client object returned by the service
                connector.

        Returns:
            The GCP Secrets Manager client.
        """
        return SecretManagerServiceClient(credentials=client)

    # ------
    # Secrets
    # ------

    @property
    def parent_name(self) -> str:
        """Construct the GCP parent path to the secret manager.

        Returns:
            The parent path to the secret manager
        """
        return f"projects/{self.config.project_id}"

    def _get_gcp_secret_name(
        self,
        secret_id: UUID,
    ) -> str:
        """Get the GCP secret name for the given secret.

        The convention used for GCP secret names is to use the ZenML
        secret UUID prefixed with `zenml` as the AWS secret name,
        i.e. `zenml/<secret_uuid>`.

        Args:
            secret_id: The ZenML secret ID.

        Returns:
            The GCP secret name.
        """
        return f"{GCP_ZENML_SECRET_NAME_PREFIX}-{str(secret_id)}"

    def store_secret_values(
        self,
        secret_id: UUID,
        secret_values: Dict[str, str],
    ) -> None:
        """Store secret values for a new secret.

        Args:
            secret_id: ID of the secret.
            secret_values: Values for the secret.

        Raises:
            RuntimeError: if the GCP Secrets Manager API returns an unexpected
                error.
        """
        secret_value = json.dumps(secret_values)

        labels = self._get_secret_metadata(secret_id=secret_id)

        try:
            gcp_secret = self.client.create_secret(
                request={
                    "parent": self.parent_name,
                    "secret_id": self._get_gcp_secret_name(secret_id),
                    "secret": {
                        "replication": {"automatic": {}},
                        "labels": labels,
                    },
                }
            )

            logger.debug(f"Created empty GCP parent secret: {gcp_secret.name}")

            self.client.add_secret_version(
                request={
                    "parent": gcp_secret.name,
                    "payload": {"data": secret_value.encode()},
                }
            )

            logger.debug(f"Added value to GCP secret {gcp_secret.name}")
        except Exception as e:
            raise RuntimeError(f"Failed to create secret.: {str(e)}") from e

        logger.debug(f"Created GCP secret {gcp_secret.name}")

    def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
        """Get the secret values for an existing secret.

        Args:
            secret_id: ID of the secret.

        Returns:
            The secret values.

        Raises:
            KeyError: if no secret values for the given ID are stored in the
                secrets store.
            RuntimeError: if the GCP Secrets Manager API returns an unexpected
                error.
        """
        gcp_secret_name = self.client.secret_path(
            self.config.project_id,
            self._get_gcp_secret_name(secret_id=secret_id),
        )

        try:
            secret = self.client.get_secret(name=gcp_secret_name)
            secret_version_values = self.client.access_secret_version(
                name=f"{gcp_secret_name}/versions/latest"
            )
        except google_exceptions.NotFound as e:
            raise KeyError(
                f"Can't find the secret values for secret ID '{secret_id}' "
                f"in the secrets store back-end: {str(e)}"
            ) from e
        except Exception as e:
            raise RuntimeError(
                f"Error fetching secret with ID {secret_id} {e}"
            )

        # The GCP secret labels do not really behave like a dictionary: when
        # a key is not found, it does not raise a KeyError, but instead
        # returns an empty string. That's why we make this conversion.
        metadata = dict(secret.labels)

        # The _verify_secret_metadata method raises a KeyError if the
        # secret is not valid or does not belong to this server. Here we
        # simply pass the exception up the stack, as if the secret was not found
        # in the first place.
        self._verify_secret_metadata(
            secret_id=secret_id,
            metadata=metadata,
        )

        secret_values = json.loads(
            secret_version_values.payload.data.decode("UTF-8")
        )

        if not isinstance(secret_values, dict):
            raise RuntimeError(
                f"Google secret values for secret ID {gcp_secret_name} could "
                "not be decoded: expected a dictionary."
            )

        logger.debug(f"Fetched GCP secret: {gcp_secret_name}")

        return secret_values

    def update_secret_values(
        self,
        secret_id: UUID,
        secret_values: Dict[str, str],
    ) -> None:
        """Updates secret values for an existing secret.

        Args:
            secret_id: The ID of the secret to be updated.
            secret_values: The new secret values.

        Raises:
            RuntimeError: if the GCP Secrets Manager API returns an unexpected
                error.
        """
        gcp_secret_name = self.client.secret_path(
            self.config.project_id,
            self._get_gcp_secret_name(secret_id=secret_id),
        )

        # Convert the ZenML secret metadata to GCP labels
        metadata = self._get_secret_metadata(secret_id)

        try:
            # Update the secret metadata
            update_secret = {
                "name": gcp_secret_name,
                "labels": metadata,
            }
            update_mask = {"paths": ["labels"]}
            gcp_updated_secret = self.client.update_secret(
                request={
                    "secret": update_secret,
                    "update_mask": update_mask,
                }
            )
            # Add a new secret version
            secret_value = json.dumps(secret_values)
            self.client.add_secret_version(
                request={
                    "parent": gcp_updated_secret.name,
                    "payload": {"data": secret_value.encode()},
                }
            )
        except Exception as e:
            raise RuntimeError(f"Error updating secret: {e}") from e

        logger.debug(f"Updated GCP secret: {gcp_secret_name}")

    def delete_secret_values(self, secret_id: UUID) -> None:
        """Deletes secret values for an existing secret.

        Args:
            secret_id: The ID of the secret.

        Raises:
            KeyError: if no secret values for the given ID are stored in the
                secrets store.
            RuntimeError: if the GCP Secrets Manager API returns an unexpected
                error.
        """
        gcp_secret_name = self.client.secret_path(
            self.config.project_id,
            self._get_gcp_secret_name(secret_id=secret_id),
        )

        try:
            self.client.delete_secret(request={"name": gcp_secret_name})
        except google_exceptions.NotFound:
            raise KeyError(f"Secret with ID {secret_id} not found")
        except Exception as e:
            raise RuntimeError(f"Failed to delete secret: {str(e)}") from e

        logger.debug(f"Deleted GCP secret: {gcp_secret_name}")
