# Copyright 2023-2024 Broadcom
# SPDX-License-Identifier: Apache-2.0
import json
import logging
import time
from typing import Optional

from requests import post
from requests.auth import HTTPBasicAuth
from requests_oauthlib import OAuth2Session
from vdk.plugin.control_api_auth.auth_config import AuthConfig
from vdk.plugin.control_api_auth.auth_config import CredentialsCache
from vdk.plugin.control_api_auth.auth_config import LocalFolderCredentialsCache
from vdk.plugin.control_api_auth.auth_exception import VDKAuthException
from vdk.plugin.control_api_auth.auth_request_values import AuthRequestValues
from vdk.plugin.control_api_auth.login_types import LoginTypes

log = logging.getLogger(__name__)


class AuthenticationCache:
    """
    Class used to keep the fields in the vdk-cred.json file from which
    authorization and authentication tokens can be taken. Used only when
    caching is enabled.
    """

    def __init__(
        self,
        authorization_url=None,
        refresh_token=None,
        access_token=None,
        access_token_expiration_time=0,
        client_id=None,
        client_secret=None,
        api_token=None,
        api_token_authorization_url=None,
        auth_type=None,
        **ignored_kwargs,
    ):
        """
        :param authorization_url:
            Token or Authorization URI used to exchange grant for access token
        :param refresh_token:
            Used in refresh token grant type flow;
            See https://datatracker.ietf.org/doc/html/rfc6749#section-1.5
        :param access_token:
            Cached access token used during requests;
            See https://datatracker.ietf.org/doc/html/rfc6749#section-1.4
        :param access_token_expiration_time:
            When the access token expires
            (epoch seconds -- same as time.time() + "expires_in" value returned
            by the authorization server that generates the access token.)
        :param client_id:
            The client identifier;
            See https://datatracker.ietf.org/doc/html/rfc6749#section-2.3.1
        :param api_token:
            Token used to grant access to the API,
            including exchange for access token.
        :param api_token_authorization_url:
            The URL which exchanges api_token for access token.
        """
        #
        self.authorization_url = authorization_url
        self.refresh_token = refresh_token
        self.access_token = access_token
        self.access_token_expiration_time = access_token_expiration_time
        self.client_id = client_id
        self.client_secret = client_secret
        self.api_token = api_token
        self.api_token_authorization_url = api_token_authorization_url
        self.auth_type = auth_type

    def __eq__(self, other):
        return self.__dict__ == other.__dict__


class AuthenticationCacheSerDe:
    """Class used for cache serialization and deserialization."""

    @staticmethod
    def serialize(cache) -> str:
        return json.dumps(cache.__dict__)

    @staticmethod
    def deserialize(content: str) -> AuthenticationCache:
        if content:
            return AuthenticationCache(**json.loads(content))
        else:
            return AuthenticationCache()


