import base64
import json
import time
from urllib.parse import unquote

import jwt
import requests
from flask import (
    Blueprint,
    jsonify,
    redirect,
    current_app,
    request,
    make_response,
    session,
    url_for,
)
from flask_cors import cross_origin

from .helpers import (
    get_public_key,
    get_allowed_root_domain,
    add_url_params,
    get_safe_repr,
)
from .logs import logger
from .settings import Settings

proxy = Blueprint("proxy", __name__)

ALLOWED_FORWARD_PARAMS = ("no_mfa_account", "dev_token", "prompt")
OMIT_REDIRECT_CHECK_ENV = "development"


@proxy.route("/login")
@proxy.route("/")
def login():
    """
    There are 2 login flows:

    1.  Relevant mostly to mobile applications, that supports
    authentication via the popup browser and catching the final
    redirect with data.
    Triggers by `response_type=token` param.
    Validates `redirect_uri` over the whitelist, stores `state` param
    and `redirect_url` into session for later processing after the
    authentication.

    2. Relevant to other webapps that can just use regular cookies
    or make extra request to get tokens from `/token` endpoint
    to get content from the http only cookies.
    """

    session.clear()

    next_url = request.args.get("next") or request.referrer

    if request.args.get("response_type") == "token":
        redirect_url = request.args.get("redirect_uri")
        state = request.args.get("state")
        if redirect_url and (
            redirect_url in current_app.config.get("ALLOWED_REDIRECT_URLS")
            or current_app.config.get("ENV") == OMIT_REDIRECT_CHECK_ENV
        ):
            session[
                current_app.config.get("REDIRECT_URL_SESSION_NAME")
            ] = redirect_url
            session[current_app.config.get("STATE_SESSION_NAME")] = state
        else:
            logger.bind(redirect_url=redirect_url).warning(
                "Invalid redirect url on login."
            )
            return "Invalid request data: invalid redirect_uri", 400
    else:
        session[current_app.config.get("REFERRER_SESSION_NAME")] = next_url

    params = {
        **{
            k: v
            for k, v in request.args.items()
            if k in ALLOWED_FORWARD_PARAMS
        },
        "auth_proxy": url_for(
            "proxy.process", _external=True, _scheme=request.scheme
        ),
    }
    if next_url:
        params["next"] = next_url

    url = add_url_params(
        current_app.config.get("LOGIN_URL"),
        params,
    )
    return redirect(url)


@proxy.route("/logout")
def logout():
    """
    Just making a cleanup for everything and redirecting
    to Atlas logout url.
    """

    session.clear()
    redirect_url = (
        request.args.get("redirect_uri")
        or request.args.get("next")
        or request.referrer
    )
    session[current_app.config.get("REFERRER_SESSION_NAME")] = redirect_url

    url = add_url_params(
        current_app.config.get("LOGOUT_URL"),
        {
            "auth_proxy": url_for(
                "proxy.process", _external=True, _scheme=request.scheme
            )
        },
    )

    resp = make_response(redirect(url))
    resp.set_cookie(
        current_app.config.get("DNA_BEARER_TOKEN_COOKIE_NAME"),
        "",
        expires=0,
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
    )
    resp.set_cookie(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME"),
        "",
        expires=0,
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
    )
    return resp


