"""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.

It also takes the reader rules (petrolyst_reader.ReaderAccessMixin), as does the
local stack's database-login manager.
"""

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 petrolyst_reader import ReaderAccessMixin
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(ReaderAccessMixin, 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 []),
        }