"""
Utilities for handling authentication in the Jobbergate system.
"""

import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Iterable

from loguru import logger
import pendulum
from pydantic import BaseModel

from jobbergate_core.auth.exceptions import AuthenticationError
from jobbergate_core.auth.token import Token, TokenType
from jobbergate_core.tools.requests import Client, RequestHandler


class DeviceCodeData(BaseModel):
    """
    A model representing the data that is returned from the OIDC provider's device code endpoint.
    """

    verification_uri_complete: str
    interval: int
    device_code: str
    expires_in: float


class TokenInformation(BaseModel):
    access_token: str | None = None
    refresh_token: str | None = None


class IdentityData(BaseModel):
    """
    A model representing the identifying data for a user from an auth token.
    """

    email: str
    client_id: str
    organization_id: str | None = None


@dataclass
class TimedIterator:
    """
    An iterator that runs for a given time interval, yielding the current iteration number.
    """

    total: int
    step: int

    def __post_init__(self):
        now = pendulum.now()
        interval = pendulum.interval(now, now.add(seconds=self.total))
        self.range = interval.range("seconds", self.step)

    def __iter__(self):
        for i, date_time in enumerate(self.range):
            time_to_sleep = date_time.diff(pendulum.now()).in_seconds()
            time.sleep(time_to_sleep)
            yield i

    def __len__(self):
        return self.total // self.step + 1


def print_login_url(device_code_data: DeviceCodeData):
    """Basic way to handle the login url."""
    print(f"Login Here: {device_code_data.verification_uri_complete}")


@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_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_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_client_id: str = "default"
    login_client_secret: str | None = None
    login_url_handler: Callable[[DeviceCodeData], None] = print_login_url
    login_sequence_handler: Callable[[Iterable], Iterable] = lambda i: i

    _client: Client = field(init=False, repr=False)
    _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)
        self._client = Client(base_url=self.login_domain, headers={"content-type": "application/x-www-form-urlencoded"})

    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.
        """
        if self._access_token.is_valid():
            return self._access_token.bearer_token

        logger.debug("Acquiring access token")

        for procedure_name in (
            "load_from_cache",
            "refresh_tokens",
            "get_access_from_secret",
            "login",
        ):
            procedure = getattr(self, procedure_name)
            try:
                procedure()
            except AuthenticationError as err:
                logger.debug("{} failed due to: {}", procedure_name, str(err))
            if self._access_token.is_valid():
                return self._access_token.bearer_token
        raise AuthenticationError("Unable to acquire the access token, all attempts failed")

    def get_access_from_secret(self) -> None:
        """
        Get the access token using the client secret.
        """
        logger.debug("Getting access token from client secret")
        AuthenticationError.require_condition(self.login_client_secret is not None, message="Client secret is unset")

        with AuthenticationError.handle_errors("Failed to get access token from client secret"):
            data = (
                RequestHandler(
                    client=self._client,
                    url_path="/protocol/openid-connect/token",
                    method="POST",
                    request_kwargs={
                        "data": {
                            "client_id": self.login_client_id,
                            "client_secret": self.login_client_secret,
                            "grant_type": "client_credentials",
                        }
                    },
                    sensitive_keys={"access_token", "client_secret"},
                )
                .raise_for_status()
                .to_model(TokenInformation)
            )
            self._update_tokens_from_info(data)

    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())
        with AuthenticationError.handle_errors("Failed to load tokens from cache"):
            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())
        with AuthenticationError.handle_errors("Failed to save tokens to cache"):
            self.cache_directory.mkdir(parents=True, exist_ok=True)
            if self._access_token.is_valid():
                self._access_token.save_to_cache()
            if self._refresh_token.is_valid():
                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")
        with AuthenticationError.handle_errors("Failed to logout 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")
        with AuthenticationError.handle_errors("Failed to login to Jobbergate"):
            login_info = self._get_device_code()
            token_info = self._wait_for_login_confirmation(login_info)
            self._update_tokens_from_info(token_info)
        logger.success("Login completed")

    def _wait_for_login_confirmation(self, device_code_data: DeviceCodeData) -> TokenInformation:
        self.login_url_handler(device_code_data)
        for counter in self.login_sequence_handler(
            TimedIterator(int(device_code_data.expires_in), device_code_data.interval)
        ):
            request_handler = RequestHandler(
                client=self._client,
                url_path="/protocol/openid-connect/token",
                method="POST",
                request_kwargs={
                    "data": {
                        "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                        "device_code": device_code_data.device_code,
                        "client_id": self.login_client_id,
                    }
                },
                sensitive_keys={"access_token", "cookie", "device_code", "refresh_token"},
            ).check_status_code(200, 400)

            if request_handler.response.is_success:
                return request_handler.to_model(TokenInformation)

            logger.debug(
                "Login not completed completed on attempt #{}, waiting {} seconds",
                counter + 1,
                device_code_data.interval,
            )

        raise AuthenticationError("Login process was not completed in time. Please try again.")

    def _get_device_code(self) -> DeviceCodeData:
        return (
            RequestHandler(
                client=self._client,
                url_path="/protocol/openid-connect/auth/device",
                method="POST",
                request_kwargs={
                    "data": {
                        "client_id": self.login_client_id,
                        "grant_type": "client_credentials",
                    }
                },
                sensitive_keys={"device_code"},
            )
            .raise_for_status()
            .check_status_code(200)
            .to_model(DeviceCodeData)
        )

    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 AuthenticationError("Session can no be refreshed since the refresh token is unavailable")
        if self._refresh_token.is_expired():
            raise AuthenticationError("Session can no be refreshed since the refresh token is expired")

        with AuthenticationError.handle_errors("Failed to refresh the session"):
            token_info = self._get_refresh_token()
            self._update_tokens_from_info(token_info)

        logger.success("Tokens refreshed successfully")

    def _get_refresh_token(self) -> TokenInformation:
        return (
            RequestHandler(
                client=self._client,
                url_path="/protocol/openid-connect/token",
                method="POST",
                request_kwargs={
                    "data": {
                        "client_id": self.login_client_id,
                        "grant_type": "refresh_token",
                        "refresh_token": self._refresh_token.content,
                    },
                },
                sensitive_keys={"access_token", "refresh_token"},
            )
            .raise_for_status()
            .to_model(TokenInformation)
        )

    def _update_tokens_from_info(self, token_information: TokenInformation):
        """
        Update the tokens with the new content.
        """
        if token_information.access_token:
            self._access_token = self._access_token.replace(content=token_information.access_token)
        if token_information.refresh_token:
            self._refresh_token = self._refresh_token.replace(content=token_information.refresh_token)
        self.save_to_cache()

    def get_identity_data(self) -> IdentityData:
        if not self._access_token.is_valid():
            self.acquire_access()
        token_data = self._access_token.data
        email = AuthenticationError.enforce_defined(
            token_data.get("email"),
            "Could not retrieve user email from token",
        )
        client_id = AuthenticationError.enforce_defined(
            token_data.get("azp"), "Could not retrieve client_id from token"
        )
        org_dict = token_data.get("organization", {})
        AuthenticationError.require_condition(
            len(org_dict) <= 1,
            message="More than one organization id found in token payload",
        )
        organization_id = set(org_dict.keys()).pop() if org_dict else None
        return IdentityData(
            email=email,
            client_id=client_id,
            organization_id=organization_id,
        )
