"""
Utilities for handling auth in jobbergate-cli.
"""

import webbrowser
from time import sleep
from typing import Dict, Optional, cast

from jose import jwt
from jose.exceptions import ExpiredSignatureError
from loguru import logger
from pydantic import ValidationError
from rich.console import Console

from jobbergate_cli.config import settings
from jobbergate_cli.exceptions import Abort, JobbergateCliError
from jobbergate_cli.render import terminal_message
from jobbergate_cli.requests import make_request
from jobbergate_cli.schemas import DeviceCodeData, IdentityData, JobbergateContext, Persona, TokenSet
from jobbergate_cli.text_tools import copy_to_clipboard, unwrap
from jobbergate_cli.time_loop import TimeLoop


def validate_token_and_extract_identity(token_set: TokenSet) -> IdentityData:
    """
    Validate the access_token from a TokenSet and extract the user's identity data.

    Validations:
        * Checks if access_token is not empty.
        * Checks timestamp on the access token.
        * Checks that the client_id is present
        * Checks that email is present

    Reports an error in the logs and to the user if there is an issue with the access_token.
    """
    logger.debug("Validating access token")

    token_file_is_empty = not token_set.access_token
    if token_file_is_empty:
        logger.debug("Access token file exists but it is empty")
        raise Abort(
            """
            Access token file exists but it is empty.

            Please try logging in again.
            """,
            subject="Empty access token file",
            support=True,
            log_message="Empty access token file",
            sentry_context=dict(access_token=dict(access_token=token_set.access_token)),
        )

    try:
        token_data = jwt.decode(
            token_set.access_token,
            None,
            options=dict(
                verify_signature=False,
                verify_aud=False,
                verify_exp=True,
            ),
        )
    except ExpiredSignatureError:
        raise  # Will be handled in calling context
    except Exception as err:
        raise Abort(
            """
            There was an unknown error while validating the access token.

            Please try logging in again.
            """,
            subject="Invalid access token",
            support=True,
            log_message=f"Unknown error while validating access access token: {err}",
            sentry_context=dict(access_token=dict(access_token=token_set.access_token)),
            original_error=err,
        )

    logger.debug("Extracting identity data from the access token")
    try:
        identity = IdentityData(
            email=token_data.get("email"),
            client_id=token_data.get("azp"),
        )
    except ValidationError as err:
        raise Abort(
            """
            There was an error extracting the user's identity from the access token.

            Please try logging in again.
            """,
            subject="Missing user data",
            support=True,
            log_message=f"Token data could not be extracted to identity: {err}",
            original_error=err,
        )
    return identity


def load_tokens_from_cache() -> TokenSet:
    """
    Loads an access token (and a refresh token if one exists) from the cache.
    """

    # Make static type checkers happy
    assert settings.JOBBERGATE_ACCESS_TOKEN_PATH is not None
    assert settings.JOBBERGATE_REFRESH_TOKEN_PATH is not None

    Abort.require_condition(
        settings.JOBBERGATE_ACCESS_TOKEN_PATH.exists(),
        "Please login with your auth token first using the `jobbergate login` command",
        raise_kwargs=dict(subject="You need to login"),
    )

    logger.debug("Retrieving access token from cache")
    token_set = TokenSet(access_token=settings.JOBBERGATE_ACCESS_TOKEN_PATH.read_text())

    if settings.JOBBERGATE_REFRESH_TOKEN_PATH.exists():
        logger.debug("Retrieving refresh token from cache")
        token_set.refresh_token = settings.JOBBERGATE_REFRESH_TOKEN_PATH.read_text()

    return token_set