@proxy.route("/process", methods=["POST"])
def process():
    """
    There are 2 auth processing flows:

    1.  Relevant mostly to mobile applications, that supports
    authentication via the popup browser and catching the final redirect
    with data.
    Triggers by redirect url, that was previously validated over the
    whitelist and saved into session.
    It results in a redirect to mentioned url adding there
    state param previously stored in session and base64
    encoded tokens payload to the location hash.
    The state later might be also validated on the mobile
    app side, e.g. like expo-auth-session do.

    2. Relevant to other webapps that can just use regular cookies
    or make extra request to get tokens from `/token` endpoint
    to get content from the http only cookies.
    """

    access_token = request.form.get("access_token", "")
    refresh_token = request.form.get("refresh_token", "")
    next_url = session.get(
        current_app.config.get("REFERRER_SESSION_NAME")
    ) or request.form.get("next", "")
    redirect_url = session.get(
        current_app.config.get("REDIRECT_URL_SESSION_NAME")
    )
    state = session.get(current_app.config.get("STATE_SESSION_NAME"), "")
    if not next_url and not redirect_url:
        logger.warning("No next/redirect url on the process request")
        return "Invalid request data: no next/redirect url detected", 400

    # on logout or token issue redirecting to the app
    if not access_token and next_url:
        return redirect(unquote(next_url))

    rsa_key = get_public_key()
    try:
        claimset = jwt.decode(
            access_token,
            rsa_key,
            algorithms=current_app.config.get("JWT_ALGORITHMS"),
        )
    except jwt.PyJWTError as e:
        logger.bind(
            access_token=get_safe_repr(access_token),
            refresh_token=get_safe_repr(refresh_token),
            error=e,
        ).warning("Unable to validate the token")
        return "Invalid request data: invalid token", 400

    if redirect_url:
        if (
            redirect_url not in current_app.config.get("ALLOWED_REDIRECT_URLS")
            and current_app.config.get("ENV") != OMIT_REDIRECT_CHECK_ENV
        ):
            logger.bind(redirect_url=redirect_url).warning(
                "Invalid redirect url on process."
            )
            return "Invalid redirect url detected", 401

        expires_in = claimset["exp"] - int(time.time())
        payload = {
            "access_token": access_token,
            "refresh_token": refresh_token,
            "token_type": "Bearer",
            "expires_in": expires_in,
        }
        string_payload = base64.urlsafe_b64encode(
            json.dumps(payload).encode()
        ).decode()
        url = f"{unquote(redirect_url)}?state={state}#{string_payload}"
        resp = make_response(redirect(url))
    else:
        resp = make_response(redirect(unquote(next_url)))

    cookie_max_age = current_app.config.get("TOKEN_COOKIES_MAX_AGE")
    resp.set_cookie(
        current_app.config.get("DNA_BEARER_TOKEN_COOKIE_NAME"),
        access_token,
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
        max_age=cookie_max_age,
    )
    resp.set_cookie(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME"),
        refresh_token,
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
        max_age=cookie_max_age,
    )
    session.clear()
    return resp


@proxy.route("/token", methods=["GET", "OPTIONS", "POST"])
@cross_origin(supports_credentials=True, origins=Settings.ORIGINS)
def token():
    """
    Intended for apps that want to keep tokens in
    other places than http-only cookies, thus need
    to get them as in a response body.
    """

    access_token = request.cookies.get(
        current_app.config.get("DNA_BEARER_TOKEN_COOKIE_NAME")
    )
    refresh_token = request.cookies.get(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME")
    )
    if not access_token:
        logger.bind(
            access_token=get_safe_repr(access_token),
            refresh_token=get_safe_repr(refresh_token),
        ).warning("No access token in cookies")
        return "Access denied", 401

    rsa_key = get_public_key()
    try:
        claimset = jwt.decode(
            access_token,
            rsa_key,
            algorithms=current_app.config.get("JWT_ALGORITHMS"),
        )
    except jwt.PyJWTError:
        logger.bind(
            access_token=get_safe_repr(access_token),
            refresh_token=get_safe_repr(refresh_token),
        ).warning("Unable to validate the token")
        return "Access denied", 401

    expires_in = claimset["exp"] - int(time.time())

    return jsonify(
        {
            "access_token": access_token,
            "refresh_token": refresh_token,
            "token_type": "Bearer",
            "expires_in": expires_in,
        }
    )


@proxy.route("/token/refresh", methods=["OPTIONS", "POST"])
@cross_origin(supports_credentials=True, origins=Settings.ORIGINS)
def refresh_token_view():
    """
    Returns refreshed tokens in the response body,
    and also updates them in cookies.
    """

    refresh_token = request.form.get("refresh_token") or request.cookies.get(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME")
    )
    if not refresh_token:
        logger.warning("No refresh token in the request")
        return "Access denied", 401

    atlas_resp = requests.post(
        current_app.config.get("REFRESH_URL"),
        data={"refresh_token": refresh_token},
    )
    if not atlas_resp.ok:
        logger.bind(refresh_token=get_safe_repr(refresh_token)).warning(
            "Unable to refresh token in Atlas"
        )
        return "Access denied", 401
    tokens_data = atlas_resp.json()

    cookie_max_age = current_app.config.get("TOKEN_COOKIES_MAX_AGE")

    resp = make_response(jsonify(tokens_data))
    resp.set_cookie(
        current_app.config.get("DNA_BEARER_TOKEN_COOKIE_NAME"),
        tokens_data.get("access_token"),
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
        max_age=cookie_max_age,
    )
    resp.set_cookie(
        current_app.config.get("DNA_REFRESH_TOKEN_COOKIE_NAME"),
        tokens_data.get("refresh_token"),
        httponly=True,
        domain=get_allowed_root_domain(),
        secure=True,
        max_age=cookie_max_age,
    )
    return resp


@proxy.route("/health")
def health():
    return jsonify({"status": "ok"})