"""Keycloak sign-in for Superset at atoma.petrolyst.com.

Set as CUSTOM_SECURITY_MANAGER in superset_config.py. Checked against
Flask-AppBuilder 5.0.2 as bundled in apache/superset:6.1.0. Three changes to
the stock behaviour:

- User info and roles come from the ID token that Keycloak issues for client
  "petrolyst" (claim petrolyst_roles). FAB's built-in Keycloak branch would read
  a "groups" claim from the userinfo endpoint instead.
- /login/ goes straight to Keycloak. Superset 6.1 normally registers its React
  login page on /login/, which only shows a "Sign in with keycloak" button.
- /logout/ also ends the Keycloak session, so the next visit asks for the
  password and OTP again.
"""

import logging
from typing import Any, Optional
from urllib.parse import urlencode

from flask import g, redirect, request, url_for
from flask_appbuilder import expose
from flask_appbuilder.security.views import AuthOAuthView
from flask_appbuilder.utils.base import get_safe_redirect
from flask_login import logout_user
from werkzeug.wrappers import Response as WerkzeugResponse

from superset.security import SupersetSecurityManager

logger = logging.getLogger(__name__)

PROVIDER = "keycloak"
ROLE_CLAIM = "petrolyst_roles"


class PetrolystAuthOAuthView(AuthOAuthView):
    @expose("/login/")
    @expose("/login/<provider>")
    def login(self, provider: Optional[str] = None) -> WerkzeugResponse:
        authenticated = g.user is not None and g.user.is_authenticated
        if provider is None and not authenticated:
            # FAB encodes the query arguments into the OAuth state and runs
            # get_safe_redirect on "next" again after the callback; checking it
            # here as well keeps an unsafe value out of the state altogether.
            args = {}
            if request.args.get("next"):
                args["next"] = get_safe_redirect(request.args["next"])
            return redirect(url_for(".login", provider=PROVIDER, **args))
        return super().login(provider)

    @expose("/logout/")
    def logout(self) -> WerkzeugResponse:
        logout_user()
        remote = self.appbuilder.sm.oauth_remotes[PROVIDER]
        end_session = remote.load_server_metadata().get("end_session_endpoint")
        if not end_session:
            logger.warning("Keycloak metadata has no end_session_endpoint")
            return redirect(self.appbuilder.get_url_for_index)
        # client_id instead of id_token_hint: Keycloak then asks the user to
        # confirm the logout once, and no ID token has to live in the cookie.
        params = {
            "client_id": remote.client_id,
            "post_logout_redirect_uri": request.host_url,
        }
        return redirect(f"{end_session}?{urlencode(params)}")


class PetrolystSecurityManager(SupersetSecurityManager):
    authoauthview = PetrolystAuthOAuthView
    # Without this, Superset registers its React login page on /login/ before
    # FAB's OAuth view, and that page wins the route.
    register_superset_auth_view = False

    def oauth_user_info(
        self, provider: str, response: Optional[dict[str, Any]] = None
    ) -> dict[str, Any]:
        if provider != PROVIDER:
            raise ValueError(f"unsupported OAuth provider {provider!r}")
        # Authlib verifies the ID token (signature, issuer, audience, nonce)
        # and puts its claims under "userinfo".
        claims = (response or {}).get("userinfo")
        if not claims:
            claims = self.oauth_remotes[provider].userinfo()
        username = claims.get("preferred_username")
        if not username:
            raise ValueError("Keycloak token has no preferred_username")
        return {
            "username": username,
            "email": claims.get("email", ""),
            "first_name": claims.get("given_name", ""),
            "last_name": claims.get("family_name", ""),
            "role_keys": list(claims.get(ROLE_CLAIM) or []),
        }