"""
Utilities for handling authentication in the Jobbergate system.
"""
import time
from collections import namedtuple
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict

import httpx
from loguru import logger

from jobbergate_core.auth.exceptions import AuthenticationError, TokenError
from jobbergate_core.auth.token import Token, TokenType


_LoginInformation = namedtuple(
    "_LoginInformation",
    ["verification_url", "wait_interval", "device_code", "expires_at"],
)


@dataclass
class JobbergateAuthHandler:
    """
    High-level class used to manage authentication to requests to the Jobbergate-API

    After an instance of this class is created, it can be used to authenticate requests
    from both `requests`_  and `httpx`_ packages by passing it to the ``auth``
    parameter on the request (see examples below).

    It just works out of the box. Behind the scenes, this procedure calls the
    :meth:`JobbergateAuthHandler.acquire_access` method to load the available tokens from the cache directory,
    it tries to refresh them if they are expired, or provides an URL to the user to login on the system.

    Notice all steps above are also available individually as public methods, allowing a fine control for
    advanced users.

    .. _requests: https://requests.readthedocs.io/en/latest/
    .. _httpx: https://www.python-httpx.org/

    Arguments:
        cache_directory: Directory to be used for the caching tokens.
        login_domain: Domain used for the login.
        login_audience: Audience of the login.
        login_client_id: Client ID used for login.

    Note:
        These values depend on the identity provider used for authentication. Consult your system
        administrator or contact Omnivector support <support@omnivector.solutions> for further assistance.

    Note:
        This class can interoperate with the tokens generated by the ``jobbergate-cli`` package, as long
        as they are stored in the same cache directory.

    Examples:

        The following example shows how to use the :meth:`JobbergateAuthHandler`
        class to authenticate a request:

        >>> from pathlib import Path
        >>> import requests
        >>> from jobbergate_core import JobbergateAuthHandler
        >>> jobbergate_auth = JobbergateAuthHandler(
        ...     cache_directory=Path("."),
        ...     login_domain="http://keycloak.local:8080/realms/jobbergate-local",
        ...     login_audience="https://local.omnivector.solutions",
        ...     login_client_id="cli",
        ... )
        >>> jobbergate_base_url = "http://localhost:8000/jobbergate"
        >>> response = requests.get(
        ...     f"{jobbergate_base_url}/applications",
        ...     auth=jobbergate_auth # this is the important part
        )
        Login Here: http://keycloak.local:8080/realms/jobbergate-local/device?user_code=LMVJ-XOLG
        >>> response.raise_for_status()
        >>> print(f"response = {response.json()}")
    """

    cache_directory: Path
    login_domain: str
    login_audience: str
    login_client_id: str = "default"
    _access_token: Token = field(init=False, repr=False)
    _refresh_token: Token = field(init=False, repr=False)

    def __post_init__(self):
        self._access_token = Token(cache_directory=self.cache_directory, label=TokenType.ACCESS.value)
        self._refresh_token = Token(cache_directory=self.cache_directory, label=TokenType.REFRESH.value)

    def __call__(self, request):
        """
        This internal method allows the integration with the ``requests`` library.

        It is called automatically when the instance is passed to the ``auth`` parameter,
        see the examples in the class docstring.

        It adds the ``Authorization`` header to the request with the access token.
        """
        logger.debug("Authenticating request")
        request.headers["Authorization"] = self.acquire_access()
        return request

    def acquire_access(self) -> str:
        """
        High-level method to acquire a valid access token.

        This method will attempt, in order:

        * Use the internal access token from the instance
        * Load the tokens from the cache directory (see :meth:`JobbergateAuthHandler.load_from_cache`)
        * If the access token is unavailable or expired, refresh both tokens
          using the refresh token grant type (see :meth:`JobbergateAuthHandler.refresh_tokens`)
        * If the refresh token is unavailable or expired, login to generate both tokens
          (see :meth:`JobbergateAuthHandler.login`)

        Returns:
            The bearer access token.

        Raises:
            AuthenticationError: If all of the steps above fail to acquire a valid access token.
        """
        logger.debug("Acquiring access token")

        for procedure_name in (None, "load_from_cache", "refresh_tokens", "login"):
            if procedure_name:
                procedure = getattr(self, procedure_name)
                try:
                    procedure()
                except TokenError:
                    logger.debug("{} failed, moving to the next procedure", procedure_name)
            if self._access_token.is_valid():
                return self._access_token.bearer_token
        raise AuthenticationError("Unable to acquire the access token")

    def load_from_cache(self) -> None:
        """
        Load the tokens that are available at the cache directory.
        """
        logger.debug("Loading tokens from cache directory: {}", self.cache_directory.as_posix())
        self._access_token = self._access_token.load_from_cache()
        self._refresh_token = self._refresh_token.load_from_cache()

    def save_to_cache(self) -> None:
        """
        Save the tokens to the cache directory.

        Note:
            This method will create the cache directory if it does not exist.
        """
        logger.debug(
            "Saving tokens to cache directory: {}",
            self.cache_directory.as_posix(),
        )
        self.cache_directory.mkdir(parents=True, exist_ok=True)
        self._access_token.save_to_cache()
        self._refresh_token.save_to_cache()

    def logout(self) -> None:
        """
        Logout from Jobbergate by clearing the loaded tokens and their cache on the disk.
        """
        logger.debug("Logging out from Jobbergate")
        self._access_token = self._access_token.replace(content="")
        self._access_token.clear_cache()
        self._refresh_token = self._refresh_token.replace(content="")
        self._refresh_token.clear_cache()

    def login(self) -> None:
        """
        Login to Jobbergate.

        An URL will be printed to the console, the user must open it in a browser and provide
        their access credentials.

        After the login is completed, the tokens will be saved to the cache directory.
        """
        logger.debug("Preparing to login to Jobbergate")
        login_info = self._get_login_information()
        response = self._wait_for_login_confirmation(login_info)
        self._process_tokens_from_response(response)
        logger.success("Login completed")

    def _wait_for_login_confirmation(self, login_info: _LoginInformation) -> httpx.Response:
        print(f"Login Here: {login_info.verification_url}")
        while True:
            AuthenticationError.require_condition(
                login_info.expires_at > time.time(), "Login expired, please try again"
            )
            response = self._get_login_confirmation(login_info)
            try:
                response.raise_for_status()
                break
            except httpx.HTTPStatusError:
                logger.debug(
                    "    Login not completed yet, waiting {} seconds",
                    login_info.wait_interval,
                )
                time.sleep(login_info.wait_interval)
        logger.debug("Preparing to login to Jobbergate")
        return response

    def _get_login_confirmation(self, login_info: _LoginInformation) -> httpx.Response:
        response = httpx.post(
            f"{self.login_domain}/protocol/openid-connect/token",
            data=dict(
                grant_type="urn:ietf:params:oauth:grant-type:device_code",
                device_code=login_info.device_code,
                client_id=self.login_client_id,
            ),
        )

        return response

    def _get_login_information(self) -> _LoginInformation:
        with AuthenticationError.handle_errors(
            "Unexpected error while fetching the tokens",
        ):
            response = httpx.post(
                f"{self.login_domain}/protocol/openid-connect/auth/device",
                data=dict(
                    client_id=self.login_client_id,
                    grant_type="client_credentials",
                    audience=self.login_audience,
                ),
            )
            response.raise_for_status()

        device_code_data = response.json()
        with AuthenticationError.handle_errors(
            f"Error processing the request data after fetching the token, {device_code_data=}",
        ):
            verification_url = device_code_data["verification_uri_complete"]
            wait_interval = device_code_data["interval"]
            device_code = device_code_data["device_code"]
            expires_at = time.time() + device_code_data["expires_in"]
        return _LoginInformation(
            verification_url,
            wait_interval,
            device_code,
            expires_at,
        )

    def refresh_tokens(self) -> None:
        """
        Refresh the tokens.

        After the refresh operation is completed, the tokens will be saved to the cache directory.

        Raises:
            AuthenticationError: If the refresh token is missing or expired.
        """
        logger.debug("Preparing to refresh the tokens")

        if not self._refresh_token.content:
            raise TokenError("The refresh is unavailable, please login again")
        if self._refresh_token.is_expired():
            raise TokenError("Refresh token is expired, please login again")

        response = self._get_refresh_token()
        self._process_tokens_from_response(response)

        logger.success("Tokens refreshed successfully")

    def _get_refresh_token(self):
        with AuthenticationError.handle_errors(
            "Unexpected error while refreshing the tokens",
        ):
            response = httpx.post(
                f"{self.login_domain}/protocol/openid-connect/token",
                data=dict(
                    client_id=self.login_client_id,
                    audience=self.login_audience,
                    grant_type="refresh_token",
                    refresh_token=self._refresh_token.content,
                ),
            )
            response.raise_for_status()
        return response

    def _process_tokens_from_response(self, response):
        response_data = response.json()

        tokens_content = {t: response_data.get(f"{t.value}_token") for t in TokenType}
        AuthenticationError.require_condition(
            all(tokens_content.values()), "Not all tokens were included in the response"
        )
        self._update_tokens(tokens_content)
        self.save_to_cache()

    def _update_tokens(self, tokens_content: Dict[TokenType, str]):
        """
        Update the tokens with the new content.
        """
        access_token = tokens_content.get(TokenType.ACCESS, "")
        if access_token:
            self._access_token = self._access_token.replace(content=access_token)
        refresh_token = tokens_content.get(TokenType.REFRESH, "")
        if refresh_token:
            self._refresh_token = self._refresh_token.replace(content=refresh_token)
