# oidc_app/core/backends.py

import logging

from django.core.exceptions import SuspiciousOperation
from django.urls import reverse
from mozilla_django_oidc.auth import OIDCAuthenticationBackend
from mozilla_django_oidc.utils import absolutify


LOGGER = logging.getLogger(__name__)


class CustomOIDCAuthenticationBackend(OIDCAuthenticationBackend):
    """Custom OIDC authentication backend."""

    def authenticate(self, request, **kwargs):
        """Authenticates a user based on the OIDC code flow."""

        self.request = request
        if not self.request:
            return None

        state = self.request.GET.get("state")
        code = self.request.GET.get("code")
        nonce = kwargs.pop("nonce", None)

        if not code or not state:
            return None

        reverse_url = self.get_settings(
            "OIDC_AUTHENTICATION_CALLBACK_URL", "oidc_authentication_callback"
        )

        token_payload = {
            "client_id": self.OIDC_RP_CLIENT_ID,
            "client_secret": self.OIDC_RP_CLIENT_SECRET,
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": absolutify(self.request, reverse(reverse_url)),
        }

        # Get the token
        token_info = self.get_token(token_payload)
        id_token = token_info.get("id_token")
        access_token = token_info.get("access_token")
        refresh_token = token_info.get("refresh_token")

        # Validate the token
        payload = self.verify_token(id_token, nonce=nonce)

        if payload:
            self.store_tokens(access_token, id_token, refresh_token) # <--- HERE: store tokens
            try:
                return self.get_or_create_user(access_token, id_token, payload)
            except SuspiciousOperation as exc:
                LOGGER.warning("failed to get or create user: %s", exc)
                return None

        return None

    def store_tokens(self, access_token, id_token, refresh_token): # <--- HERE: store tokens
        """Store OIDC tokens."""
        session = self.request.session

        if self.get_settings("OIDC_STORE_ACCESS_TOKEN", True):
            session["oidc_access_token"] = access_token

        if self.get_settings("OIDC_STORE_ID_TOKEN", False):
            session["oidc_id_token"] = id_token

        if self.get_settings("OIDC_STORE_REFRESH_TOKEN", True): # <--- HERE: Add refresh token option
            session["oidc_refresh_token"] = refresh_token

# oidc_app/settings.py

...
AUTHENTICATION_BACKENDS = (
    "oidc_app.core.backends.CustomOIDCAuthenticationBackend",
    "django.contrib.auth.backends.ModelBackend",
)
...
OIDC_RENEW_ID_TOKEN_EXPIRY_SECONDS = 60 * 60  # 1 hour
OIDC_STORE_ACCESS_TOKEN = os.environ.get("OIDC_STORE_ACCESS_TOKEN", True) # Store the access token in the OIDC backend
OIDC_STORE_ID_TOKEN = os.environ.get("OIDC_STORE_ID_TOKEN", True) # Store the ID token in the OIDC backend
OIDC_STORE_REFRESH_TOKEN = os.environ.get("OIDC_STORE_REFRESH_TOKEN", True) # Store the refresh token in the OIDC backend
OIDC_RP_SCOPES = os.environ.get("OIDC_RP_SCOPES", "openid profile email offline_access")  # The OIDC scopes to request

...

# oidc_app/authentication/views.py
...
from django.contrib.auth import logout
from django.http import HttpResponseRedirect
from django.contrib.auth.views import LogoutView
from django.utils.decorators import method_decorator
from django.views.decorators.cache import never_cache
from oidc_app.core.logout import oidc_logout # to be implemented
...
...
class LogoutViewSet(LogoutView):
    """
    API endpoint that handles user logout
    """
    @method_decorator(never_cache)
    def dispatch(self, request, *args, **kwargs):
        """
        Override the dispatch method to add the okta logout functionality
        """

        oidc_logout(request) # This is the function that does the oidc logout
        logout(request) # This is the django logout
        redirect_to = self.get_success_url()
        if redirect_to != request.get_full_path():
            # Redirect to target page once the session has been cleared.
            return HttpResponseRedirect(redirect_to)
        return super().dispatch(request, *args, **kwargs)

# oidc_app/urls.py
...
from oidc_app.authentication import views as auth_views

