"""Log in to Weights & Biases.

This authenticates your machine to log data to your account.
"""

import enum
import os
from typing import Literal, Optional, Tuple

import click

import wandb
from wandb.errors import AuthenticationError, UsageError
from wandb.old.settings import Settings as OldSettings

from ..apis import InternalApi
from .internal.internal_api import Api
from .lib import apikey
from .wandb_settings import Settings


def _handle_host_wandb_setting(host: Optional[str], cloud: bool = False) -> None:
    """Write the host parameter to the global settings file.

    This takes the parameter from wandb.login or wandb login for use by the
    application's APIs.
    """
    _api = InternalApi()
    if host == "https://api.wandb.ai" or (host is None and cloud):
        _api.clear_setting("base_url", globally=True, persist=True)
        # To avoid writing an empty local settings file, we only clear if it exists
        if os.path.exists(OldSettings._local_path()):
            _api.clear_setting("base_url", persist=True)
    elif host:
        host = host.rstrip("/")
        # force relogin if host is specified
        _api.set_setting("base_url", host, globally=True, persist=True)


def login(
    anonymous: Optional[Literal["must", "allow", "never"]] = None,
    key: Optional[str] = None,
    relogin: Optional[bool] = None,
    host: Optional[str] = None,
    force: Optional[bool] = None,
    timeout: Optional[int] = None,
    verify: bool = False,
) -> bool:
    """Set up W&B login credentials.

    By default, this will only store credentials locally without
    verifying them with the W&B server. To verify credentials, pass
    `verify=True`.

    Args:
        anonymous: (string, optional) Can be "must", "allow", or "never".
            If set to "must", always log a user in anonymously. If set to
            "allow", only create an anonymous user if the user
            isn't already logged in. If set to "never", never log a
            user anonymously. Default set to "never".
        key: (string, optional) The API key to use.
        relogin: (bool, optional) If true, will re-prompt for API key.
        host: (string, optional) The host to connect to.
        force: (bool, optional) If true, will force a relogin.
        timeout: (int, optional) Number of seconds to wait for user input.
        verify: (bool) Verify the credentials with the W&B server.

    Returns:
        bool: if key is configured

    Raises:
        AuthenticationError - if api_key fails verification with the server
        UsageError - if api_key cannot be configured and no tty
    """
    _handle_host_wandb_setting(host)
    if wandb.setup()._settings._noop:
        return True

    configured = _login(
        anonymous=anonymous,
        key=key,
        relogin=relogin,
        host=host,
        force=force,
        timeout=timeout,
    )

    if verify:
        from . import wandb_setup

        singleton = wandb_setup._WandbSetup._instance
        assert singleton is not None
        viewer = singleton._server._viewer
        if not viewer:
            raise AuthenticationError(
                "API key verification failed. Make sure your API key is valid."
            )
    return True if configured else False


class ApiKeyStatus(enum.Enum):
    VALID = 1
    NOTTY = 2
    OFFLINE = 3
    DISABLED = 4


