from functools import wraps
import json
import logging

import flask
from flask import current_app, request, render_template
from flask.helpers import url_for
from oic import rndstr
from oic.oic.provider import Provider
from oic.utils.sdb import create_session_db
from oic.oic.message import EndSessionRequest
from oic.utils.authn.authn_context import AuthnBroker
from oic.utils.authn.client import verify_client
from oic.utils.authz import AuthzHandling
from oic.utils.keyio import keyjar_init
from urllib.parse import parse_qsl
from werkzeug.utils import redirect
from .AuthorizationResponseHandler import AuthResponseHandler, AuthorizationResponse, AuthResponseProcessError
from typing import List, Dict
from stores import IdentityProviderStore, SessionStore, UserStore
from models import IdentityProvider, TokenType, TokenSource, User
from .IdentityProviderContext import IdentityProviderContext
from .PyOIDCHandler import PyOIDCHandler
from .OpenIdService import OpenIdService
from models import IdServerConfiguration, Session
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)


class IdentityService:
    RequiredClaims = ['sub', 'email', 'given_name']
    """
    Identity Service
    """

    """
    Args:
        provider_configurations (Mapping[str, ProviderConfiguration]):
            provider configurations by name
        app (flask.app.Flask): optional Flask app
    """
    def __init__(self, identity_server_configuration: IdServerConfiguration, 
                       openid_service: OpenIdService,
                       identity_provider_store: IdentityProviderStore, 
                       session_store: SessionStore,
                       user_store: UserStore,
                       app: flask.Flask=None):
        #self._provider_configurations = provider_configurations

        self.identity_server_configuration = identity_server_configuration
        self.openid_service = openid_service
        self.identity_provider_store = identity_provider_store
        self.session_store = session_store
        self.user_store = user_store
        
        self.identity_providers: Dict[str, IdentityProvider] = {}
        self.identity_contexts: Dict[str, IdentityProviderContext] = {}

        self.clients: Dict[str, PyOIDCHandler]= {}
        self._logout_view = None
        self._error_view = None
        self._redirect_uri_endpoint = None

        if app:
            self.init_app(app)

    def init_app(self, app):
        self.identity_providers = self.identity_provider_store.get_providers()
        for identity_provider_name in self.identity_providers:
            identity_context = IdentityProviderContext(self.identity_providers[identity_provider_name])
            
            app.add_url_rule('/' + identity_context.redirect_uri_endpoint,
                    identity_context.redirect_uri_endpoint,
                    self._handle_authentication_response,
                    methods=['GET', 'POST'])
            
            self.identity_contexts[identity_provider_name] = identity_context
            
            with app.app_context():
                self.clients[identity_provider_name] = PyOIDCHandler(self.identity_server_configuration.address, identity_context)

        if len(self.identity_providers) > 1:
            app.add_url_rule('/' + self.identity_server_configuration.providers_form_route,
                        self.identity_server_configuration.providers_form_route,
                        self._show_providers_form,
                        methods=['GET'])
                        
    def _show_providers_form(self):
        return render_template('providers_form.html', login_endpoint='/user/login', 
                                                    providers=self.identity_providers.values(), 
                                                    use_builtin_identity_provider=self.identity_server_configuration.use_builtin_identity_provider)

    def _get_post_logout_redirect_uri(self, client):
        if client.post_logout_redirect_uris:
            return client.post_logout_redirect_uris[0]
        return self._get_url_for_logout_view()

    def _get_url_for_logout_view(self):
        return url_for(self._logout_view.__name__, _external=True) if self._logout_view else None

    def _register_client(self, client):
        def default_post_logout_redirect_uris():
            url_for_logout_view = self._get_url_for_logout_view()
            if url_for_logout_view:
                return [url_for_logout_view]
            return []

        client_registration_args = {}
        post_logout_redirect_uris = client._provider_configuration._client_registration_info.get(
            'post_logout_redirect_uris',
            default_post_logout_redirect_uris())
        if post_logout_redirect_uris:
            logger.debug('registering with post_logout_redirect_uris=%s', post_logout_redirect_uris)
            client_registration_args['post_logout_redirect_uris'] = post_logout_redirect_uris
        client.register(client_registration_args)

    def _authenticate(self, client: PyOIDCHandler, interactive: bool=True):
        if not client.is_registered():
            self._register_client(client)

        flask.session['destination'] = flask.request.url.split('?')[0]
        flask.session['state'] = rndstr()
        flask.session['nonce'] = rndstr()
        # flask.session['current_provider'] = client.identity_provider_context.identity_provider.provider_name

        # Use silent authentication for session refresh
        # This will not show login prompt to the user
        extra_auth_params: Dict[str, Any] = {}
        if not interactive:
            extra_auth_params['prompt'] = 'none'

        extra_auth_params["scope"] = client.identity_provider_context.identity_provider.scopes
       
        for auth_param in client.identity_provider_context.identity_provider.extra_authorization_parameters:
            if auth_param not in extra_auth_params:
                extra_auth_params[auth_param] = client.identity_provider_context.identity_provider.extra_authorization_parameters[auth_param]

        login_url = client.authentication_request(flask.session['state'],
                                                  flask.session['nonce'],
                                                  extra_auth_params)

        auth_params = dict(parse_qsl(login_url.split('?')[1]))
        flask.session['fragment_encoded_response'] = AuthResponseHandler.expect_fragment_encoded_response(auth_params)
        return redirect(login_url)

    def _handle_authentication_response(self):
        has_error = flask.request.args.get('error', False, lambda x: bool(int(x)))
        if has_error:
            if 'error' in flask.session:
                return self._show_error_response(flask.session.pop('error'))
            return 'Something went wrong.'

        session = self.get_session()
        if not session.provider_name:
            return "Bad Request", 400

        client = self.clients[session.provider_name]

        if flask.session.pop('fragment_encoded_response', False):
            # return importlib_resources.read_binary('flask_pyoidc', 'parse_fragment.html').decode('utf-8')
            return render_template('parse_fragment.html', redirect_uri=client.redirect_uri)
            
        is_processing_fragment_encoded_response = flask.request.method == 'POST'

        if is_processing_fragment_encoded_response:
            auth_resp = flask.request.form
        else:
            auth_resp = flask.request.args

        authn_resp = client.parse_authentication_response(auth_resp)
        logger.debug('received authentication response: %s', authn_resp.to_json())

        try:
            result = AuthResponseHandler(client).process_auth_response(authn_resp,
                                                                       flask.session.pop('state'),
                                                                       flask.session.pop('nonce'))
        except AuthResponseErrorResponseError as e:
            return self._handle_error_response(e.error_response, is_processing_fragment_encoded_response)
        except AuthResponseProcessError as e:
            return self._handle_error_response({'error': 'unexpected_error', 'error_description': str(e)},
                                               is_processing_fragment_encoded_response)

        if current_app.config.get('OIDC_SESSION_PERMANENT', True):
            flask.session.permanent = True

        session.update_provider_token(result.access_token, 
                                    result.id_token_jwt, 
                                    result.refresh_token)

        destination = flask.session.pop('destination')
        if is_processing_fragment_encoded_response:
            # if the POST request was from the JS page handling fragment encoded responses we need to return
            # the destination URL as the response body
            return destination

        # create local jwt and account here
        if not 'sub' in result.id_token_claims:
            raise Exception("Subject not set in provider token !")

        provider_subject = result.id_token_claims['sub']

        if not 'email' in result.id_token_claims:
            raise Exception("Email not set in provider token !")
        
        user_email = result.id_token_claims['email']

        # get user from db
        user = self.user_store.get_user_by_email(client.identity_provider_context.identity_provider.provider_name, user_email)

        local_jwt_claims: Dict[str, Any] = {}
        for claim_mapping in client.identity_provider_context.identity_provider.user_claim_mappings:
            if claim_mapping.source == TokenSource.id_token:
                if claim_mapping.source_claim in result.id_token_claims:
                    local_jwt_claims[claim_mapping.destination_claim] = result.id_token_claims[claim_mapping.source_claim]
                elif claim_mapping.required:
                    raise Exception("Missing required claim") # todo specific exception
        
        if not user:
            user = User(provider_subject=provider_subject, 
                    email=user_email, 
                    provider_name=client.identity_provider_context.identity_provider.provider_name,
                    given_name=result.id_token_claims["given_name"],
                    family_name=result.id_token_claims["family_name"], # todo family_name is optional
                    claims=local_jwt_claims)

            self.user_store.add_user(user)
        else:
            # check if provider sub is matching
            if user.provider_subject != result.id_token_claims["sub"]:
                raise("User Provider Subject has changed !")
            
            user.claims = local_jwt_claims
            self.user_store.update_user(user)

        # local jwt creation
        # todo

        return redirect(destination)

    def _handle_error_response(self, error_response, should_redirect=False):
        if should_redirect:
            # if the current request was from the JS page handling fragment encoded responses we need to return
            # a URL for the error page to redirect to
            flask.session['error'] = error_response
            return '/' + self._redirect_uri_endpoint + '?error=1'
        return self._show_error_response(error_response)

    def _show_error_response(self, error_response):
        logger.error(json.dumps(error_response))
        if self._error_view:
            error = {k: error_response[k] for k in ['error', 'error_description'] if k in error_response}
            return self._error_view(**error)

        return 'Something went wrong with the authentication, please try to login again.'

    def get_session(self) -> Session:
        if 'session_id' not in flask.session:
            session = self.session_store.create_session()
            flask.session['session_id'] = session.session_id
            return session

        session_id: str = flask.session["session_id"]
        session = self.session_store.get_session(session_id) or self.session_store.create_session()
        flask.session['session_id'] = session.session_id

        return session

    def oidc_auth(self, view_func):
        @wraps(view_func)
        def wrapper(*args, **kwargs):
            session = self.get_session()

            if len(self.clients) == 0:
                raise Exception("No OpenId Provider configured")
            client = next(iter(self.clients.values()))

            if session.should_refresh(client.session_refresh_interval_seconds):
                # todo review logic
                logger.debug('user auth will be refreshed "silently"')
                return self._authenticate(client, interactive=False)
            elif session.is_authenticated():
                logger.debug('user is already authenticated')
                return view_func(*args, **kwargs)
            else:
                logger.debug('user not authenticated, start flow')
                if len(self.clients) > 1:
                    if 'provider' in request.args:
                        provider_name = request.args.get('provider', str, None)
                        if provider_name is None or provider_name not in self.identity_providers:
                            return "Bad Request", 400
                        client = self.clients[provider_name]
                        session.provider_name = provider_name
                        return self._authenticate(client)
                    else:
                        # flask.session["current_provider"] = client.identity_provider_context.identity_provider.provider_name
                        return redirect('/' + self.identity_server_configuration.providers_form_route)
                else:
                    return self._authenticate(client)
        return wrapper

    def _logout(self):
        logger.debug('user logout')
        try:
            session = self.get_session()
        except UninitialisedSession as e:
            logger.info('user was already logged out, doing nothing')
            return None

        id_token_jwt = session.provider_token.id_token
        client = self.clients[session.provider_name]
        flask.session.pop('session_id')

        if client.provider_end_session_endpoint:
            flask.session['end_session_state'] = rndstr()

            end_session_request = EndSessionRequest(id_token_hint=id_token_jwt,
                                                    post_logout_redirect_uri=self._get_post_logout_redirect_uri(client),
                                                    state=flask.session['end_session_state'])

            logger.debug('send endsession request: %s', end_session_request.to_json())

            return redirect(end_session_request.request(client.provider_end_session_endpoint), 303)
        return None

    def oidc_logout(self, view_func):
        self._logout_view = view_func

        @wraps(view_func)
        def wrapper(*args, **kwargs):
            if 'state' in flask.request.args:
                # returning redirect from provider
                if flask.request.args['state'] != flask.session.pop('end_session_state'):
                    logger.error("Got unexpected state '%s' after logout redirect.", flask.request.args['state'])
                return view_func(*args, **kwargs)

            redirect_to_provider = self._logout()
            if redirect_to_provider:
                return redirect_to_provider

            return view_func(*args, **kwargs)

        return wrapper

    def error_view(self, view_func):
        self._error_view = view_func
        return view_func