#  Copyright (c) ZenML GmbH 2024. 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.
"""ZenML login credentials store support."""

import os
from datetime import timedelta
from typing import Dict, List, Optional, Tuple, Union, cast

from zenml.config.global_config import GlobalConfiguration
from zenml.constants import (
    ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
    handle_bool_env_var,
)
from zenml.io import fileio
from zenml.logger import get_logger
from zenml.login.credentials import APIToken, ServerCredentials, ServerType
from zenml.login.pro.constants import ZENML_PRO_API_URL
from zenml.login.pro.workspace.models import WorkspaceRead
from zenml.models import OAuthTokenResponse, ServerModel
from zenml.utils import yaml_utils
from zenml.utils.singleton import SingletonMetaClass
from zenml.utils.time_utils import utc_now_tz_aware

logger = get_logger(__name__)

CREDENTIALS_STORE_FILENAME = "credentials.yaml"

# How long to keep tokens in the credentials store after their expiration date before evicting them
TOKEN_STORE_EVICTION_TIME = 60 * 60 * 24 * 7  # 1 week


class CredentialsStore(metaclass=SingletonMetaClass):
    """Login credentials store.

    This is a singleton object that maintains a cache of all API tokens and API
    keys that are configured for the ZenML servers that the client connects to
    throughout its lifetime.

    The cache is persistent and it is backed by a `credentials.yaml` YAML file
    kept in the global configuration location. The Credentials Store cache is
    populated mainly in the following ways:

        1. when the user runs `zenml login` to authenticate to a ZenML Pro
        server, the ZenML Pro API token fetched from the web login flow is
        stored in the Credentials Store.
        2. when the user runs `zenml login` to authenticate to a ZenML Pro
        server using an API key, the ZenML Pro API key is stored in the
        Credentials Store.
        3. when the user runs `zenml login` to authenticate to a regular ZenML
        server with the web login flow, the ZenML server API token fetched
        through the web login flow is stored in the Credentials Store
        4. when the user runs `zenml login` to authenticate to any ZenML server
        using an API key, the API key is stored in the Credentials Store
        5. when the REST zen store is initialized, it starts up not yet
        authenticated. Then, if/when it needs to authenticate or re-authenticate
        to the remote server, it will use whatever form of credentials it finds
        in the Credentials Store, in order of priority:

            * if it finds an API token that is not expired (e.g. authorized
            device API tokens fetched through the web login flow or short-lived
            session API tokens fetched through some other means of
            authentication), it will use that to authenticate
            * for ZenML servers that use an API key to authenticate, it will use
            that to fetch a short-lived ZenML Pro server API token that it also
            stores in the Credentials Store
            * for ZenML Pro servers:
                * if a valid ZenML Pro API token is found, it exchanges it into
                a short lived ZenML Pro server API token
                * otherwise, if a ZenML Pro API key is found, it will use that
                to fetch a short-lived ZenML Pro server API token that it also
                stores in the Credentials Store

    Alongside credentials, the Credentials Store is also used to store
    additional server information:
        * ZenML Pro workspace information populated by the `zenml login` command
        * ZenML server information populated by the REST zen store by fetching
        the server's information endpoint after authenticating

    """

    credentials: Dict[str, ServerCredentials]
    last_modified_time: Optional[float] = None

    def __init__(self) -> None:
        """Initializes the login credentials store with values loaded from the credentials YAML file.

        CredentialsStore is a singleton class: only one instance can exist.
        Calling this constructor multiple times will always yield the same
        instance.
        """
        self.credentials = {}
        self._load_credentials()

    @classmethod
    def reset_instance(
        cls, store: Optional["CredentialsStore"] = None
    ) -> None:
        """Reset the singleton instance of the CredentialsStore.

        Args:
            store: Optional instance of the CredentialsStore to set as the
                singleton instance. If None, a new instance will be created.
        """
        current_store = cls.get_instance()
        if current_store is not None and current_store is not store:
            # Delete the credentials file from disk if it exists, otherwise
            # the credentials will be reloaded from the file when the new
            # instance is created and this call will have no effect
            current_store._delete_credentials_file()

        cls._clear(store)  # type: ignore[arg-type]
        if store:
            store._save_credentials()

    @classmethod
    def get_instance(cls) -> Optional["CredentialsStore"]:
        """Get the singleton instance of the CredentialsStore.

        Returns:
            The singleton instance of the CredentialsStore.
        """
        return cast(CredentialsStore, cls._instance())

    @property
    def _credentials_file(self) -> str:
        """Path to the file where the credentials are stored.

        Returns:
            The path to the file where the credentials are stored.
        """
        config_path = GlobalConfiguration().config_directory
        return os.path.join(config_path, CREDENTIALS_STORE_FILENAME)

    def _load_credentials(self) -> None:
        """Load the credentials from the YAML file if it exists."""
        if handle_bool_env_var(ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING):
            return

        credentials_file = self._credentials_file
        credentials_store = {}

        if fileio.exists(credentials_file):
            try:
                credentials_store = yaml_utils.read_yaml(credentials_file)
            except Exception as e:
                logger.error(
                    f"Failed to load credentials file {credentials_file}: {e}. "
                )
            self.last_modified_time = os.path.getmtime(credentials_file)

        if credentials_store is None:
            # This can happen for example if the config file is empty
            credentials_store = {}
        elif not isinstance(credentials_store, dict):
            logger.warning(
                f"The credentials file {credentials_file} is corrupted. "
                "Creating a new credentials file."
            )
            credentials_store = {}

        self.credentials = {}
        for server_url, token_data in credentials_store.items():
            try:
                self.credentials[server_url] = ServerCredentials(**token_data)
            except ValueError as e:
                logger.warning(
                    f"Failed to load credentials for {server_url}: {e}. "
                    "Ignoring this token."
                )

    def _save_credentials(self) -> None:
        """Dump the current credentials store to the YAML file."""
        if handle_bool_env_var(ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING):
            return
        credentials_file = self._credentials_file
        credentials_store = {
            server_url: credential.model_dump(
                mode="json", exclude_none=True, exclude_unset=True
            )
            for server_url, credential in self.credentials.items()
            # Evict tokens that have expired past the eviction time
            # and have no API key or username/password to fall back on
            if credential.api_key
            or credential.username
            and credential.password is not None
            or credential.api_token
            and (
                not credential.api_token.expires_at
                or credential.api_token.expires_at
                + timedelta(seconds=TOKEN_STORE_EVICTION_TIME)
                > utc_now_tz_aware()
            )
        }
        yaml_utils.write_yaml(credentials_file, credentials_store)
        self.last_modified_time = os.path.getmtime(credentials_file)

    def _delete_credentials_file(self) -> None:
        """Delete the credentials file."""
        if handle_bool_env_var(ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING):
            return
        credentials_file = self._credentials_file
        if fileio.exists(credentials_file):
            fileio.remove(credentials_file)
            self.last_modified_time = None

    def check_and_reload_from_file(self) -> None:
        """Check if the credentials file has been modified and reload it if necessary."""
        if handle_bool_env_var(ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING):
            return
        if not self.last_modified_time:
            return
        credentials_file = self._credentials_file
        try:
            last_modified_time = os.path.getmtime(credentials_file)
        except FileNotFoundError:
            # The credentials file has been deleted
            self.last_modified_time = None
            return
        if last_modified_time != self.last_modified_time:
            self._load_credentials()

    def get_password(
        self, server_url: str
    ) -> Tuple[Optional[str], Optional[str]]:
        """Retrieve the username and password from the credentials store for a specific server URL.

        Args:
            server_url: The server URL for which to retrieve the username and
                password.

        Returns:
            The stored username and password if they exist, None otherwise.
        """
        self.check_and_reload_from_file()
        credential = self.credentials.get(server_url)
        if credential:
            return credential.username, credential.password
        return None, None

    def get_api_key(self, server_url: str) -> Optional[str]:
        """Retrieve an API key from the credentials store for a specific server URL.

        Args:
            server_url: The server URL for which to retrieve the API key.

        Returns:
            The stored API key if it exists, None otherwise.
        """
        self.check_and_reload_from_file()
        credential = self.credentials.get(server_url)
        if credential:
            return credential.api_key
        return None

    def get_token(
        self, server_url: str, allow_expired: bool = False
    ) -> Optional[APIToken]:
        """Retrieve a valid token from the credentials store for a specific server URL.

        Args:
            server_url: The server URL for which to retrieve the token.
            allow_expired: Whether to allow expired tokens to be returned. The
                default behavior is to return None if a token does exist but is
                expired.

        Returns:
            The stored token if it exists and is not expired, None otherwise.
        """
        self.check_and_reload_from_file()
        credential = self.credentials.get(server_url)
        if credential:
            token = credential.api_token
            if token and (not token.expired or allow_expired):
                return token
        return None

    def get_credentials(self, server_url: str) -> Optional[ServerCredentials]:
        """Retrieve the credentials for a specific server URL.

        Args:
            server_url: The server URL for which to retrieve the credentials.

        Returns:
            The stored credentials if they exist, None otherwise.
        """
        self.check_and_reload_from_file()
        return self.credentials.get(server_url)

    def has_valid_credentials(self, server_url: str) -> bool:
        """Check if valid credentials exist for a specific server URL.

        Args:
            server_url: The server URL for which to check the authentication.

        Returns:
            True if the credentials store contains credentials that can be used
            to login to the given server URL, False otherwise.
        """
        credentials = self.get_credentials(server_url)
        if credentials and (
            credentials.has_valid_token or credentials.can_refresh_token
        ):
            return True
        return False

    def get_pro_api_key(self, pro_api_url: str) -> Optional[str]:
        """Retrieve an API key from the credentials store for a ZenML Pro API server.

        Args:
            pro_api_url: The URL of the ZenML Pro API server.

        Returns:
            The stored API key if it exists, None otherwise.
        """
        credential = self.get_pro_credentials(pro_api_url)
        if credential and credential.api_key:
            return credential.api_key
        return None

    def get_pro_token(
        self, pro_api_url: str, allow_expired: bool = False
    ) -> Optional[APIToken]:
        """Retrieve a valid token from the credentials store for a ZenML Pro API server.

        Args:
            pro_api_url: The URL of the ZenML Pro API server.
            allow_expired: Whether to allow expired tokens to be returned. The
                default behavior is to return None if a token does exist but is
                expired.

        Returns:
            The stored token if it exists and is not expired, None otherwise.
        """
        credential = self.get_pro_credentials(pro_api_url)
        if credential:
            token = credential.api_token
            if token and (not token.expired or allow_expired):
                return token
        return None

    def get_pro_credentials(
        self, pro_api_url: str
    ) -> Optional[ServerCredentials]:
        """Retrieve valid credentials from the credentials store for a ZenML Pro API server.

        Args:
            pro_api_url: The URL of the ZenML Pro API server.

        Returns:
            The stored credentials if they exist, None otherwise.
        """
        credential = self.get_credentials(pro_api_url)
        if credential and credential.type == ServerType.PRO_API:
            return credential
        return None

    def has_valid_pro_credentials(self, pro_api_url: str) -> bool:
        """Check if valid credentials exist for a ZenML Pro API server.

        Args:
            pro_api_url: The URL of the ZenML Pro API server.

        Returns:
            True if the credentials store contains credentials that can be used
            to login to the given ZenML Pro API server, False otherwise.
        """
        credentials = self.get_pro_credentials(pro_api_url)
        if credentials and (
            credentials.has_valid_token or credentials.can_refresh_token
        ):
            return True
        return False

    def clear_pro_credentials(self, pro_api_url: str) -> None:
        """Delete the token from the store for a ZenML Pro API server.

        Args:
            pro_api_url: The URL of the ZenML Pro API server.
        """
        self.clear_token(pro_api_url)

    def clear_all_credentials(self) -> None:
        """Delete all stored credentials."""
        self.credentials = {}
        self._save_credentials()

    def can_login(self, server_url: str) -> bool:
        """Check if credentials to login to the given server exist.

        Args:
            server_url: The server URL for which to check the authentication.

        Returns:
            True if the credentials store contains credentials that can be used
            to login to the given server URL, False otherwise.
        """
        credentials = self.get_credentials(server_url)
        if not credentials:
            return False

        if credentials.can_refresh_token:
            return True

        if credentials.type == ServerType.PRO:
            pro_api_url = credentials.pro_api_url or ZENML_PRO_API_URL
            pro_credentials = self.get_pro_credentials(pro_api_url)
            if pro_credentials and (
                pro_credentials.has_valid_token
                or pro_credentials.can_refresh_token
            ):
                return True

        return False

    def set_api_key(
        self,
        server_url: str,
        api_key: str,
        is_zenml_pro: bool = False,
    ) -> None:
        """Store an API key in the credentials store for a specific server URL.

        If an API token or a password is already stored for the server URL, they
        will be replaced by the API key.

        Args:
            server_url: The server URL for which the token is to be stored.
            api_key: The API key to store.
            is_zenml_pro: Whether the API key is for the ZenML Pro API.
        """
        self.check_and_reload_from_file()
        credential = self.credentials.get(server_url)
        if credential and credential.api_key != api_key:
            # Reset the API token if a new or updated API key is set, because
            # the current token might have been issued for a different account
            credential.api_token = None
            credential.api_key = api_key
            credential.username = None
            credential.password = None
            if is_zenml_pro:
                credential.pro_api_url = server_url
        else:
            self.credentials[server_url] = ServerCredentials(
                url=server_url, api_key=api_key, pro_api_url=server_url
            )

        self._save_credentials()

    def set_password(
        self,
        server_url: str,
        username: str,
        password: str,
    ) -> None:
        """Store a username and password in the credentials store for a specific server URL.

        If an API token is already stored for the server URL, it will be
        replaced by the username and password.

        Args:
            server_url: The server URL for which the token is to be stored.
            username: The username to store.
            password: The password to store.
        """
        self.check_and_reload_from_file()
        credential = self.credentials.get(server_url)
        if credential and (
            credential.username != username or credential.password != password
        ):
            # Reset the API token if a new or updated password is set, because
            # the current token might have been issued for a different account
            credential.api_token = None
            credential.username = username
            credential.password = password
            credential.api_key = None
        else:
            self.credentials[server_url] = ServerCredentials(
                url=server_url, username=username, password=password
            )

        self._save_credentials()

    def set_token(
        self,
        server_url: str,
        token_response: OAuthTokenResponse,
        is_zenml_pro: bool = False,
    ) -> APIToken:
        """Store an API token received from an OAuth2 server.

        Args:
            server_url: The server URL for which the token is to be stored.
            token_response: Token response received from an OAuth2 server.
            is_zenml_pro: Whether the token is for a ZenML Pro server.

        Returns:
            APIToken: The stored token.
        """
        self.check_and_reload_from_file()
        if token_response.expires_in:
            expires_at = utc_now_tz_aware() + timedelta(
                seconds=token_response.expires_in
            )
            # Best practice to calculate the leeway depending on the token
            # expiration time:
            #
            # - for short-lived tokens (less than 1 hour), use a fixed leeway of
            # a few seconds (e.g., 30 seconds)
            # - for longer-lived tokens (e.g., 1 hour or more), use a
            # percentage-based leeway of 5-10%
            if token_response.expires_in < 3600:
                leeway = 30
            else:
                leeway = token_response.expires_in // 20
        else:
            expires_at = None
            leeway = None

        api_token = APIToken(
            access_token=token_response.access_token,
            expires_in=token_response.expires_in,
            expires_at=expires_at,
            leeway=leeway,
            device_id=token_response.device_id,
            device_metadata=token_response.device_metadata,
        )

        credential = self.credentials.get(server_url)
        if credential:
            credential.api_token = api_token
        else:
            credential = self.credentials[server_url] = ServerCredentials(
                url=server_url, api_token=api_token
            )

        if is_zenml_pro:
            # This is how we encode that the token is for a ZenML Pro server
            credential.pro_api_url = server_url

        self._save_credentials()

        return api_token

    def set_bare_token(
        self,
        server_url: str,
        token: str,
    ) -> APIToken:
        """Store a bare API token.

        Args:
            server_url: The server URL for which the token is to be stored.
            token: The token to store.

        Returns:
            APIToken: The stored token.
        """
        self.check_and_reload_from_file()

        api_token = APIToken(
            access_token=token,
        )

        credential = self.credentials.get(server_url)
        if credential:
            credential.api_token = api_token
        else:
            self.credentials[server_url] = ServerCredentials(
                url=server_url, api_token=api_token
            )

        self._save_credentials()

        return api_token

    def update_server_info(
        self,
        server_url: str,
        server_info: Union[ServerModel, WorkspaceRead],
    ) -> None:
        """Update the server information stored for a specific server URL.

        Args:
            server_url: The server URL for which the server information is to be
                updated.
            server_info: Updated server information.
        """
        self.check_and_reload_from_file()

        credential = self.credentials.get(server_url)
        if not credential:
            # No credentials stored for this server URL, nothing to update
            return

        credential.update_server_info(server_info)
        self._save_credentials()

    def clear_token(self, server_url: str) -> None:
        """Delete a token from the store for a specific server URL.

        Args:
            server_url: The server URL for which to delete the token.
        """
        self.check_and_reload_from_file()

        if server_url in self.credentials:
            credential = self.credentials[server_url]
            if (
                not credential.api_key
                and not credential.username
                and not credential.password is not None
            ):
                # Only delete the credential entry if there is no API key or
                # username/password to fall back on
                del self.credentials[server_url]
            else:
                credential.api_token = None
            self._save_credentials()

    def clear_credentials(self, server_url: str) -> None:
        """Delete all credentials from the store for a specific server URL.

        Args:
            server_url: The server URL for which to delete the credentials.
        """
        self.check_and_reload_from_file()

        if server_url in self.credentials:
            del self.credentials[server_url]
            self._save_credentials()

    def list_credentials(
        self, type: Optional[ServerType] = None
    ) -> List[ServerCredentials]:
        """Get all credentials stored in the credentials store.

        Args:
            type: Optional server type to filter the credentials by.

        Returns:
            A list of all credentials stored in the credentials store.
        """
        self.check_and_reload_from_file()

        credentials = list(self.credentials.values())

        if type is not None:
            credentials = [c for c in credentials if c and c.type == type]

        return credentials


def get_credentials_store() -> CredentialsStore:
    """Get the global credentials store instance.

    Returns:
        The global credentials store instance.
    """
    return CredentialsStore()