urlpatterns = [
   ...
    path("logout", auth_views.LogoutViewSet.as_view(), name="logout"),

# oidc_app/settings.py
OIDC_OP_TOKEN_REVOKE_ENDPOINT = (
    f"https://{OKTA_DOMAIN}/oauth2/default/v1/revoke"  # The OIDC token revocation endpoint
)

# oidc_app/core/oidc_logout.py

import requests
import logging
from django.conf import settings

LOGGER = logging.getLogger(__name__)


def revoke_token(token_type, token):
    """Revoke an OIDC token."""

    token_revoke_payload = {
        "client_id": settings.OIDC_RP_CLIENT_ID,
        "client_secret": settings.OIDC_RP_CLIENT_SECRET,
        "token": token,
        "token_type_hint": token_type,
    }

    try:
        response = requests.post(
            settings.OIDC_OP_TOKEN_REVOKE_ENDPOINT, data=token_revoke_payload
        )
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        LOGGER.error("Failed to revoke token: %s", e)


def oidc_logout(request):
    """Logout the user."""

    token_types = {
        "refresh_token": request.session.get("oidc_refresh_token"),
        "access_token": request.session.get("oidc_access_token"),
    }

    for token_type, token in token_types.items():
        if token:
            LOGGER.info("Revoking token of type %s", token_type)
            revoke_token(token_type, token)

# oidc_app/core/middleware.py

import logging
import time
from re import Pattern as re_Pattern

import requests
from django.urls import reverse
from django.utils.functional import cached_property
from mozilla_django_oidc.middleware import SessionRefresh as OIDCSessionRefresh

LOGGER = logging.getLogger(__name__)


class OIDCSessionRefreshMiddleware(OIDCSessionRefresh):
    @cached_property
    def exempt_urls(self):
        """Generate and return a set of url paths to exempt from SessionRefresh

        This takes the value of ``settings.OIDC_EXEMPT_URLS`` and appends three
        urls that mozilla-django-oidc uses. These values can be view names or
        absolute url paths.

        :returns: list of url paths (for example "/oidc/callback/")

        """
        exempt_urls = []
        for url in self.OIDC_EXEMPT_URLS:
            if not isinstance(url, re_Pattern):
                exempt_urls.append(url)
        return set(
            [url if url.startswith("/") else reverse(url) for url in exempt_urls]
        )

    def refresh_session(self, request):
        """Refresh the session with new data from the request session store."""

        refresh_token = request.session.get("oidc_refresh_token", None)

        token_refresh_payload = {
            "refresh_token": refresh_token,
            "client_id": self.get_settings("OIDC_RP_CLIENT_ID"),
            "client_secret": self.get_settings("OIDC_RP_CLIENT_SECRET"),
            "grant_type": "refresh_token",
        }

        try:
            response = requests.post(
                self.get_settings("OIDC_OP_TOKEN_ENDPOINT"), data=token_refresh_payload
            )
            response.raise_for_status()
        except requests.exceptions.RequestException as e:
            LOGGER.error("Failed to refresh session: %s", e)
            return False
        data = response.json()
        request.session.update(
            {
                "oidc_access_token": data.get("access_token"),
                "oidc_id_token_expiration": time.time() + data.get("expires_in"),
                "oidc_refresh_token": data.get("refresh_token"),
            }
        )
        return True

    def process_request(self, request):
        if not self.is_refreshable_url(request):
            LOGGER.debug("request is not refreshable")
            return

        expiration = request.session.get("oidc_id_token_expiration", 0)
        now = time.time()
        if expiration > now:
            # The id_token is still valid, so we don't have to do anything.
            LOGGER.debug("id token is still valid (%s > %s)", expiration, now)
            return

        LOGGER.debug("id token has expired")
        if not self.refresh_session(request):
            # If we can't refresh the session, then we need to reauthenticate the user.
            # As per the default OIDCSessionRefresh implementation.
            return super().process_request(request)

        LOGGER.debug("session refreshed")

# oidc_app/settings.py

MIDDLEWARE = [
    ...
   "django.contrib.auth.middleware.AuthenticationMiddleware",
    "oidc_app.core.middleware.OIDCSessionRefreshMiddleware",
    ...
]

OIDC_EXEMPT_URLS = [
    "oidc_authentication_init",
    "oidc_authentication_callback",
    "logout",
]