def save_tokens_to_cache(token_set: TokenSet):
    """
    Saves tokens from a token_set to the cache.
    """

    # Make static type checkers happy
    assert settings.JOBBERGATE_ACCESS_TOKEN_PATH is not None
    assert settings.JOBBERGATE_REFRESH_TOKEN_PATH is not None

    logger.debug(f"Caching access token at {settings.JOBBERGATE_ACCESS_TOKEN_PATH}")
    settings.JOBBERGATE_ACCESS_TOKEN_PATH.write_text(token_set.access_token)
    settings.JOBBERGATE_ACCESS_TOKEN_PATH.chmod(0o600)

    if token_set.refresh_token is not None:
        logger.debug(f"Caching refresh token at {settings.JOBBERGATE_REFRESH_TOKEN_PATH}")
        settings.JOBBERGATE_REFRESH_TOKEN_PATH.write_text(token_set.refresh_token)
        settings.JOBBERGATE_REFRESH_TOKEN_PATH.chmod(0o600)


def clear_token_cache():
    """
    Clears the token cache.
    """
    logger.debug("Clearing cached tokens")

    logger.debug(f"Removing access token at {settings.JOBBERGATE_ACCESS_TOKEN_PATH}")
    if settings.JOBBERGATE_ACCESS_TOKEN_PATH.exists():
        settings.JOBBERGATE_ACCESS_TOKEN_PATH.unlink()

    logger.debug(f"Removing refresh token at {settings.JOBBERGATE_REFRESH_TOKEN_PATH}")
    if settings.JOBBERGATE_REFRESH_TOKEN_PATH.exists():
        settings.JOBBERGATE_REFRESH_TOKEN_PATH.unlink()


def init_persona(ctx: JobbergateContext, token_set: Optional[TokenSet] = None):
    """
    Initializes the "persona" which contains the tokens and email address for a user.

    Retrieves the access token for the user from the cache.

    Token is retrieved from the cache, validated, and user email is extracted.

    If the access token is expired, a new one will be acquired via the cached refresh token (if there is one).

    Saves token_set to cache.

    Returns the persona.
    """
    if token_set is None:
        token_set = load_tokens_from_cache()

    try:
        identity_data = validate_token_and_extract_identity(token_set)
    except ExpiredSignatureError:
        Abort.require_condition(
            token_set.refresh_token is not None,
            "The auth token is expired. Please retrieve a new and log in again.",
            raise_kwargs=dict(
                subject="Expired access token",
                support=True,
            ),
        )

        logger.debug("The access token is expired. Attempting to refresh token")
        refresh_access_token(ctx, token_set)
        identity_data = validate_token_and_extract_identity(token_set)

    logger.debug(f"Persona created with identity_data: {identity_data}")

    save_tokens_to_cache(token_set)

    return Persona(
        token_set=token_set,
        identity_data=identity_data,
    )


def refresh_access_token(ctx: JobbergateContext, token_set: TokenSet):
    """
    Attempt to fetch a new access token given a refresh token in a token_set.

    Sets the access token in-place.

    If refresh fails, notify the user that they need to log in again.
    """
    url = f"https://{settings.OIDC_DOMAIN}/protocol/openid-connect/token"
    logger.debug(f"Requesting refreshed access token from {url}")

    JobbergateCliError.require_condition(
        ctx.client is not None,
        "Attempted to refresh with a null client. This should not happen",
    )

    # Make static type-checkers happy
    assert ctx.client is not None

    refreshed_token_set = cast(
        TokenSet,
        make_request(
            # Can this even work? this client should be for the armada api...
            ctx.client,
            "/protocol/openid-connect/token",
            "POST",
            abort_message="The auth token could not be refreshed. Please try logging in again.",
            abort_subject="EXPIRED ACCESS TOKEN",
            support=True,
            response_model_cls=TokenSet,
            data=dict(
                client_id=settings.OIDC_CLIENT_ID,
                audience=settings.OIDC_AUDIENCE,
                grant_type="refresh_token",
                refresh_token=token_set.refresh_token,
            ),
        ),
    )

    token_set.access_token = refreshed_token_set.access_token
    if refreshed_token_set.refresh_token is not None:
        token_set.refresh_token = refreshed_token_set.refresh_token


def open_on_browser(url: str) -> bool:
    """Open the url on the browser using webbrowser."""
    try:
        browser = webbrowser.get()
        browser.open(url)
        return True
    except Exception as e:
        logger.warning(f"Couldn't open login url on browser due to -- {str(e)}")
        return False