class _WandbLogin:
    def __init__(self):
        self._settings: Optional[Settings] = None
        self._backend = None
        self._silent: Optional[bool] = None
        self._entity: Optional[str] = None
        self._wl = None
        self._key = None
        self._relogin = None

    def setup(
        self,
        *,
        anonymous: Optional[Literal["allow", "must", "never"]] = None,
        key: Optional[str] = None,
        relogin: Optional[bool] = None,
        host: Optional[str] = None,
        force: Optional[bool] = None,
        timeout: Optional[int] = None,
    ) -> None:
        """Updates login-related settings on the global setup object."""
        self._relogin = relogin

        login_settings = {
            "anonymous": anonymous,
            "api_key": key,
            "base_url": host,
            "force": force,
            "login_timeout": timeout,
        }

        # make sure they are applied globally
        self._wl = wandb.setup(
            settings=wandb.Settings(
                **{k: v for k, v in login_settings.items() if v is not None}
            )
        )
        self._settings = self._wl.settings

    def is_apikey_configured(self) -> bool:
        """Returns whether an API key is set or can be inferred."""
        return apikey.api_key(settings=self._settings) is not None

    def should_use_identity_token(self):
        return self._settings.identity_token_file is not None

    def set_backend(self, backend):
        self._backend = backend

    def set_silent(self, silent: bool) -> None:
        self._silent = silent

    def set_entity(self, entity: str) -> None:
        self._entity = entity

    def login(self) -> bool:
        """Returns whether the user is logged in (i.e. an API key exists).

        If the user is logged in, this also prints an informational message.
        """
        apikey_configured = self.is_apikey_configured()
        if self._settings.relogin or self._relogin:
            apikey_configured = False
        if not apikey_configured:
            return False

        if not self._silent:
            self._print_logged_in_message()

        return apikey_configured

    def _print_logged_in_message(self) -> None:
        """Prints a message telling the user they are logged in."""
        username = self._wl._get_username()

        if username:
            # check to see if we got an entity from the setup call or from the user
            entity = self._entity or self._wl._get_entity()

            entity_str = ""
            # check if entity exist, valid (is part of a certain team) and different from the username
            if entity and entity in self._wl._get_teams() and entity != username:
                entity_str = f" ({click.style(entity, fg='yellow')})"

            login_state_str = f"Currently logged in as: {click.style(username, fg='yellow')}{entity_str}"
        else:
            login_state_str = "W&B API key is configured"

        login_info_str = (
            f"Use {click.style('`wandb login --relogin`', bold=True)} to force relogin"
        )
        wandb.termlog(
            f"{login_state_str}. {login_info_str}",
            repeat=False,
        )

    def configure_api_key(self, key: str) -> None:
        """Saves the API key and updates the the global setup object."""
        if self._settings._notebook and not self._settings.silent:
            wandb.termwarn(
                "If you're specifying your api key in code, ensure this "
                "code is not shared publicly.\nConsider setting the "
                "WANDB_API_KEY environment variable, or running "
                "`wandb login` from the command line."
            )
        apikey.write_key(self._settings, key)
        self.update_session(key)
        self._key = key

    def update_session(
        self,
        key: Optional[str],
        status: ApiKeyStatus = ApiKeyStatus.VALID,
    ) -> None:
        """Updates mode and API key settings on the global setup object.

        If we're online, this also pulls in user settings from the server.
        """
        login_settings = dict()
        if status == ApiKeyStatus.OFFLINE:
            login_settings = dict(mode="offline")
        elif status == ApiKeyStatus.DISABLED:
            login_settings = dict(mode="disabled")
        elif key:
            login_settings = dict(api_key=key)
        self._wl._settings.update_from_dict(login_settings)
        # Whenever the key changes, make sure to pull in user settings
        # from server.
        if not self._wl.settings._offline:
            self._wl._update_user_settings()

    def _prompt_api_key(self) -> Tuple[Optional[str], ApiKeyStatus]:
        api = Api(self._settings)
        while True:
            try:
                key = apikey.prompt_api_key(
                    self._settings,
                    api=api,
                    no_offline=self._settings.force if self._settings else None,
                    no_create=self._settings.force if self._settings else None,
                )
            except ValueError as e:
                # invalid key provided, try again
                wandb.termerror(e.args[0])
                continue
            except TimeoutError:
                wandb.termlog("W&B disabled due to login timeout.")
                return None, ApiKeyStatus.DISABLED
            if key is False:
                return None, ApiKeyStatus.NOTTY
            if not key:
                return None, ApiKeyStatus.OFFLINE
            return key, ApiKeyStatus.VALID

    def prompt_api_key(self) -> None:
        """Updates the global API key by prompting the user."""
        key, status = self._prompt_api_key()
        if status == ApiKeyStatus.NOTTY:
            directive = (
                "wandb login [your_api_key]"
                if self._settings.x_cli_only_mode
                else "wandb.login(key=[your_api_key])"
            )
            raise UsageError("api_key not configured (no-tty). call " + directive)

        self.update_session(key, status=status)
        self._key = key


def _login(
    *,
    anonymous: Optional[Literal["allow", "must", "never"]] = None,
    key: Optional[str] = None,
    relogin: Optional[bool] = None,
    host: Optional[str] = None,
    force: Optional[bool] = None,
    timeout: Optional[int] = None,
    _backend=None,
    _silent: Optional[bool] = None,
    _disable_warning: Optional[bool] = None,
    _entity: Optional[str] = None,
):
    if wandb.run is not None:
        if not _disable_warning:
            wandb.termwarn("Calling wandb.login() after wandb.init() has no effect.")
        return True

    wlogin = _WandbLogin()

    if _backend:
        wlogin.set_backend(_backend)

    if _silent:
        wlogin.set_silent(_silent)

    if _entity:
        wlogin.set_entity(_entity)

    # configure login object
    wlogin.setup(
        anonymous=anonymous,
        key=key,
        relogin=relogin,
        host=host,
        force=force,
        timeout=timeout,
    )

    if wlogin._settings._offline:
        return False
    elif wandb.util._is_kaggle() and not wandb.util._has_internet():
        wandb.termerror(
            "To use W&B in kaggle you must enable internet in the settings panel on the right."
        )
        return False

    if wlogin.should_use_identity_token():
        return True

    # perform a login
    logged_in = wlogin.login()

    if key:
        wlogin.configure_api_key(key)

    if logged_in:
        return logged_in

    if not key:
        wlogin.prompt_api_key()

    return wlogin._key or False