import datetime
import json
import os
from dataclasses import dataclass
from typing import Dict, Optional, Tuple, Type

import jwt
from jupyter_server.base.handlers import APIHandler
from jupyter_server.utils import url_path_join
from tornado.httputil import HTTPServerRequest
from tornado.routing import PathMatches, Rule, _RuleList

from aext_shared.backend_proxy import ProxyResponse, backend_proxy
from aext_shared.config import get_config
from aext_shared.errors import BackendError, UnauthorizedError

# TODO: This can be removed if/when we add anaconda-cloud-auth as a required dependency
try:
    from anaconda_cloud_auth.config import AuthConfig
    from anaconda_cloud_auth.exceptions import TokenNotFoundError
    from anaconda_cloud_auth.token import TokenInfo
except ImportError:  # pragma: no cover
    TokenInfo = None
    AuthConfig = None
    TokenNotFoundError = None

config = get_config()
debug_refresh_token = os.environ.get("DEBUG_REFRESH_TOKEN", None)
extra_anaconda_headers: Dict[str, str] = json.loads(os.environ.get("ANACONDA_CLOUD_EXTRA_HEADERS", "{}"))


@dataclass
class UserAccessCredentials:
    access_token: str = None
    new_access_token: str = None
    api_key: str = None
    username: str = None


def _get_api_key() -> Optional[str]:
    if TokenInfo is None:  # pragma: no cover
        # anaconda-cloud-auth is not installed
        return None
    auth_config = AuthConfig()
    try:
        return TokenInfo.load(domain=auth_config.domain).api_key
    except TokenNotFoundError:
        # We need to log in
        return None


def _get_username(token: Optional[str] = None) -> Optional[str]:
    try:
        if not token:
            token = _get_api_key()
        user_info = jwt.decode(token, algorithms=["RS256"], options={"verify_signature": False})
        return user_info.get("sub")
    except (TokenNotFoundError, jwt.exceptions.DecodeError):
        return None


class BackendHandler(APIHandler):
    def write_error(self, status_code, **kwargs):
        (err_cls, err, traceback) = kwargs.get("exc_info", (None, None, None))
        if isinstance(err, BackendError):
            self.set_status(err.status_code)
            self.finish(err.data)
            return
        self.set_status(status_code)
        self.finish({"reason": "unknown"})

    @classmethod
    def _is_token_expired(cls, token: str) -> bool:
        dt = datetime.datetime.now(datetime.timezone.utc)
        utc_timestamp = dt.replace(tzinfo=datetime.timezone.utc).timestamp()
        try:
            decoded_token = jwt.decode(token, options={"verify_signature": False})
        except jwt.DecodeError:
            return False
        # checking if the token expired - if the key is not found assume it is expired by returning utc_timestamp
        return utc_timestamp >= decoded_token.get("exp", utc_timestamp)

    async def get_user_access_credentials(self, auth_optional: bool = False) -> UserAccessCredentials:
        """
        Retrieves the access token used to proxy requests to anaconda.cloud API.

        Returns:
            - access_token: The current access_token
            - new_access_token: Returned in case it is a new access_token
            - api_key: Returns an API key in case it is a local stored API key
        """
        access_token = _get_cookie(self.request, "access_token")
        refresh_token = _get_cookie(self.request, "refresh_token")
        api_key = None
        new_access_token = None
        username = None

        if not access_token and not refresh_token and not debug_refresh_token:
            # Only attempt to get the local API key if neither cookie is set
            access_token = api_key = _get_api_key()

        if api_key is None:
            if not access_token or self._is_token_expired(access_token):
                # No matter what, try to update the access token
                # If auth is not required, ignore any failures that happen because of it
                try:
                    new_access_token, expires_in = await get_access_token(self.request)
                    access_token = new_access_token
                except BackendError as e:  # pragma: no cover
                    if not auth_optional:
                        raise e

        username = _get_username(access_token or new_access_token or api_key)
        return UserAccessCredentials(
            access_token=access_token,
            new_access_token=new_access_token,
            api_key=api_key,
            username=username,
        )

    async def anaconda_cloud_proxy(self, path: str, **kwargs) -> ProxyResponse:
        """
        This function proxies a request to anaconda.cloud API

        path: endpoint that should be hit
        kwargs: can be used to pass some request specifics such as:
                * headers
                * json: sets content-type to json
                * method: http verb [get, post, put, patch, delete]
        """
        auth_optional = kwargs.pop("auth_optional", False)
        expires_in = None
        user_access_credentials = await self.get_user_access_credentials(auth_optional=auth_optional)
        headers: Dict[str, str] = dict(kwargs.get("headers", {}))

        headers.update(extra_anaconda_headers)
        kwargs["headers"] = headers

        if user_access_credentials.access_token:
            headers["Authorization"] = f"Bearer {user_access_credentials.access_token}"

        # Enforce that the bearer token only gets sent to anaconda.cloud/api by setting the full URL in the helper method
        url = f"{config['cloud']['url']}/{path}"

        try:
            resp = await backend_proxy(self.request, url, **dict(kwargs))
        except BackendError as e:
            if user_access_credentials.api_key is not None:  # pragma: no cover
                # If we are using an API key, then we don't want to do the retry after refresh
                raise e

            if auth_optional or e.data.get("remote_status_code") not in [401, 403]:
                # Some unknown error, not due to expired token
                raise e

            # Refresh the token and retry the request
            user_access_credentials.new_access_token, expires_in = await get_access_token(self.request)
            headers["Authorization"] = f"Bearer {user_access_credentials.new_access_token}"
            resp = await backend_proxy(self.request, url, **dict(kwargs))
        finally:
            if user_access_credentials.new_access_token:
                self.set_cookie(
                    "access_token",
                    user_access_credentials.new_access_token,
                    max_age=expires_in,
                    secure=True,
                    samesite="strict",
                )
        return resp


def create_rules(base_url: str, prefix: str, handlers: Dict[str, Type[BackendHandler]]) -> _RuleList:
    rules: _RuleList = [
        Rule(PathMatches(url_path_join(base_url, prefix, key)), value) for key, value in handlers.items()
    ]
    return rules


def _get_cookie(request: HTTPServerRequest, name: str) -> Optional[str]:
    cookie = request.cookies.get(name)
    return cookie.value if cookie else None


async def get_access_token(incoming_request: HTTPServerRequest) -> Tuple[str, int]:
    refresh_token = _get_cookie(incoming_request, "refresh_token")
    if not refresh_token:
        refresh_token = debug_refresh_token

    if not refresh_token:
        raise UnauthorizedError({"reason": "missing refresh_token"})

    payload = {
        "grant_type": "refresh_token",
        "refresh_token": refresh_token,
    }
    try:
        proxy_response = await backend_proxy(
            incoming_request,
            f"{config['cloud']['url']}/iam/token",
            method="POST",
            json=payload,
            headers=dict(extra_anaconda_headers),
        )
    except BackendError as e:
        raise UnauthorizedError(e.data) from e
    access_token: str = proxy_response["remote_data"]["access_token"]
    expires_in: int = proxy_response["remote_data"]["expires_in"]
    return access_token, expires_in
