chat_backend/llm_be/chat_backend/views_oauth.py

"""OAuth SSO start + callback views (#24)."""

from __future__ import annotations

import logging
from urllib.parse import urlencode

from django.conf import settings
from django.http import HttpResponseRedirect
from django.urls import reverse
from rest_framework import permissions, status
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework_simplejwt.tokens import RefreshToken

from .models import OAuthIdentity
from .oauth import (
    OAuthError,
    build_authorization_url,
    configured_providers,
    dump_oauth_state,
    exchange_code_for_profile,
    load_oauth_state,
    provider_configured,
    resolve_user_from_profile,
    upsert_identity,
)

logger = logging.getLogger(__name__)


def _callback_redirect_uri(request, provider: str) -> str:
    """Absolute backend callback URL registered with the IdP."""
    path = reverse("oauth_callback", kwargs={"provider": provider})
    base = (settings.OAUTH_CALLBACK_BASE_URL or "").rstrip("/")
    if base:
        return f"{base}{path}"
    return request.build_absolute_uri(path)


def _frontend_callback_url(**params: str) -> str:
    base = settings.FRONTEND_BASE_URL.rstrip("/")
    query = urlencode({k: v for k, v in params.items() if v is not None and v != ""})
    return f"{base}/auth/callback/?{query}"


def _redirect_error(code: str, message: str = "") -> HttpResponseRedirect:
    return HttpResponseRedirect(
        _frontend_callback_url(error=code, error_description=message or code)
    )


class OAuthStartView(APIView):
    """Redirect the browser to Google / Microsoft authorize URL."""

    permission_classes = (permissions.AllowAny,)
    authentication_classes = ()
    http_method_names = ["get"]

    def get(self, request, provider: str):
        provider = (provider or "").lower()
        if provider not in OAuthIdentity.Provider.values:
            return Response(
                {"detail": "Unsupported OAuth provider."},
                status=status.HTTP_404_NOT_FOUND,
            )
        if not provider_configured(provider):
            return Response(
                {"detail": f"{provider} OAuth is not configured."},
                status=status.HTTP_503_SERVICE_UNAVAILABLE,
            )

        intent = (request.query_params.get("intent") or "login").lower()
        if intent not in {"login", "signup"}:
            return Response(
                {"detail": "intent must be 'login' or 'signup'."},
                status=status.HTTP_400_BAD_REQUEST,
            )
        if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION:
            return Response(
                {"detail": "Account registration is disabled."},
                status=status.HTTP_403_FORBIDDEN,
            )

        state = dump_oauth_state(provider=provider, intent=intent)
        redirect_uri = _callback_redirect_uri(request, provider)
        try:
            auth_url = build_authorization_url(
                provider=provider, redirect_uri=redirect_uri, state=state
            )
        except OAuthError as exc:
            return Response({"detail": exc.message}, status=status.HTTP_400_BAD_REQUEST)

        return HttpResponseRedirect(auth_url)


class OAuthCallbackView(APIView):
    """IdP redirect target — exchange code, create/link user, send JWTs to FE."""

    permission_classes = (permissions.AllowAny,)
    authentication_classes = ()
    http_method_names = ["get"]

    def get(self, request, provider: str):
        provider = (provider or "").lower()
        if provider not in OAuthIdentity.Provider.values:
            return _redirect_error("invalid_provider", "Unsupported OAuth provider.")

        error = request.query_params.get("error")
        if error:
            description = request.query_params.get("error_description") or error
            code = "access_denied" if error == "access_denied" else "provider_error"
            return _redirect_error(code, description)

        code = request.query_params.get("code")
        state = request.query_params.get("state")
        if not code or not state:
            return _redirect_error("missing_code", "Missing OAuth code or state.")

        try:
            state_data = load_oauth_state(state)
            if state_data["provider"] != provider:
                raise OAuthError("invalid_state", "OAuth provider mismatch.")

            redirect_uri = _callback_redirect_uri(request, provider)
            profile = exchange_code_for_profile(
                provider=provider, code=code, redirect_uri=redirect_uri
            )
            user, created = resolve_user_from_profile(
                profile=profile, intent=state_data["intent"]
            )
            # Refresh stored tokens on every successful login.
            upsert_identity(user, profile)
        except OAuthError as exc:
            logger.info("OAuth callback failed (%s): %s", exc.code, exc.message)
            return _redirect_error(exc.code, exc.message)
        except Exception:
            logger.exception("Unexpected OAuth callback failure")
            return _redirect_error("server_error", "Unexpected OAuth error.")

        refresh = RefreshToken.for_user(user)
        from finance.services.plans import needs_checkout as user_needs_checkout

        needs_checkout = "1" if (created and user_needs_checkout(user)) else "0"
        return HttpResponseRedirect(
            _frontend_callback_url(
                access=str(refresh.access_token),
                refresh=str(refresh),
                created="1" if created else "0",
                needs_checkout=needs_checkout,
            )
        )


def oauth_public_flags() -> dict:
    """Feature flags for /public/settings/."""
    return configured_providers()