#  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.
"""Base secrets store class used for all secrets stores that use a service connector."""

import json
from abc import abstractmethod
from threading import Lock
from typing import (
    Any,
    ClassVar,
    Dict,
    Optional,
    Type,
)

from pydantic import Field, model_validator

from zenml.config.secrets_store_config import SecretsStoreConfiguration
from zenml.logger import get_logger
from zenml.models import (
    ServiceConnectorConfiguration,
    ServiceConnectorRequest,
)
from zenml.service_connectors.service_connector import ServiceConnector
from zenml.service_connectors.service_connector_registry import (
    service_connector_registry,
)
from zenml.utils.pydantic_utils import before_validator_handler
from zenml.zen_stores.secrets_stores.base_secrets_store import (
    BaseSecretsStore,
)

logger = get_logger(__name__)


class ServiceConnectorSecretsStoreConfiguration(SecretsStoreConfiguration):
    """Base configuration for secrets stores that use a service connector.

    Attributes:
        auth_method: The service connector authentication method to use.
        auth_config: The service connector authentication configuration.
    """

    auth_method: str
    auth_config: Dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="before")
    @classmethod
    @before_validator_handler
    def validate_auth_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
        """Convert the authentication configuration if given in JSON format.

        Args:
            data: The configuration values.

        Returns:
            The validated configuration values.

        Raises:
            ValueError: If the authentication configuration is not a valid
                JSON object.
        """
        if isinstance(data.get("auth_config"), str):
            try:
                data["auth_config"] = json.loads(data["auth_config"])
            except json.JSONDecodeError as e:
                raise ValueError(
                    f"The authentication configuration is not a valid JSON "
                    f"object: {e}"
                )
        return data


class ServiceConnectorSecretsStore(BaseSecretsStore):
    """Base secrets store class for service connector-based secrets stores.

    All secrets store implementations that use a Service Connector to
    authenticate and connect to the secrets store back-end should inherit from
    this class and:

    * implement the `_initialize_client_from_connector` method
    * use a configuration class that inherits from
    `ServiceConnectorSecretsStoreConfiguration`
    * set the `SERVICE_CONNECTOR_TYPE` to the service connector type used
    to connect to the secrets store back-end
    * set the `SERVICE_CONNECTOR_RESOURCE_TYPE` to the resource type used
    to connect to the secrets store back-end
    """

    config: ServiceConnectorSecretsStoreConfiguration
    CONFIG_TYPE: ClassVar[Type[ServiceConnectorSecretsStoreConfiguration]]
    SERVICE_CONNECTOR_TYPE: ClassVar[str]
    SERVICE_CONNECTOR_RESOURCE_TYPE: ClassVar[str]

    _connector: Optional[ServiceConnector] = None
    _client: Optional[Any] = None
    _lock: Optional[Lock] = None

    def _initialize(self) -> None:
        """Initialize the secrets store."""
        self._lock = Lock()
        # Initialize the client early, just to catch any configuration or
        # authentication errors early, before the Secrets Store is used.
        _ = self.client

    def _get_client(self) -> Any:
        """Initialize and return the secrets store API client.

        Returns:
            The secrets store API client object.
        """
        if self._connector is not None:
            # If the client connector expires, we'll try to get a new one.
            if self._connector.has_expired():
                self._connector = None
                self._client = None

        if self._connector is None:
            # Initialize a base service connector with the credentials from
            # the configuration.
            request = ServiceConnectorRequest(
                name="secrets-store",
                connector_type=self.SERVICE_CONNECTOR_TYPE,
                resource_types=[self.SERVICE_CONNECTOR_RESOURCE_TYPE],
                auth_method=self.config.auth_method,
                configuration=ServiceConnectorConfiguration(
                    **self.config.auth_config
                ),
            )
            base_connector = service_connector_registry.instantiate_connector(
                model=request
            )
            # Set the `allow_implicit_auth_methods` flag to `True` to allow
            # implicit authentication methods to be used even when not globally
            # enabled via the `ZENML_ENABLE_IMPLICIT_AUTH_METHODS` environment
            # variable.
            base_connector.allow_implicit_auth_methods = True
            self._connector = base_connector.get_connector_client()

        if self._client is None:
            # Use the connector to get a client object.
            client = self._connector.connect(
                # Don't verify again because we already did that when we
                # initialized the connector.
                verify=False
            )

            self._client = self._initialize_client_from_connector(client)
        return self._client

    @property
    def lock(self) -> Lock:
        """Get the lock used to treat the client initialization as a critical section.

        Returns:
            The lock instance.
        """
        assert self._lock is not None
        return self._lock

    @property
    def client(self) -> Any:
        """Get the secrets store API client.

        Returns:
            The secrets store API client instance.
        """
        # Multiple API calls can be made to this secrets store in the context
        # of different threads. We want to make sure that we only initialize
        # the client once, and then reuse it. We have to use a lock to treat
        # this method as a critical section.
        with self.lock:
            return self._get_client()

    @abstractmethod
    def _initialize_client_from_connector(self, client: Any) -> Any:
        """Initialize the client from the service connector.

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

        Returns:
            The initialized client instance.
        """
