from tornado import web
from jupyterhub.auth import Authenticator
from jupyterhub.handlers.login import LogoutHandler
from traitlets import Dict, List, Type, Instance, Unicode, default
from ._jhub_shared import HeaderLoginHandler, UserDataHandler, Parser


class HeaderAuthenticator(Authenticator):
    """
    Authenticates the user via the allowed_headers
    defined 'auth' HTTP/HTTPS header.
    """

    allowed_headers = Dict(
        default_value={"auth": "Remote-User"},
        allow_none=False,
        help="""Dict of HTTP/HTTPS headers that the authenticator will process during login.
         The 'auth' key must be set to a header, default is Remote-User.
         Additional headers will be stored in the user's 'auth_state' session store
        """,
    ).tag(config=True)

    header_parser_classes = Dict(
        default_value={"auth": Parser},
        traits={Unicode(): Type(Parser, Parser)},
        allow_none=True,
        help="""Dict of classes that should be used to parse the allowed_headers dictionary
        """,
    ).tag(config=True)

    header_parsers = Dict({"auth": Instance(Parser)})

    @default("header_parsers")
    def _header_parsers_default(self):
        return {
            parser_key: parser_val(parent=self)
            for parser_key, parser_val in self.header_parser_classes.items()
        }

    spawner_shared_headers = List(
        default_value=[],
        allow_none=False,
        help="""List of headers that should be shared from the auth_state dict as environment
         variables to the spawner via the pre_spawn_start hook.
        """,
    ).tag(config=True)

    user_external_allow_attributes = List(
        default_value=[],
        traits=[Unicode()],
        allow_none=True,
        help="""List of user attributes that are allowed to be defined externally
         via a JSON post request. Submit via /set-user-data
        """,
    ).tag(config=True)

    def __init__(self, **kwargs):
        if "auth" not in self.allowed_headers:
            self.log.error(
                "HeaderAuthenticator - requires that the allowed_headers "
                "has an 'auth' key"
            )
            raise KeyError("Missing required 'auth' key in allowed_headers")
        super().__init__(**kwargs)

    def get_handlers(self, app):
        return [
            (r"/login", HeaderLoginHandler),
            (r"/logout", LogoutHandler),
            (r"/set-user-data", UserDataHandler),
        ]

    async def authenticate(self, handler, data):
        self.log.debug(
            "HeaderAuthenticator - Request authentication with "
            "handler: {}, data: {} of type: {}".format(handler, data, type(data))
        )
        user_data = {}
        # Process remaining allowed_headers, save valid in user_data
        for allowed_index, allowed_value in self.allowed_headers.items():
            auth_data = data.get(allowed_value, "")
            if auth_data:
                prepared_data = None
                if allowed_index in self.header_parsers:
                    prepared_data = self.header_parsers[allowed_index].parse(auth_data)
                else:
                    prepared_data = data
                if prepared_data:
                    user_data[allowed_value] = prepared_data

        self.log.debug(
            "HeaderAuthenticator - Prepared user_data: {} "
            "for auth check".format(user_data)
        )
        if self.allowed_headers["auth"] not in user_data:
            self.log.error(
                "HeaderAuthenticator - Failed to find the 'auth' key "
                "in the user_data dictionary, this is required "
                "to set the authenticated user's username"
            )
            raise web.HTTPError(
                401,
                "Authentication failed, "
                "missing information to authenticate you with.",
            )

        user = {"name": user_data.pop(self.allowed_headers["auth"], None)}
        # Something left in user_data, put in auth_state
        if user_data:
            user.update({"auth_state": user_data})

        self.log.info("Authenticated: {} - Login".format(user))
        return user

    async def pre_spawn_start(self, user, spawner):
        """Pass upstream_token to spawner via environment variable"""
        auth_state = await user.get_auth_state()
        if not auth_state:
            self.log.debug(
                "HeaderAuthenticator - pre_spawn_hook, "
                "auth_state: {}".format(auth_state)
            )
            # auth_state not enabled
            return

        self.log.debug(
            "HeaderAuthenticator - pre_spawn_hook, "
            "loaded auth_state {}".format(auth_state)
        )
        # Share permitted headers
        if not self.spawner_shared_headers:
            self.log.debug(
                "HeaderAuthenticator - no headers were "
                "shared with spawner environment: {}".format(
                    self.spawner_shared_headers
                )
            )
            return

        for auth_key, auth_val in auth_state.items():
            if auth_key in self.spawner_shared_headers:
                spawner.environment[auth_key] = auth_val
        self.log.debug(
            "HeaderAuthenticator - shared auth_state headers: {} with "
            "spawner environment: {}".format(
                self.spawner_shared_headers, spawner.environment
            )
        )