class BaseAuth:
    REFRESH_TOKEN_GRANT_TYPE = "refresh_token"  # nosec

    def __init__(
        self,
        conf: CredentialsCache = LocalFolderCredentialsCache(),
        vdk_config=AuthConfig(),
    ):
        self._conf = conf
        self.__vdk_config = vdk_config
        self._cache = self.__load_cache()

    def __load_cache(self):
        content = self._conf.read_credentials()
        return AuthenticationCacheSerDe.deserialize(content)

    def __update_cache(self):
        self._conf.save_credentials(AuthenticationCacheSerDe.serialize(self._cache))

    def __exchange_api_for_access_token(self):
        try:
            log.debug(
                f"Refresh API token against {self._cache.api_token_authorization_url} "
            )
            client = OAuth2Session()
            token_response = client.refresh_token(
                self._cache.api_token_authorization_url,
                refresh_token=self._cache.api_token,
            )
            log.debug(
                f"Token response received: "
                f"token_type: {token_response.get('token_type', None)}; "
                f"expires_in: {token_response.get('expires_in', None)}; "
                f"scope: {token_response.get('scope', None)}"
            )
        except Exception as e:
            raise VDKAuthException(
                what="Failed to login",
                why=f"Authorization server at {self._cache.api_token_authorization_url} returned error: {str(e)}",
                consequence="Your credentials are not refreshed and VDK control operations that require authentication "
                "will not work.",
                countermeasure="Check error message and follow instructions in it. Check your network connectivity."
                " Check if you can access the Oauth2 server. Retry to login a few times again."
                " Contact SRE Team operating Oauth2 server for help if nothing else works.",
            ) from e
        # Response per spec in https://tools.ietf.org/html/rfc6750#section-4
        # see also https://www.oauth.com/oauth2-servers/access-tokens/access-token-response
        self.update_access_token(
            token_response[AuthRequestValues.ACCESS_TOKEN_KEY.value]
        )
        self.update_access_token_expiration_time(
            time.time()
            + int(token_response.get(AuthRequestValues.EXPIRATION_TIME_KEY.value, "0"))
        )

    def __exchange_refresh_for_access_token(self):
        basic_auth = HTTPBasicAuth(self._cache.client_id, self._cache.client_secret)
        headers = {
            AuthRequestValues.CONTENT_TYPE_HEADER.value: AuthRequestValues.CONTENT_TYPE_URLENCODED.value,
        }
        data = (
            f"grant_type={self.REFRESH_TOKEN_GRANT_TYPE}&"
            + f"refresh_token={self._cache.refresh_token}"
        )
        log.debug(
            f"Refresh access token against {self._cache.authorization_url} grant_type={self.REFRESH_TOKEN_GRANT_TYPE}"
        )
        response = post(
            self._cache.authorization_url, data=data, headers=headers, auth=basic_auth
        )
        if response.status_code >= 400:
            raise VDKAuthException(
                what="Failed to refresh access token",
                why="Authorization server returned status code: {response.status_code}, error response: {response.text}",
                consequence="Subsequent API calls may fail.",
                countermeasure="Try to log in again the usual way. "
                "If using the CLI, you can login using the 'vdk login' command. "
                "If using a notebook, use the Login Button.",
            )

        json_data = json.loads(response.text)
        log.debug(
            f"Refresh access token response received: status: {response.status_code}"
        )
        self.update_access_token(json_data[AuthRequestValues.ACCESS_TOKEN_KEY.value])
        self.update_access_token_expiration_time(
            time.time() + int(json_data[AuthRequestValues.EXPIRATION_TIME_KEY.value])
        )

    def read_access_token(self) -> Optional[str]:
        """
        Read access token from cache or fetch it from Authorization server.
        If not available in cache it will get it using provided configuration during VDK CLI login to fetch it.
        If it detects that token is about to expire it will try to refresh it.
        :return: the access token or None if it cannot detect any credentials.
        """
        if (
            not self._cache.access_token
            or self._cache.access_token_expiration_time < time.time() + 60
        ):
            log.debug("Acquire access token (it's either expired or missing) ...")
            self.acquire_and_cache_access_token()
        return self._cache.access_token

    def read_cached_access_token_only(self) -> Optional[str]:
        """
        Read access token from credentials cache and will return it if available.
        Will not check if it is valid (for example not expired) but return it directly.
        Will not try to fetch it form Authorization provider if it is not available. Instead, will simply return None
        """
        return self._cache.access_token

    def acquire_and_cache_access_token(self):
        """
        Acquires and caches access token
        """
        log.debug(f"Using auth type {self._cache.auth_type} to acquire access token")
        if (
            self._cache.auth_type == LoginTypes.API_TOKEN.value
            and self._configured_api_token()
        ):
            self.__exchange_api_for_access_token()
        elif (
            self._cache.auth_type == LoginTypes.CREDENTIALS.value
            and self._configured_refresh_token()
        ):
            self.__exchange_refresh_for_access_token()
        elif (
            self.__vdk_config.api_token
            and self.__vdk_config.api_token_authorization_url
        ):
            self.update_api_token(self.__vdk_config.api_token)
            self.update_api_token_authorization_url(
                self.__vdk_config.api_token_authorization_url
            )
            self.__exchange_api_for_access_token()
        else:
            log.debug(
                "No authentication mechanism found. Will not cache access token."
                "If Control Service authentication is enabled, API calls will fail."
            )

    def update_api_token_authorization_url(self, api_token_authorization_url):
        """
        Updates and caches authorization URL
        """
        self._cache.api_token_authorization_url = api_token_authorization_url
        self.__update_cache()

    def update_api_token(self, api_token):
        """
        Updates and caches refresh token
        """
        self._cache.api_token = api_token
        self.__update_cache()

    def update_refresh_token(self, refresh_token):
        """
        Updates and caches refresh token
        """
        self._cache.refresh_token = refresh_token
        self.__update_cache()

    def update_client_id(self, client_id):
        """
        Updates and caches refresh token
        """
        self._cache.client_id = client_id
        self.__update_cache()

    def update_client_secret(self, client_secret):
        """
        Updates and caches refresh token
        """
        self._cache.client_secret = client_secret
        self.__update_cache()

    def update_access_token(self, access_token):
        """
        Updates and caches refresh token
        """
        self._cache.access_token = access_token
        self.__update_cache()

    def update_access_token_expiration_time(self, access_token_expiration_time):
        """
        Updates and caches refresh token
        """
        self._cache.access_token_expiration_time = access_token_expiration_time
        self.__update_cache()

    def update_oauth2_authorization_url(self, authorization_url):
        self._cache.authorization_url = authorization_url
        self.__update_cache()

    def update_auth_type(self, auth_type):
        self._cache.auth_type = auth_type
        self.__update_cache()

    def _configured_api_token(self):
        return self._cache.api_token and self._cache.api_token_authorization_url

    def _configured_refresh_token(self):
        return (
            self._cache.refresh_token
            and self._cache.authorization_url
            and self._cache.client_id
            and self._cache.client_secret
        )
