import json
import os
from urllib.request import urlopen

import jwt
from flask import Flask, abort, request, render_template, jsonify
from jwt.utils import base64url_decode

ALGORITHMS = "RS256"
PUBLIC_KEY = "https://dev-static-um.atlas.stream/pem/public.pem"
API_GATEWAY_MODE = os.environ.get("API_GATEWAY_MODE")
ATLAS_BEARER_TOKEN_COOKIE_NAME = os.environ.get(
    "ATLAS_BEARER_TOKEN_COOKIE_NAME", "dna_bearer_token"
)


app = Flask(__name__, template_folder="templates")


@app.route("/")
def index():
    """Available for all."""

    return render_template("spa.html")


@app.route("/health")
def health():
    """Available for all."""

    return jsonify({"status": "ok"})


@app.route("/api/protected")
def protected():
    """Available only for authorized users."""

    user = get_user()

    if not is_authorized(
        user,
        {"decibel/role": "user", "dna/role": "user", "rti/role": "user"}
    ):
        abort(401)

    token = get_token()

    return f"Protected endpoint: authorized for {user.get('sub')} / " \
           f"{user.get('email')} / {user.get('name')} account / " \
           f"using token: {token}"


def get_user():
    """Decoding and validating token with user information."""

    if API_GATEWAY_MODE:
        try:
            payload = json.loads(
                base64url_decode(request.headers.get("X-Userinfo"))
            )
        except Exception:
            return None

        return payload

    token = get_token()
    rsa_key = get_public_key()

    if not token:
        return None

    try:
        payload = jwt.decode(
            token,
            rsa_key,
            algorithms=ALGORITHMS,
        )
    except jwt.PyJWTError as e:
        print(f"Token validation error: {e}")
        return None

    return payload


def get_token():
    """Extract token from possible places of the request."""

    # case with token in cookie
    token = request.cookies.get(ATLAS_BEARER_TOKEN_COOKIE_NAME)

    # case with token in header, e.g. Authorization: Bearer <token_string>
    auth_header = request.headers.get("Authorization", "")
    if not token and auth_header.startswith("Bearer "):
        token = auth_header.replace("Bearer ", "")

    return token


def get_public_key():
    """Fetch public key from shared resource on the network."""

    public_key = urlopen(PUBLIC_KEY)
    public_key = public_key.read()
    key = public_key.decode()
    return key


def is_authorized(user, allowed_claims=None):
    """Checking permissions with valid token."""

    if not user:
        return False

    for claim_name, claim_value in allowed_claims.items():
        if claim_value in user.get(claim_name, []):
            return True

    return False


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=os.environ.get('PORT', 5000))