import json
import os
from typing import Dict, Optional, Tuple, Type

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", "{}"))


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


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"})

    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)
        access_token = _get_cookie(self.request, "access_token")
        refresh_token = _get_cookie(self.request, "refresh_token")
        api_key = None

        new_access_token, expires_in = (None, 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 not 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

        headers: Dict[str, str] = dict(kwargs.get("headers", {}))
        headers.update(extra_anaconda_headers)
        kwargs["headers"] = headers

        if access_token:
            headers["Authorization"] = f"Bearer {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 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
            new_access_token, expires_in = await get_access_token(self.request)
            headers["Authorization"] = f"Bearer {new_access_token}"
            resp = await backend_proxy(self.request, url, **dict(kwargs))
        finally:
            if new_access_token:
                self.set_cookie("access_token", 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
