# Copyright 2019 British Broadcasting Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from socket import getfqdn
from authlib.integrations.flask_oauth2 import AuthorizationServer, ResourceProtector
from authlib.integrations.sqla_oauth2 import (
    create_query_client_func,
    create_save_token_func,
    create_revocation_endpoint,
    create_bearer_token_validator,
)
from authlib.oauth2.rfc6749 import grants
from authlib.oauth2.rfc6749.errors import InvalidRequestError
from authlib.oauth2.rfc7523 import JWTBearerClientAssertion
from authlib.oauth2.rfc7636 import CodeChallenge
from werkzeug.security import gen_salt

from .models import db, OAuth2Client, OAuth2AuthorizationCode, OAuth2Token
from .db_utils import getResourceOwner
from .metadata import create_metadata
from .client_registration import ClientRegistrationEndpoint
from .constants import AUTH_VERSION_ROOT, TOKEN_ENDPOINT


class JWTClientAuth(JWTBearerClientAssertion):
    def validate_jti(self, claims, jti):
        # validate_jti is required by OpenID Connect
        # but it is optional by RFC7523
        # use cache to validate jti value
        return True

    def resolve_client_public_key(self, client, headers):
        if headers['alg'] == 'HS256':
            return client.client_secret
        if headers['alg'] == 'RS256':
            return client.public_key


class AuthorizationCodeGrant(grants.AuthorizationCodeGrant):
    def create_authorization_code(self, client, user, request):
        code = gen_salt(48)
        code_challenge = request.data.get('code_challenge', None)
        code_challenge_method = request.data.get('code_challenge_method', None)
        item = OAuth2AuthorizationCode(
            code=code,
            client_id=client.client_id,
            redirect_uri=request.redirect_uri,
            scope=request.scope,
            user_id=user.id,
            code_challenge=code_challenge,
            code_challenge_method=code_challenge_method
        )
        db.session.add(item)
        db.session.commit()
        return code

    def parse_authorization_code(self, code, client):
        item = OAuth2AuthorizationCode.query.filter_by(
            code=code, client_id=client.client_id).first()
        if item and not item.is_expired():
            return item

    def delete_authorization_code(self, authorization_code):
        db.session.delete(authorization_code)
        db.session.commit()

    def authenticate_user(self, authorization_code):
        return getResourceOwner(authorization_code.user_id)


class PasswordGrant(grants.ResourceOwnerPasswordCredentialsGrant):
    def authenticate_user(self, username, password):
        user = getResourceOwner(username)
        if user is None:
            raise InvalidRequestError("User Not Found")
        if user.check_password(password):
            return user


class RefreshTokenGrant(grants.RefreshTokenGrant):

    INCLUDE_NEW_REFRESH_TOKEN = True

    # Allowed client auth methods for token endpoint - allows public clients to refresh tokens
    TOKEN_ENDPOINT_AUTH_METHODS = [
        'client_secret_basic', 'client_secret_post', 'none'
    ]

    def authenticate_refresh_token(self, refresh_token):
        token = OAuth2Token.query.filter_by(refresh_token=refresh_token).first()
        if token and not token.revoked and not token.is_refresh_token_expired():
            return token

    def authenticate_user(self, credential):
        return getResourceOwner(credential.user_id)

    def revoke_old_credential(self, token_object):
        token_object.revoked = True
        db.session.commit()


class ClientCredentialsGrant(grants.ClientCredentialsGrant):
    # Allowed client auth methods for token endpoint - allows use of JWT assertion for authentication
    TOKEN_ENDPOINT_AUTH_METHODS = [
        'client_secret_basic',
        'client_secret_post',
        JWTClientAuth.CLIENT_AUTH_METHOD
    ]


query_client = create_query_client_func(db.session, OAuth2Client)
save_token = create_save_token_func(db.session, OAuth2Token)

authorization = AuthorizationServer(
    query_client=query_client,
    save_token=save_token,
)
require_oauth = ResourceProtector()


def config_oauth(app):
    authorization.init_app(app)
    authorization.metadata = create_metadata(app)

    # Register client authentication method for "client_secret_jwt " and "private_key_jwt"
    authorization.register_client_auth_method(
        JWTClientAuth.CLIENT_AUTH_METHOD,
        JWTClientAuth('https://{}{}{}'.format(getfqdn(), AUTH_VERSION_ROOT, TOKEN_ENDPOINT))
    )

    # Comment out unsupported grants
    # authorization.register_grant(grants.ImplicitGrant)
    # authorization.register_grant(PasswordGrant)
    authorization.register_grant(grants.ClientCredentialsGrant)
    authorization.register_grant(AuthorizationCodeGrant, [CodeChallenge(required=app.config["PKCE_REQUIRED"])])
    authorization.register_grant(RefreshTokenGrant)

    # Create Client Registration Endpoint
    authorization.register_endpoint(ClientRegistrationEndpoint)

    # Create Revocation Endpoint
    revocation_cls = create_revocation_endpoint(db.session, OAuth2Token)
    authorization.register_endpoint(revocation_cls)

    # Create Protected Resource Endpoint
    bearer_cls = create_bearer_token_validator(db.session, OAuth2Token)
    require_oauth.register_token_validator(bearer_cls())