def show_login_message(verification_uri: str):
    """Show a message to the user with a link to the auth provider to login."""
    console = Console()
    EXTRA_CHARS = 7  # for indentation and panel borders

    kwargs = {}

    if open_on_browser(verification_uri):
        kwargs["footer"] = "The output was opened on your browser"
    elif copy_to_clipboard(verification_uri):
        kwargs["footer"] = "The output was copied to your clipboard"

    if console.width >= len(verification_uri) + EXTRA_CHARS:
        _show_login_standard_screen(verification_uri, **kwargs)
    else:
        _show_login_narrow_screen(verification_uri, console, **kwargs)


def _show_login_narrow_screen(verification_uri: str, console: Console, **kwargs):
    """Print the link out of the panel to make it easier to copy."""
    terminal_message(
        f"""
        To complete login, please open the link bellow in a browser.

        Waiting up to {settings.OIDC_MAX_POLL_TIME / 60} minutes for you to complete the process...
        """,
        subject="Waiting for login",
        **kwargs,
    )
    console.print(verification_uri, overflow="ignore", no_wrap=True, crop=False)
    console.print()


def _show_login_standard_screen(verification_uri: str, **kwargs):
    """Print a rich panel with a link to the auth provider to login."""
    terminal_message(
        f"""
        To complete login, please open the following link in a browser:

          {verification_uri}

        Waiting up to {settings.OIDC_MAX_POLL_TIME / 60} minutes for you to complete the process...
        """,
        subject="Waiting for login",
        **kwargs,
    )


def fetch_auth_tokens(ctx: JobbergateContext) -> TokenSet:
    """
    Fetch an access token (and possibly a refresh token) from Auth0.

    Prints out a URL for the user to use to authenticate and polls the token endpoint to fetch it when
    the browser-based process finishes
    """
    # Make static type-checkers happy
    assert ctx.client is not None

    device_code_data = cast(
        DeviceCodeData,
        make_request(
            ctx.client,
            "/protocol/openid-connect/auth/device",
            "POST",
            expected_status=200,
            abort_message="There was a problem retrieving a device verification code from the auth provider",
            abort_subject="COULD NOT RETRIEVE TOKEN",
            support=True,
            response_model_cls=DeviceCodeData,
            data=dict(
                client_id=settings.OIDC_CLIENT_ID,
                grant_type="client_credentials",
                audience=settings.OIDC_AUDIENCE,
            ),
        ),
    )
    show_login_message(device_code_data.verification_uri_complete)

    for tick in TimeLoop(
        settings.OIDC_MAX_POLL_TIME,
        message="Waiting for web login",
    ):
        response_data = cast(
            Dict,
            make_request(
                ctx.client,
                "/protocol/openid-connect/token",
                "POST",
                abort_message="There was a problem retrieving a device verification code from the auth provider",
                abort_subject="COULD NOT FETCH ACCESS TOKEN",
                support=True,
                data=dict(
                    grant_type="urn:ietf:params:oauth:grant-type:device_code",
                    device_code=device_code_data.device_code,
                    client_id=settings.OIDC_CLIENT_ID,
                ),
            ),
        )
        if "error" in response_data:
            if response_data["error"] == "authorization_pending":
                logger.debug(f"Token fetch attempt #{tick.counter} failed")
                sleep(device_code_data.interval)
            else:
                # TODO: Test this failure condition
                raise Abort(
                    unwrap(
                        """
                        There was a problem retrieving a device verification code from the auth provider:
                        Unexpected failure retrieving access token.
                        """
                    ),
                    subject="Unexpected error",
                    support=True,
                    log_message=f"Unexpected error response: {response_data}",
                )
        else:
            return TokenSet(**response_data)

    raise Abort(
        "Login process was not completed in time. Please try again.",
        subject="Timed out",
        log_message="Timed out while waiting for user to complete login",
    )
