"""
Has all /sso/* routes

/sso/key/generate - handles user signing in with SSO and redirects to /sso/callback
/sso/callback - returns JWT Redirect Response that redirects to LiteLLM UI

/sso/debug/login - handles user signing in with SSO and redirects to /sso/debug/callback
/sso/debug/callback - returns the OpenID object returned by the SSO provider
"""

import asyncio
import base64
import hashlib
import inspect
import os
import re
import secrets
from copy import deepcopy
from html import escape
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    List,
    Literal,
    NoReturn,
    Optional,
    Tuple,
    Union,
    cast,
)
from urllib.parse import parse_qs, urlencode, urlparse

if TYPE_CHECKING:
    import httpx

import jwt
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from fastapi.responses import RedirectResponse

import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
    CLI_SSO_CLAIM_MAP,
    CLI_SSO_CLAIM_MAX_SCALAR_LENGTH,
    CLI_SSO_SESSION_CACHE_KEY_PREFIX,
    CLI_SSO_SESSION_TTL_SECONDS,
    LITELLM_CLI_SOURCE_IDENTIFIER,
    LITELLM_UI_SESSION_DURATION,
    MAX_SPENDLOG_ROWS_TO_QUERY,
    MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE,
    MICROSOFT_USER_EMAIL_ATTRIBUTE,
    MICROSOFT_USER_FIRST_NAME_ATTRIBUTE,
    MICROSOFT_USER_ID_ATTRIBUTE,
    MICROSOFT_USER_LAST_NAME_ATTRIBUTE,
)
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
from litellm.llms.custom_httpx.http_handler import (
    AsyncHTTPHandler,
    get_async_httpx_client,
    httpxSpecialProvider,
)
from litellm.proxy._types import (
    CommonProxyErrors,
    LiteLLM_UserTable,
    LitellmUserRoles,
    Member,
    NewTeamRequest,
    NewUserRequest,
    NewUserResponse,
    ProxyErrorTypes,
    ProxyException,
    SSOUserDefinedValues,
    TeamMemberAddRequest,
    UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object
from litellm.proxy.auth.auth_utils import (
    _get_request_ip_address,
    _has_user_setup_sso,
)
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.admin_ui_utils import (
    admin_ui_disabled,
    show_missing_vars_in_env,
)
from litellm.proxy.common_utils.html_forms.jwt_display_template import (
    jwt_display_template,
)
from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO
from litellm.proxy.management_endpoints.sso_helper_utils import (
    check_is_admin_only_access,
    has_admin_ui_access,
)
from litellm.proxy.management_endpoints.team_endpoints import new_team, team_member_add
from litellm.proxy.management_endpoints.types import (
    CustomOpenID,
    get_litellm_user_role,
    is_valid_litellm_user_role,
)
from litellm.proxy.utils import (
    PrismaClient,
    ProxyLogging,
    get_custom_url,
    get_server_root_path,
)
from litellm.repositories.table_repositories import SSOConfigRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.secret_managers.main import get_secret_bool, str_to_bool
from litellm.types.proxy.management_endpoints.ui_sso import *  # noqa: F403
from litellm.types.proxy.management_endpoints.ui_sso import (
    DefaultTeamSSOParams,
    MicrosoftGraphAPIUserGroupDirectoryObject,
    MicrosoftGraphAPIUserGroupResponse,
    MicrosoftServicePrincipalTeam,
    RoleMappings,
    TeamMappings,
)
from litellm.types.proxy.ui_sso import ParsedOpenIDResult

if TYPE_CHECKING:
    from fastapi_sso.sso.base import OpenID
else:
    from typing import Any as OpenID

router = APIRouter()

# OAuth bearer credential fields that must not appear in SSO debug responses
# (received_response is included in restricted-group error messages).
# Metadata fields (token_type, expires_in, scope) are intentionally kept so
# response convertors see the same fields in the PKCE path as in the non-PKCE path.
_OAUTH_TOKEN_FIELDS = frozenset({"access_token", "id_token", "refresh_token"})
_CLI_SSO_FLOW_CACHE_KEY_PREFIX = f"{CLI_SSO_SESSION_CACHE_KEY_PREFIX}:flow"
_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX = (
    f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:start_rate_limit"
)
_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS = 60
_CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS = 30
_CLI_SSO_USER_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
_CLI_SSO_LOGIN_ID_RE = re.compile(r"^cli-[A-Za-z0-9_-]{12,124}$")
_CLI_SSO_USER_CODE_RE = re.compile(
    rf"^[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}-[{_CLI_SSO_USER_CODE_ALPHABET}]{{4}}$"
)
_CLI_SSO_SCALAR_TYPES = (str, int, float, bool)
_CLI_SSO_DEST_KEY_RE = re.compile(r"^[A-Za-z0-9_.-]+$")
_CLI_SSO_SECRET_KEY_FRAGMENTS = frozenset(
    {
        "access_token",
        "api_key",
        "client_secret",
        "id_token",
        "password",
        "private_key",
        "refresh_token",
        "secret",
    }
)


def _hash_cli_sso_secret(secret: str) -> str:
    return hashlib.sha256(secret.encode("utf-8")).hexdigest()


def _normalize_cli_sso_user_code(user_code: str) -> str:
    return "".join(ch for ch in user_code.upper() if ch.isalnum())


def _generate_cli_sso_user_code() -> str:
    user_code = "".join(secrets.choice(_CLI_SSO_USER_CODE_ALPHABET) for _ in range(8))
    return f"{user_code[:4]}-{user_code[4:]}"


def _get_cli_sso_flow_cache_key(login_id: str) -> str:
    return f"{_CLI_SSO_FLOW_CACHE_KEY_PREFIX}:{login_id}"


def _is_valid_cli_sso_login_id(login_id: Optional[str]) -> bool:
    return isinstance(login_id, str) and bool(_CLI_SSO_LOGIN_ID_RE.fullmatch(login_id))


def _is_valid_cli_sso_user_code(user_code: str | None) -> bool:
    return isinstance(user_code, str) and bool(
        _CLI_SSO_USER_CODE_RE.fullmatch(user_code)
    )


def _cli_sso_verification_uri_complete_enabled() -> bool:
    from litellm.proxy.proxy_server import general_settings

    return bool(general_settings.get("allow_cli_sso_verification_uri_complete", False))


def _cli_sso_start_response_body(
    *,
    login_id: str,
    poll_secret: str,
    user_code: str,
    verification_uri_complete: str | None,
) -> dict[str, str | int]:
    if verification_uri_complete is None:
        return {
            "login_id": login_id,
            "poll_secret": poll_secret,
            "user_code": user_code,
            "expires_in": CLI_SSO_SESSION_TTL_SECONDS,
        }
    return {
        "login_id": login_id,
        "poll_secret": poll_secret,
        "user_code": user_code,
        "verification_uri_complete": verification_uri_complete,
        "expires_in": CLI_SSO_SESSION_TTL_SECONDS,
    }


def _get_cli_sso_start_rate_limit_cache_key(
    request: Request, use_x_forwarded_for: Optional[bool] = False
) -> str:
    client_ip = (
        _get_request_ip_address(
            request=request, use_x_forwarded_for=use_x_forwarded_for
        )
        or "unknown"
    )
    client_ip_hash = _hash_cli_sso_secret(client_ip)
    return f"{_CLI_SSO_START_RATE_LIMIT_CACHE_KEY_PREFIX}:{client_ip_hash}"


def _check_cli_sso_start_rate_limit(
    request: Request,
    cache: DualCache,
    use_x_forwarded_for: Optional[bool] = False,
) -> None:
    rate_limit_cache_key = _get_cli_sso_start_rate_limit_cache_key(
        request=request, use_x_forwarded_for=use_x_forwarded_for
    )
    current_attempts = cache.increment_cache(
        key=rate_limit_cache_key,
        value=1,
        ttl=_CLI_SSO_START_RATE_LIMIT_WINDOW_SECONDS,
    )
    if current_attempts > _CLI_SSO_START_RATE_LIMIT_MAX_ATTEMPTS:
        raise HTTPException(
            status_code=429,
            detail="Too many CLI login attempts. Try again later.",
        )


def _get_cli_sso_flow_or_raise(login_id: Optional[str], cache: DualCache) -> dict:
    if not _is_valid_cli_sso_login_id(login_id):
        raise HTTPException(status_code=400, detail="Invalid CLI login session")

    cache_key = _get_cli_sso_flow_cache_key(cast(str, login_id))
    flow = cache.get_cache(key=cache_key)
    if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
        raise HTTPException(status_code=400, detail="Invalid CLI login session")
    return flow


def _set_cli_sso_flow(login_id: str, cache: DualCache, flow: dict) -> None:
    cache.set_cache(
        key=_get_cli_sso_flow_cache_key(login_id),
        value=flow,
        ttl=CLI_SSO_SESSION_TTL_SECONDS,
    )


def _verify_cli_sso_poll_secret(flow: dict, poll_secret: Optional[str]) -> bool:
    expected_poll_secret_hash = flow.get("poll_secret_hash")
    if not isinstance(expected_poll_secret_hash, str) or not isinstance(
        poll_secret, str
    ):
        return False
    supplied_poll_secret_hash = _hash_cli_sso_secret(poll_secret)
    return secrets.compare_digest(supplied_poll_secret_hash, expected_poll_secret_hash)


def _parse_cli_sso_claim_map() -> List[Tuple[str, str]]:
    """
    Parse CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP.

    Format: comma-separated ``source_claim->metadata_key`` pairs, e.g.
    ``employment_type->acme_employment_type,org_info.department->department``.
    Destination keys may use an optional ``metadata.`` prefix; values are stored
    on the LiteLLM user's ``metadata`` JSON column.
    """
    claim_map_raw = CLI_SSO_CLAIM_MAP.strip()
    if not claim_map_raw:
        return []

    parsed: List[Tuple[str, str]] = []
    for entry in claim_map_raw.split(","):
        entry = entry.strip()
        if not entry or "->" not in entry:
            continue
        source_claim, dest_key = entry.split("->", 1)
        source_claim = source_claim.strip()
        dest_key = dest_key.strip()
        if dest_key.startswith("metadata."):
            dest_key = dest_key[len("metadata.") :]
        if source_claim and dest_key:
            parsed.append((source_claim, dest_key))
    return parsed


def _is_safe_cli_sso_metadata_dest_key(dest_key: str) -> bool:
    if not dest_key or not _CLI_SSO_DEST_KEY_RE.fullmatch(dest_key):
        return False
    lowered = dest_key.lower()
    return not any(fragment in lowered for fragment in _CLI_SSO_SECRET_KEY_FRAGMENTS)


def _is_safe_cli_sso_scalar_claim_value(value: Any) -> bool:
    if not isinstance(value, _CLI_SSO_SCALAR_TYPES):
        return False
    if isinstance(value, str):
        if len(value) > CLI_SSO_CLAIM_MAX_SCALAR_LENGTH:
            return False
        if value.startswith("eyJ") and value.count(".") >= 2:
            return False
    return True


def _sso_result_to_dict(result: Union[CustomOpenID, OpenID, dict]) -> Dict[str, Any]:
    if isinstance(result, dict):
        return result
    if hasattr(result, "model_dump"):
        dumped = result.model_dump()
        if isinstance(dumped, dict):
            return cast(Dict[str, Any], dumped)
    return {}


def _get_nested_claim_value(data: Dict[str, Any], claim_path: str) -> Any:
    """Resolve a dot-notation claim path against an SSO result dict.

    Unlike ``get_nested_value``, this does not strip a leading ``metadata.``
    prefix, since OIDC claims may legitimately use ``metadata`` as a top-level
    key.
    """
    if not claim_path:
        return None
    if claim_path in data:
        return data[claim_path]
    placeholder = "\x00"
    parts = claim_path.replace("\\.", placeholder).split(".")
    parts = [p.replace(placeholder, ".") for p in parts]
    current: Any = data
    for part in parts:
        if isinstance(current, dict) and part in current:
            current = current[part]
        else:
            return None
    return current


def _extract_sso_claim_value(
    result: Union[CustomOpenID, OpenID, dict], claim_path: str
) -> Any:
    extra_fields = getattr(result, "extra_fields", None)
    if isinstance(extra_fields, dict):
        if claim_path in extra_fields:
            return extra_fields[claim_path]
        nested = _get_nested_claim_value(extra_fields, claim_path)
        if nested is not None:
            return nested

    if isinstance(result, dict):
        return _get_nested_claim_value(result, claim_path)

    result_dict = _sso_result_to_dict(result)
    return _get_nested_claim_value(result_dict, claim_path)


def _set_nested_metadata_value(
    metadata: Dict[str, Any], key_path: str, value: Any
) -> None:
    placeholder = "\x00"
    parts = key_path.replace("\\.", placeholder).split(".")
    parts = [p.replace(placeholder, ".") for p in parts]
    current: Any = metadata
    for part in parts[:-1]:
        existing = current.get(part)
        if not isinstance(existing, dict):
            existing = {}
            current[part] = existing
        current = existing
    current[parts[-1]] = value


def _flatten_cli_sso_metadata_for_poll(
    metadata: Dict[str, Any],
) -> Dict[str, Union[str, int, float, bool]]:
    """Expose scalar attribution metadata as a flat dict for CLI poll responses."""
    flattened: Dict[str, Union[str, int, float, bool]] = {}
    stack: List[Tuple[str, Any]] = [("", metadata)]
    while stack:
        prefix, value = stack.pop()
        if isinstance(value, dict):
            for key, nested in value.items():
                nested_prefix = f"{prefix}.{key}" if prefix else key
                stack.append((nested_prefix, nested))
        elif _is_safe_cli_sso_scalar_claim_value(value):
            flattened[prefix] = value
    return flattened


def build_cli_sso_attribution_metadata(
    result: Union[CustomOpenID, OpenID, dict],
) -> Dict[str, Any]:
    """
    Build allowlisted, non-secret scalar attribution metadata from an SSO result.

    Sources are configured via CLI_SSO_CLAIM_MAP / LITELLM_CLI_SSO_CLAIM_MAP and
    may include claims captured by GENERIC_USER_EXTRA_ATTRIBUTES on CustomOpenID.
    """
    claim_map = _parse_cli_sso_claim_map()
    if not claim_map:
        return {}

    metadata: Dict[str, Any] = {}
    for source_claim, dest_key in claim_map:
        if not _is_safe_cli_sso_metadata_dest_key(dest_key):
            verbose_proxy_logger.debug(
                f"Skipping unsafe CLI SSO metadata destination key: {dest_key}"
            )
            continue

        raw_value = _extract_sso_claim_value(result=result, claim_path=source_claim)
        if not _is_safe_cli_sso_scalar_claim_value(raw_value):
            continue

        _set_nested_metadata_value(
            metadata=metadata, key_path=dest_key, value=raw_value
        )

    return metadata


def _merge_cli_sso_attribution_metadata(
    existing_metadata: Dict[str, Any], attribution_metadata: Dict[str, Any]
) -> Dict[str, Any]:
    """Merge attribution metadata into existing user metadata in-place.

    Preserves original value types (in particular, string claim values that
    happen to look numeric are NOT coerced to ``int``/``float``). Nested dicts
    are merged iteratively so attribution claims do not clobber unrelated keys
    under the same parent.
    """
    pending: List[Tuple[Dict[str, Any], Dict[str, Any]]] = [
        (existing_metadata, attribution_metadata)
    ]
    while pending:
        target, source = pending.pop()
        for key, value in source.items():
            if value is None:
                continue
            existing_value = target.get(key)
            if isinstance(value, dict) and isinstance(existing_value, dict):
                pending.append((existing_value, value))
            else:
                target[key] = value
    return existing_metadata


async def _persist_cli_sso_user_metadata(
    prisma_client: PrismaClient,
    user_id: str,
    attribution_metadata: Dict[str, Any],
) -> None:
    if not attribution_metadata:
        return

    try:
        user_row = await UserRepository(prisma_client).table.find_unique(
            where={"user_id": user_id}
        )
        existing_metadata: Dict[str, Any] = {}
        if user_row is not None:
            row_metadata = user_row.metadata
            if isinstance(row_metadata, dict):
                existing_metadata = deepcopy(row_metadata)

        merged_metadata = _merge_cli_sso_attribution_metadata(
            existing_metadata=existing_metadata,
            attribution_metadata=attribution_metadata,
        )
        await UserRepository(prisma_client).table.update_many(
            where={"user_id": user_id},
            data={"metadata": merged_metadata},
        )
        verbose_proxy_logger.info(
            f"Persisted CLI SSO attribution metadata for user {user_id}: "
            f"{list(_flatten_cli_sso_metadata_for_poll(attribution_metadata).keys())}"
        )
    except Exception as e:
        verbose_proxy_logger.error(
            f"Failed to persist CLI SSO attribution metadata for user {user_id}: {e}"
        )


def _cli_poll_attribution_metadata_from_session(
    session_data: Dict[str, Any],
) -> Dict[str, Union[str, int, float, bool]]:
    stored = session_data.get("attribution_metadata")
    if isinstance(stored, dict):
        return _flatten_cli_sso_metadata_for_poll(stored)
    return {}


def _render_cli_sso_verification_page(
    verify_url: str,
    browser_complete_token: str,
    prefill_user_code: str | None = None,
) -> str:
    escaped_verify_url = escape(verify_url, quote=True)
    escaped_browser_complete_token = escape(browser_complete_token, quote=True)
    user_code_value_attr = (
        f' value="{escape(prefill_user_code, quote=True)}"' if prefill_user_code else ""
    )
    instructions = (
        "Confirm the verification code below to finish this login."
        if prefill_user_code
        else "Enter the verification code shown in your terminal to finish this login."
    )
    return f"""
    <!doctype html>
    <html>
      <head>
        <title>LiteLLM CLI Login</title>
        <style>
          body {{
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
            margin: 0;
            min-height: 100vh;
            display: flex;
            align-items: center;
            justify-content: center;
            background: #f8fafc;
            color: #0f172a;
          }}
          main {{
            width: min(420px, calc(100vw - 32px));
            background: #ffffff;
            border: 1px solid #e2e8f0;
            border-radius: 8px;
            padding: 28px;
            box-shadow: 0 12px 32px rgba(15, 23, 42, 0.08);
          }}
          h1 {{ font-size: 22px; margin: 0 0 12px; }}
          p {{ line-height: 1.5; margin: 0 0 18px; color: #334155; }}
          label {{ display: block; font-weight: 600; margin-bottom: 8px; }}
          input {{
            box-sizing: border-box;
            width: 100%;
            padding: 12px;
            border: 1px solid #cbd5e1;
            border-radius: 6px;
            font-size: 20px;
            letter-spacing: 0.08em;
            text-transform: uppercase;
          }}
          button {{
            margin-top: 16px;
            width: 100%;
            padding: 12px;
            border: 0;
            border-radius: 6px;
            background: #0f172a;
            color: #ffffff;
            font-weight: 600;
            cursor: pointer;
          }}
        </style>
      </head>
      <body>
        <main>
          <h1>Complete CLI Login</h1>
          <p>{instructions}</p>
          <form method="post" action="{escaped_verify_url}">
            <input type="hidden" name="browser_complete_token" value="{escaped_browser_complete_token}" />
            <label for="user_code">Verification code</label>
            <input id="user_code" name="user_code" autocomplete="one-time-code"{user_code_value_attr} required autofocus />
            <button type="submit">Continue</button>
          </form>
        </main>
      </body>
    </html>
    """


@router.post("/sso/cli/start", tags=["experimental"], include_in_schema=False)
async def cli_sso_start(request: Request):
    from litellm.proxy.proxy_server import general_settings, user_api_key_cache

    _check_cli_sso_start_rate_limit(
        request=request,
        cache=user_api_key_cache,
        use_x_forwarded_for=bool(
            (general_settings or {}).get("use_x_forwarded_for", False)
        ),
    )

    login_id = f"cli-{secrets.token_urlsafe(24)}"
    poll_secret = secrets.token_urlsafe(32)
    user_code = _generate_cli_sso_user_code()

    flow = {
        "poll_secret_hash": _hash_cli_sso_secret(poll_secret),
        "user_code_hash": _hash_cli_sso_secret(_normalize_cli_sso_user_code(user_code)),
        "sso_complete": False,
        "user_code_verified": False,
        "session_data": None,
    }
    _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)

    verification_uri_complete: str | None = (
        (
            get_custom_url(
                request_base_url=str(request.base_url), route="sso/key/generate"
            )
            + "?"
            + urlencode(
                {
                    "source": LITELLM_CLI_SOURCE_IDENTIFIER,
                    "key": login_id,
                    "user_code": user_code,
                }
            )
        )
        if _cli_sso_verification_uri_complete_enabled()
        else None
    )
    return _cli_sso_start_response_body(
        login_id=login_id,
        poll_secret=poll_secret,
        user_code=user_code,
        verification_uri_complete=verification_uri_complete,
    )


@router.post(
    "/sso/cli/complete/{login_id}", tags=["experimental"], include_in_schema=False
)
async def cli_sso_complete(request: Request, login_id: str):
    from fastapi.responses import HTMLResponse

    from litellm.proxy.common_utils.html_forms.cli_sso_success import (
        render_cli_sso_success_page,
    )
    from litellm.proxy.proxy_server import user_api_key_cache

    flow = _get_cli_sso_flow_or_raise(login_id=login_id, cache=user_api_key_cache)
    if not flow.get("sso_complete") or not flow.get("session_data"):
        raise HTTPException(status_code=400, detail="CLI login is not ready")

    body = (await request.body()).decode("utf-8")
    form_values = parse_qs(body)
    supplied_user_code = (form_values.get("user_code") or [""])[0]
    supplied_browser_complete_token = (
        form_values.get("browser_complete_token") or [""]
    )[0]
    supplied_user_code_hash = _hash_cli_sso_secret(
        _normalize_cli_sso_user_code(supplied_user_code)
    )
    supplied_browser_complete_token_hash = _hash_cli_sso_secret(
        supplied_browser_complete_token
    )

    expected_user_code_hash = flow.get("user_code_hash")
    if not isinstance(expected_user_code_hash, str) or not secrets.compare_digest(
        supplied_user_code_hash, expected_user_code_hash
    ):
        raise HTTPException(status_code=400, detail="Invalid verification code")

    expected_browser_complete_token_hash = flow.get("browser_complete_token_hash")
    if not isinstance(
        expected_browser_complete_token_hash, str
    ) or not secrets.compare_digest(
        supplied_browser_complete_token_hash, expected_browser_complete_token_hash
    ):
        raise HTTPException(status_code=400, detail="Invalid verification code")

    flow["user_code_verified"] = True
    _set_cli_sso_flow(login_id=login_id, cache=user_api_key_cache, flow=flow)

    html_content = render_cli_sso_success_page()
    return HTMLResponse(content=html_content, status_code=200)


def normalize_email(email: Optional[str]) -> Optional[str]:
    """
    Normalize email address to lowercase for consistent storage and comparison.

    Email addresses should be treated as case-insensitive for SSO purposes,
    even though RFC 5321 technically allows case-sensitive local parts.
    This prevents issues where SSO providers return emails with different casing
    than what's stored in the database.

    Args:
        email: Email address to normalize, can be None

    Returns:
        Lowercased email address, or None if input is None
    """
    if email is None:
        return None
    return email.lower() if isinstance(email, str) else email


def determine_role_from_groups(
    user_groups: List[str],
    role_mappings: "RoleMappings",
) -> Optional[LitellmUserRoles]:
    """
    Determine the highest privilege role for a user based on their groups.

    Role hierarchy (highest to lowest):
    - proxy_admin
    - proxy_admin_viewer
    - internal_user
    - internal_user_viewer

    Args:
        user_groups: List of group names from the SSO token
        role_mappings: RoleMappings configuration object

    Returns:
        The highest privilege role found, or default_role if no matches, or None
    """
    if not role_mappings.roles:
        # No role mappings configured, return default_role
        return role_mappings.default_role

    # Role hierarchy (highest to lowest)
    role_hierarchy = [
        LitellmUserRoles.PROXY_ADMIN,
        LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
        LitellmUserRoles.INTERNAL_USER,
        LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
    ]

    # Convert user_groups to a set for efficient lookup
    user_groups_set = set(user_groups) if isinstance(user_groups, list) else set()

    # Find the highest privilege role the user belongs to
    for role in role_hierarchy:
        if role in role_mappings.roles:
            role_groups = role_mappings.roles[role]
            if isinstance(role_groups, list) and user_groups_set.intersection(
                set(role_groups)
            ):
                verbose_proxy_logger.debug(
                    f"User groups {user_groups} matched role '{role.value}' via groups: {role_groups}"
                )
                return role

    # No matching groups found, return default_role
    verbose_proxy_logger.debug(
        f"User groups {user_groups} did not match any role mappings, using default_role: {role_mappings.default_role}"
    )
    return role_mappings.default_role


def process_sso_jwt_access_token(
    access_token_str: Optional[str],
    sso_jwt_handler: Optional[JWTHandler],
    result: Union[OpenID, dict, None],
    role_mappings: Optional["RoleMappings"] = None,
) -> Optional[dict]:
    """
    Process SSO JWT access token and extract team IDs and user role if available.

    This function decodes the JWT access token and extracts team IDs and user
    role, then sets them on the result object. Role extraction from the access
    token is needed because some SSO providers (e.g., Keycloak) do not include
    role claims in the UserInfo endpoint response.

    Args:
        access_token_str: The JWT access token string
        sso_jwt_handler: SSO-specific JWT handler for team ID extraction
        result: The SSO result object to update with team IDs and role
        role_mappings: Optional role mappings configuration for group-based role determination

    Returns:
        The decoded access token payload dict, or None if decoding failed or
        inputs were missing. Callers can pass this to _sync_user_role_from_jwt_role_map
        so it has access to custom role claims (e.g. custom_roles) that are
        encoded inside the JWT but stripped from received_response.
    """
    if access_token_str and result:
        import jwt

        try:
            access_token_payload = jwt.decode(
                access_token_str, options={"verify_signature": False}
            )
        except jwt.exceptions.DecodeError:
            verbose_proxy_logger.debug(
                "Access token is not a valid JWT (possibly an opaque token), skipping JWT-based extraction"
            )
            return None

        # Extract team IDs from access token if sso_jwt_handler is available
        if sso_jwt_handler:
            if isinstance(result, dict):
                result_team_ids: Optional[List[str]] = result.get("team_ids", [])
                if not result_team_ids:
                    team_ids = sso_jwt_handler.get_team_ids_from_jwt(
                        access_token_payload
                    )
                    result["team_ids"] = team_ids
            else:
                result_team_ids = getattr(result, "team_ids", []) if result else []
                if not result_team_ids:
                    team_ids = sso_jwt_handler.get_team_ids_from_jwt(
                        access_token_payload
                    )
                    setattr(result, "team_ids", team_ids)

        # Extract user role from access token if not already set from UserInfo
        existing_role = (
            result.get("user_role")
            if isinstance(result, dict)
            else getattr(result, "user_role", None)
        )
        if existing_role is None:
            user_role: Optional[LitellmUserRoles] = None

            # Try role_mappings first (group-based role determination)
            if role_mappings is not None and role_mappings.roles:
                group_claim = role_mappings.group_claim
                user_groups_raw: Any = get_nested_value(
                    access_token_payload, group_claim
                )

                user_groups: List[str] = []
                if isinstance(user_groups_raw, list):
                    user_groups = [str(g) for g in user_groups_raw]
                elif isinstance(user_groups_raw, str):
                    user_groups = [
                        g.strip() for g in user_groups_raw.split(",") if g.strip()
                    ]
                elif user_groups_raw is not None:
                    user_groups = [str(user_groups_raw)]

                if user_groups:
                    user_role = determine_role_from_groups(user_groups, role_mappings)
                    verbose_proxy_logger.debug(
                        f"Determined role '{user_role}' from access token groups '{user_groups}' using role_mappings"
                    )
                elif role_mappings.default_role:
                    user_role = role_mappings.default_role

            # Fallback: try GENERIC_USER_ROLE_ATTRIBUTE on the access token payload
            if user_role is None:
                generic_user_role_attribute_name = os.getenv(
                    "GENERIC_USER_ROLE_ATTRIBUTE", "role"
                )
                user_role_from_token = get_nested_value(
                    access_token_payload, generic_user_role_attribute_name
                )
                if user_role_from_token is not None:
                    user_role = get_litellm_user_role(user_role_from_token)
                    verbose_proxy_logger.debug(
                        f"Extracted role '{user_role}' from access token field '{generic_user_role_attribute_name}'"
                    )

            if user_role is not None:
                if isinstance(result, dict):
                    result["user_role"] = user_role
                else:
                    setattr(result, "user_role", user_role)
                verbose_proxy_logger.debug(
                    f"Set user_role='{user_role}' from JWT access token"
                )

        return access_token_payload

    return None


@router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False)
async def google_login(
    request: Request,
    source: Optional[str] = None,
    key: Optional[str] = None,
    existing_key: Optional[str] = None,
    return_to: Optional[str] = None,
    user_code: str | None = None,
):
    """
    Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
    PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
    Example:
    """
    from litellm.proxy.proxy_server import (
        general_settings,
        premium_user,
        prisma_client,
        user_api_key_cache,
        user_custom_ui_sso_sign_in_handler,
    )

    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
    generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)

    ####### Check if UI is disabled #######
    _disable_ui_flag = os.getenv("DISABLE_ADMIN_UI")
    if _disable_ui_flag is not None:
        is_disabled = str_to_bool(value=_disable_ui_flag)
        if is_disabled:
            return admin_ui_disabled()

    ####### Check if user is a Enterprise / Premium User #######
    if (
        microsoft_client_id is not None
        or google_client_id is not None
        or generic_client_id is not None
    ):
        if premium_user is not True:
            # Check if under 'free SSO user' limit
            if prisma_client is not None:
                total_users = await UserRepository(prisma_client).table.count()
                if total_users and total_users > 5:
                    raise ProxyException(
                        message="You must be a LiteLLM Enterprise user to use SSO for more than 5 users. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this",
                        type=ProxyErrorTypes.auth_error,
                        param="premium_user",
                        code=status.HTTP_403_FORBIDDEN,
                    )
            else:
                raise ProxyException(
                    message=CommonProxyErrors.db_not_connected_error.value,
                    type=ProxyErrorTypes.auth_error,
                    param="premium_user",
                    code=status.HTTP_403_FORBIDDEN,
                )

    ####### Detect DB + MASTER KEY in .env #######
    missing_env_vars = show_missing_vars_in_env()
    if missing_env_vars is not None:
        return missing_env_vars

    # get url from request - always use regular callback, but set state for CLI
    redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
        request=request,
        sso_callback_route="sso/callback",
    )

    if source == LITELLM_CLI_SOURCE_IDENTIFIER:
        _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)

    # Store CLI login handle in state for OAuth flow
    cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state(
        source=source,
        key=key,
        user_code=(user_code if _cli_sso_verification_uri_complete_enabled() else None),
    )

    # check if user defined a custom auth sso sign in handler, if yes, use it
    if user_custom_ui_sso_sign_in_handler is not None:
        try:
            from litellm_enterprise.proxy.auth.custom_sso_handler import (  # type: ignore[import-untyped]
                EnterpriseCustomSSOHandler,
            )

            return await EnterpriseCustomSSOHandler.handle_custom_ui_sso_sign_in(
                request=request,
            )
        except ImportError:
            raise ValueError(
                "Enterprise features are not available. Custom UI SSO sign-in requires LiteLLM Enterprise."
            )

    # Check if we should use SSO handler
    if (
        SSOAuthenticationHandler.should_use_sso_handler(
            microsoft_client_id=microsoft_client_id,
            google_client_id=google_client_id,
            generic_client_id=generic_client_id,
        )
        is True
    ):
        verbose_proxy_logger.info(f"Redirecting to SSO login for {redirect_url}")
        sso_redirect = await SSOAuthenticationHandler.get_sso_login_redirect(
            redirect_url=redirect_url,
            microsoft_client_id=microsoft_client_id,
            google_client_id=google_client_id,
            generic_client_id=generic_client_id,
            state=cli_state,
            request=request,
        )
        if return_to is not None and sso_redirect is not None:
            if SSOAuthenticationHandler._validate_return_to(return_to):
                sso_redirect.set_cookie(
                    key="litellm_cp_return_to",
                    value=return_to,
                    max_age=600,
                    httponly=True,
                    samesite="lax",
                )
        return sso_redirect

    from fastapi.responses import HTMLResponse

    hide_default_credentials_hint = (
        os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
        or general_settings.get("hide_default_credentials_hint", False) is True
    )
    return HTMLResponse(
        content=build_ui_login_form(
            show_deprecation_banner=True,
            hide_default_credentials_hint=hide_default_credentials_hint,
        ),
        status_code=200,
    )


def generic_response_convertor(
    response,
    jwt_handler: JWTHandler,
    sso_jwt_handler: Optional[JWTHandler] = None,
    role_mappings: Optional["RoleMappings"] = None,
    team_mappings: Optional["TeamMappings"] = None,
) -> CustomOpenID:
    generic_user_id_attribute_name = os.getenv(
        "GENERIC_USER_ID_ATTRIBUTE", "preferred_username"
    )
    generic_user_display_name_attribute_name = os.getenv(
        "GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub"
    )
    generic_user_email_attribute_name = os.getenv(
        "GENERIC_USER_EMAIL_ATTRIBUTE", "email"
    )

    generic_user_first_name_attribute_name = os.getenv(
        "GENERIC_USER_FIRST_NAME_ATTRIBUTE", "first_name"
    )
    generic_user_last_name_attribute_name = os.getenv(
        "GENERIC_USER_LAST_NAME_ATTRIBUTE", "last_name"
    )

    generic_provider_attribute_name = os.getenv(
        "GENERIC_USER_PROVIDER_ATTRIBUTE", "provider"
    )

    generic_user_role_attribute_name = os.getenv("GENERIC_USER_ROLE_ATTRIBUTE", "role")

    generic_user_extra_attributes = os.getenv("GENERIC_USER_EXTRA_ATTRIBUTES", None)

    verbose_proxy_logger.debug(
        f" generic_user_id_attribute_name: {generic_user_id_attribute_name}\n generic_user_email_attribute_name: {generic_user_email_attribute_name}"
    )

    all_teams = []
    if sso_jwt_handler is not None:
        team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response))
        all_teams.extend(team_ids)

    if team_mappings is not None and team_mappings.team_ids_jwt_field is not None:
        team_ids_from_db_mapping: Optional[List[str]] = get_nested_value(
            data=cast(dict, response),
            key_path=team_mappings.team_ids_jwt_field,
            default=[],
        )
        if team_ids_from_db_mapping:
            all_teams.extend(team_ids_from_db_mapping)
            verbose_proxy_logger.debug(
                f"Loaded team_ids from DB team_mappings.team_ids_jwt_field='{team_mappings.team_ids_jwt_field}': {team_ids_from_db_mapping}"
            )
    else:
        team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response))
        all_teams.extend(team_ids)

    # Determine user role based on role_mappings if available
    # Only apply role_mappings for GENERIC SSO provider
    user_role: Optional[LitellmUserRoles] = None

    if role_mappings is not None and role_mappings.provider.lower() in [
        "generic",
        "okta",
    ]:
        # Use role_mappings to determine role from groups
        group_claim = role_mappings.group_claim
        user_groups_raw: Any = get_nested_value(response, group_claim)

        # Handle different formats: could be a list, string (comma-separated), or single value
        user_groups: List[str] = []
        if isinstance(user_groups_raw, list):
            user_groups = [str(g) for g in user_groups_raw]
        elif isinstance(user_groups_raw, str):
            # Handle comma-separated string
            user_groups = [g.strip() for g in user_groups_raw.split(",") if g.strip()]
        elif user_groups_raw is not None:
            # Single value
            user_groups = [str(user_groups_raw)]

        if user_groups:
            user_role = determine_role_from_groups(user_groups, role_mappings)
            verbose_proxy_logger.debug(
                f"Determined role '{user_role.value if user_role else None}' from groups '{user_groups}' using role_mappings"
            )
        else:
            # No groups found, use default_role
            user_role = role_mappings.default_role
            verbose_proxy_logger.debug(
                f"No groups found in '{group_claim}', using default_role: {role_mappings.default_role}"
            )

    # Fallback to existing logic if role_mappings not used
    if user_role is None:
        user_role_from_sso = get_nested_value(
            response, generic_user_role_attribute_name
        )
        if user_role_from_sso is not None:
            role = get_litellm_user_role(user_role_from_sso)
            if role is not None:
                user_role = role
                verbose_proxy_logger.debug(
                    f"Found valid LitellmUserRoles '{role.value}' from SSO attribute '{generic_user_role_attribute_name}'"
                )

    # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified
    extra_fields: Optional[Dict[str, Any]] = None
    if generic_user_extra_attributes:
        extra_fields = {}
        for attr_name in generic_user_extra_attributes.split(","):
            attr_name = attr_name.strip()
            extra_fields[attr_name] = get_nested_value(response, attr_name)

    return CustomOpenID(
        id=get_nested_value(response, generic_user_id_attribute_name),
        display_name=get_nested_value(
            response, generic_user_display_name_attribute_name
        ),
        email=normalize_email(
            get_nested_value(response, generic_user_email_attribute_name)
        ),
        first_name=get_nested_value(response, generic_user_first_name_attribute_name),
        last_name=get_nested_value(response, generic_user_last_name_attribute_name),
        provider=get_nested_value(response, generic_provider_attribute_name),
        team_ids=all_teams,
        user_role=user_role,
        extra_fields=extra_fields,
    )


def _setup_generic_sso_env_vars(
    generic_client_id: str, redirect_url: str
) -> Tuple[str, List[str], str, str, str, bool]:
    """Setup and validate Generic SSO environment variables."""
    generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
    generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(" ")
    generic_authorization_endpoint = os.getenv("GENERIC_AUTHORIZATION_ENDPOINT", None)
    generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
    generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
    generic_include_client_id = (
        os.getenv("GENERIC_INCLUDE_CLIENT_ID", "false").lower() == "true"
    )

    # Validate required environment variables
    if generic_client_secret is None:
        raise ProxyException(
            message="GENERIC_CLIENT_SECRET not set. Set it in .env file",
            type=ProxyErrorTypes.auth_error,
            param="GENERIC_CLIENT_SECRET",
            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    if generic_authorization_endpoint is None:
        raise ProxyException(
            message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file",
            type=ProxyErrorTypes.auth_error,
            param="GENERIC_AUTHORIZATION_ENDPOINT",
            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    if generic_token_endpoint is None:
        raise ProxyException(
            message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file",
            type=ProxyErrorTypes.auth_error,
            param="GENERIC_TOKEN_ENDPOINT",
            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    if generic_userinfo_endpoint is None:
        raise ProxyException(
            message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file",
            type=ProxyErrorTypes.auth_error,
            param="GENERIC_USERINFO_ENDPOINT",
            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )

    verbose_proxy_logger.debug(
        f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}"
    )
    verbose_proxy_logger.debug(
        f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n"
    )

    return (
        generic_client_secret,
        generic_scope,
        generic_authorization_endpoint,
        generic_token_endpoint,
        generic_userinfo_endpoint,
        generic_include_client_id,
    )


async def _setup_team_mappings() -> Optional["TeamMappings"]:
    """Setup team mappings from SSO database settings."""
    team_mappings: Optional["TeamMappings"] = None
    try:
        from litellm.proxy.utils import get_prisma_client_or_throw

        prisma_client = get_prisma_client_or_throw(
            "Prisma client is None, connect a database to your proxy"
        )

        sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(
            where={"id": "sso_config"}
        )

        if sso_db_record and sso_db_record.sso_settings:
            sso_settings_dict = dict(sso_db_record.sso_settings)
            team_mappings_data = sso_settings_dict.get("team_mappings")

            if team_mappings_data:
                from litellm.types.proxy.management_endpoints.ui_sso import TeamMappings

                if isinstance(team_mappings_data, dict):
                    team_mappings = TeamMappings(**team_mappings_data)
                elif isinstance(team_mappings_data, TeamMappings):
                    team_mappings = team_mappings_data

                if team_mappings and team_mappings.team_ids_jwt_field:
                    verbose_proxy_logger.debug(
                        f"Loaded team_mappings with team_ids_jwt_field: '{team_mappings.team_ids_jwt_field}'"
                    )
    except Exception as e:
        verbose_proxy_logger.debug(
            f"Could not load team_mappings from database: {e}. Continuing with config-based team mapping."
        )

    return team_mappings


async def _setup_role_mappings() -> Optional["RoleMappings"]:
    """Setup role mappings from SSO database settings."""
    role_mappings: Optional["RoleMappings"] = None
    try:
        from litellm.proxy.utils import get_prisma_client_or_throw

        prisma_client = get_prisma_client_or_throw(
            "Prisma client is None, connect a database to your proxy"
        )

        sso_db_record = await SSOConfigRepository(prisma_client).table.find_unique(
            where={"id": "sso_config"}
        )

        if sso_db_record and sso_db_record.sso_settings:
            sso_settings_dict = dict(sso_db_record.sso_settings)
            role_mappings_data = sso_settings_dict.get("role_mappings")

            if role_mappings_data:
                from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings

                if isinstance(role_mappings_data, dict):
                    role_mappings = RoleMappings(**role_mappings_data)
                elif isinstance(role_mappings_data, RoleMappings):
                    role_mappings = role_mappings_data

                if role_mappings:
                    verbose_proxy_logger.debug(
                        f"Loaded role_mappings for provider '{role_mappings.provider}'"
                    )
    except Exception as e:
        verbose_proxy_logger.debug(
            f"Could not load role_mappings from database: {e}. Continuing with existing role logic."
        )

    generic_role_mappings = os.getenv("GENERIC_ROLE_MAPPINGS_ROLES", None)
    generic_role_mappings_group_claim = os.getenv(
        "GENERIC_ROLE_MAPPINGS_GROUP_CLAIM", None
    )
    generic_role_mappings_default_role = os.getenv(
        "GENERIC_ROLE_MAPPINGS_DEFAULT_ROLE", None
    )
    if generic_role_mappings is not None:
        verbose_proxy_logger.debug(
            "Found role_mappings for generic provider in environment variables"
        )
        import ast

        try:
            generic_user_role_mappings_data: Dict[LitellmUserRoles, List[str]] = (
                ast.literal_eval(generic_role_mappings)
            )
            if isinstance(generic_user_role_mappings_data, dict):
                from litellm.types.proxy.management_endpoints.ui_sso import RoleMappings

                role_mappings_data = {
                    "provider": "generic",
                    "group_claim": generic_role_mappings_group_claim,
                    "default_role": generic_role_mappings_default_role,
                    "roles": generic_user_role_mappings_data,
                }

                role_mappings = RoleMappings(**role_mappings_data)
                verbose_proxy_logger.debug(
                    f"Loaded role_mappings from environments for provider '{role_mappings.provider}'."
                )
                return role_mappings
        except TypeError as e:
            verbose_proxy_logger.warning(
                f"Error decoding role mappings from environment variables: {e}. Continuing with existing role logic."
            )
    return role_mappings


def _parse_generic_sso_headers() -> dict:
    """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict."""
    raw = os.getenv("GENERIC_SSO_HEADERS", None)
    if raw is None:
        return {}
    result: Dict[str, str] = {}
    for header in raw.split(","):
        header = header.strip()
        if header:
            key, value = header.split("=")
            result[key] = value
    return result


def _handle_generic_sso_error(
    e: Exception,
    generic_authorization_endpoint: Optional[str],
    generic_token_endpoint: Optional[str],
    additional_headers: dict,
) -> NoReturn:
    """Handle errors from generic SSO verify_and_process. Always re-raises."""
    error_message = str(e)

    # Surface a helpful PKCE misconfiguration hint only when:
    # 1. The error mentions PKCE/code verifier, AND
    # 2. PKCE is not currently configured (GENERIC_CLIENT_USE_PKCE != true)
    pkce_configured = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"
    if not pkce_configured and (
        "PKCE" in error_message or "code verifier" in error_message.lower()
    ):
        is_okta = (
            generic_authorization_endpoint
            and "okta" in generic_authorization_endpoint.lower()
        ) or (generic_token_endpoint and "okta" in generic_token_endpoint.lower())
        provider_name = "Okta" if is_okta else "Your OAuth provider"

        detailed_message = (
            f"SSO authentication failed: {provider_name} requires PKCE (Proof Key for Code Exchange) "
            f"but it's not enabled in your LiteLLM configuration.\n\n"
            f"SOLUTION: Add this environment variable and restart your proxy:\n"
            f"  GENERIC_CLIENT_USE_PKCE=true\n\n"
        )
        if is_okta:
            detailed_message += (
                "For AWS ECS: Add the environment variable to your task definition.\n"
                "For Docker: Add -e GENERIC_CLIENT_USE_PKCE=true to your docker run command.\n"
                "For .env file: Add GENERIC_CLIENT_USE_PKCE=true to your .env file.\n\n"
            )
        detailed_message += f"Original error: {error_message}"

        raise ProxyException(
            message=detailed_message,
            type=ProxyErrorTypes.auth_error,
            param="GENERIC_CLIENT_USE_PKCE",
            code=status.HTTP_401_UNAUTHORIZED,
        )

    if isinstance(e, ProxyException):
        verbose_proxy_logger.error(
            "SSO authentication failed: %s. Passed in headers: %s",
            e,
            additional_headers,
        )
    else:
        verbose_proxy_logger.exception(
            "Error verifying and processing generic SSO: %s. Passed in headers: %s",
            e,
            additional_headers,
        )
    raise e


async def get_generic_sso_response(
    request: Request,
    jwt_handler: JWTHandler,
    sso_jwt_handler: Optional[
        JWTHandler
    ],  # sso specific jwt handler - used for restricted sso group access control
    generic_client_id: str,
    redirect_url: str,
) -> Tuple[
    Union[OpenID, dict], Optional[dict], Optional[dict]
]:  # (result, received_response, access_token_payload)
    # make generic sso provider
    from fastapi_sso.sso.base import DiscoveryDocument
    from fastapi_sso.sso.generic import create_provider

    received_response: Optional[dict] = None

    # Setup environment variables
    (
        generic_client_secret,
        generic_scope,
        generic_authorization_endpoint,
        generic_token_endpoint,
        generic_userinfo_endpoint,
        generic_include_client_id,
    ) = _setup_generic_sso_env_vars(generic_client_id, redirect_url)

    discovery = DiscoveryDocument(
        authorization_endpoint=generic_authorization_endpoint,
        token_endpoint=generic_token_endpoint,
        userinfo_endpoint=generic_userinfo_endpoint,
    )

    role_mappings = await _setup_role_mappings()
    team_mappings = await _setup_team_mappings()

    def response_convertor(response, client):
        nonlocal received_response  # return for user debugging
        received_response = response
        return generic_response_convertor(
            response=response,
            jwt_handler=jwt_handler,
            sso_jwt_handler=sso_jwt_handler,
            role_mappings=role_mappings,
            team_mappings=team_mappings,
        )

    SSOProvider = create_provider(
        name="oidc",
        discovery_document=discovery,
        response_convertor=response_convertor,
    )
    generic_sso = SSOProvider(
        client_id=generic_client_id,
        client_secret=generic_client_secret,
        redirect_uri=redirect_url,
        allow_insecure_http=True,
        scope=generic_scope,
    )
    verbose_proxy_logger.debug("calling generic_sso.verify_and_process")
    additional_generic_sso_headers_dict = _parse_generic_sso_headers()

    code_verifier: Optional[str] = (
        None  # assigned inside try; initialized for type tracking
    )
    access_token_payload: Optional[dict] = None  # decoded JWT access token claims

    try:
        token_exchange_params = (
            await SSOAuthenticationHandler.prepare_token_exchange_parameters(
                request=request,
                generic_include_client_id=generic_include_client_id,
            )
        )

        # Extract code_verifier (and the cache key for deferred deletion) before calling fastapi-sso
        code_verifier = token_exchange_params.pop("code_verifier", None)
        pkce_cache_key = token_exchange_params.pop("_pkce_cache_key", None)

        # Get authorization code from query params (only used in the PKCE path below;
        # the non-PKCE path delegates to verify_and_process which handles OAuth error
        # callbacks — user-denied, CSRF mismatch — internally).
        authorization_code = request.query_params.get("code")

        if code_verifier:
            # State-to-session-cookie binding.  The non-PKCE branch below
            # delegates to fastapi-sso's ``verify_and_process``, which
            # performs its own session-cookie check.  The PKCE branch
            # bypasses that helper, so we validate the URL ``state``
            # against the ``litellm_oauth_state`` cookie set on the
            # redirect response — without this an attacker can pre-mint
            # a state + cached PKCE verifier and hijack a victim's auth
            # code (Login-CSRF / token theft).
            url_state = request.query_params.get("state")
            cookie_state = request.cookies.get("litellm_oauth_state")
            if (
                not url_state
                or not cookie_state
                or not secrets.compare_digest(url_state, cookie_state)
            ):
                raise ProxyException(
                    message=(
                        "Invalid OAuth state parameter — does not match "
                        "the browser-bound state cookie."
                    ),
                    type=ProxyErrorTypes.auth_error,
                    param="state",
                    code=status.HTTP_400_BAD_REQUEST,
                )
            if not authorization_code:
                raise ProxyException(
                    message="Missing authorization code in callback",
                    type=ProxyErrorTypes.auth_error,
                    param="code",
                    code=status.HTTP_400_BAD_REQUEST,
                )
            if not generic_client_id:
                raise ProxyException(
                    message="GENERIC_CLIENT_ID must be set when PKCE is enabled",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_CLIENT_ID",
                    code=status.HTTP_401_UNAUTHORIZED,
                )
            if not generic_token_endpoint:
                raise ProxyException(
                    message="GENERIC_TOKEN_ENDPOINT must be set when PKCE is enabled",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_TOKEN_ENDPOINT",
                    code=status.HTTP_401_UNAUTHORIZED,
                )
            # All guards above raise, so authorization_code is a non-empty str here.
            # Use an explicit type guard rather than assert (assert is a no-op with -O).
            if not isinstance(authorization_code, str):
                raise ProxyException(
                    message="Missing authorization code in callback",
                    type=ProxyErrorTypes.auth_error,
                    param="code",
                    code=status.HTTP_400_BAD_REQUEST,
                )
            combined_response = await SSOAuthenticationHandler._pkce_token_exchange(
                authorization_code=authorization_code,
                code_verifier=code_verifier,
                client_id=generic_client_id,
                client_secret=generic_client_secret,
                token_endpoint=generic_token_endpoint,
                userinfo_endpoint=generic_userinfo_endpoint,
                include_client_id=generic_include_client_id,
                redirect_url=redirect_url,
                additional_headers=additional_generic_sso_headers_dict,
            )
            # Pass the full response so custom response_convertor implementations
            # can access all fields (including id_token for claim extraction).
            result = response_convertor(combined_response, generic_sso)
            # Strip bearer credentials from combined_response before storing in
            # received_response. received_response may appear in restricted-group
            # error messages — bearer tokens (access_token, id_token, refresh_token)
            # must not be exposed to callers.
            # Assign directly rather than relying on nonlocal mutation so that Pyright
            # can track that received_response is non-None from this point on.
            received_response = {
                k: v
                for k, v in combined_response.items()
                if k not in _OAUTH_TOKEN_FIELDS
            }
            # In the PKCE path verify_and_process is skipped, so generic_sso.access_token
            # is never set. Read the token directly from the exchange response instead so
            # process_sso_jwt_access_token can extract JWT-embedded roles/teams.
            access_token_str: Optional[str] = combined_response.get("access_token")
        else:
            result = await generic_sso.verify_and_process(
                request,
                params=token_exchange_params,
                headers=additional_generic_sso_headers_dict,
            )
            access_token_str = generic_sso.access_token

        access_token_payload = process_sso_jwt_access_token(
            access_token_str, sso_jwt_handler, result, role_mappings=role_mappings
        )
        # Delete the single-use PKCE verifier only after all downstream processing
        # (response_convertor and process_sso_jwt_access_token) has completed
        # successfully.  Deleting earlier would consume the verifier on a transient
        # failure, forcing the user to restart the entire OAuth flow from scratch.
        if pkce_cache_key:
            await SSOAuthenticationHandler._delete_pkce_verifier(pkce_cache_key)

    except Exception as e:
        _handle_generic_sso_error(
            e,
            generic_authorization_endpoint,
            generic_token_endpoint,
            additional_generic_sso_headers_dict,
        )
    verbose_proxy_logger.debug("generic result: %s", result)
    return result or {}, received_response, access_token_payload


async def create_team_member_add_task(team_id, user_info):
    """Create a task for adding a member to a team."""
    try:
        member = Member(user_id=user_info.user_id, role="user")
        team_member_add_request = TeamMemberAddRequest(
            member=member,
            team_id=team_id,
        )
        return await team_member_add(
            data=team_member_add_request,
            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
        )
    except Exception as e:
        verbose_proxy_logger.debug(
            f"[Non-Blocking] Error trying to add sso user to db: {e}"
        )


async def add_missing_team_member(
    user_info: Union[NewUserResponse, LiteLLM_UserTable], sso_teams: List[str]
):
    """
    - Get missing teams (diff b/w user_info.team_ids and sso_teams)
    - Add missing user to missing teams
    """
    # Handle None as empty list for new users
    user_teams = user_info.teams if user_info.teams is not None else []
    missing_teams = set(sso_teams) - set(user_teams)
    missing_teams_list = list(missing_teams)
    tasks = []
    tasks = [
        create_team_member_add_task(team_id, user_info)
        for team_id in missing_teams_list
    ]

    try:
        await asyncio.gather(*tasks)
    except Exception as e:
        verbose_proxy_logger.debug(
            f"[Non-Blocking] Error trying to add sso user to db: {e}"
        )


def get_disabled_non_admin_personal_key_creation():
    key_generation_settings = litellm.key_generation_settings
    if key_generation_settings is None:
        return False
    personal_key_generation = (
        key_generation_settings.get("personal_key_generation") or {}
    )
    allowed_user_roles = personal_key_generation.get("allowed_user_roles") or []
    return bool("proxy_admin" in allowed_user_roles)


async def get_existing_user_info_from_db(
    user_id: Optional[str],
    user_email: Optional[str],
    prisma_client: PrismaClient,
    user_api_key_cache: UserApiKeyCache,
    proxy_logging_obj: ProxyLogging,
) -> Optional[LiteLLM_UserTable]:
    try:
        user_info = await get_user_object(
            user_id=user_id,
            user_email=user_email,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            user_id_upsert=False,
            parent_otel_span=None,
            proxy_logging_obj=proxy_logging_obj,
            sso_user_id=user_id,
        )
    except Exception as e:
        verbose_proxy_logger.debug(f"Error getting user object: {e}")
        user_info = None

    return user_info


async def get_user_info_from_db(
    result: Union[CustomOpenID, OpenID, dict],
    prisma_client: PrismaClient,
    user_api_key_cache: UserApiKeyCache,
    proxy_logging_obj: ProxyLogging,
    user_email: Optional[str],
    user_defined_values: Optional[SSOUserDefinedValues],
    alternate_user_id: Optional[str] = None,
) -> Optional[Union[LiteLLM_UserTable, NewUserResponse]]:
    try:
        potential_user_ids = []
        if alternate_user_id is not None:
            potential_user_ids.append(alternate_user_id)
        if not isinstance(result, dict):
            _id = getattr(result, "id", None)
            if _id is not None and isinstance(_id, str):
                potential_user_ids.append(_id)
        else:
            _id = result.get("id", None)
            if _id is not None and isinstance(_id, str):
                potential_user_ids.append(_id)

        user_email = normalize_email(
            getattr(result, "email", None)
            if not isinstance(result, dict)
            else result.get("email", None)
        )

        user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = None

        for user_id in potential_user_ids:
            user_info = await get_existing_user_info_from_db(
                user_id=user_id,
                user_email=user_email,
                prisma_client=prisma_client,
                user_api_key_cache=user_api_key_cache,
                proxy_logging_obj=proxy_logging_obj,
            )
            if user_info is not None:
                break

        verbose_proxy_logger.debug(
            f"user_info: {user_info}; litellm.default_internal_user_params: {litellm.default_internal_user_params}"
        )

        # Upsert SSO User to LiteLLM DB
        user_info = await SSOAuthenticationHandler.upsert_sso_user(
            result=result,
            user_info=user_info,
            user_email=user_email,
            user_defined_values=user_defined_values,
            prisma_client=prisma_client,
        )

        await SSOAuthenticationHandler.add_user_to_teams_from_sso_response(
            result=result,
            user_info=user_info,
        )

        return user_info
    except Exception as e:
        verbose_proxy_logger.exception(
            f"[Non-Blocking] Error trying to add sso user to db: {e}"
        )

    return None


def _should_use_role_from_sso_response(sso_role: Optional[str]) -> bool:
    """returns true if SSO upsert should use the 'role' defined on the SSO response"""
    if sso_role is None:
        return False

    if not is_valid_litellm_user_role(sso_role):
        verbose_proxy_logger.debug(
            f"SSO role '{sso_role}' is not a valid LiteLLM user role. "
            "Ignoring role from SSO response. See LitellmUserRoles enum for valid roles."
        )
        return False
    return True


def _build_sso_user_update_data(
    result: Optional[Union["CustomOpenID", OpenID, dict]],
    user_email: Optional[str],
    user_id: Optional[str],
) -> dict:
    """
    Build the update data dictionary for SSO user upsert.

    Args:
        result: The SSO response containing user information
        user_email: The user's email from SSO
        user_id: The user's ID for logging purposes

    Returns:
        dict: Update data containing user_email and optionally user_role if valid
    """
    update_data: dict = {"user_email": normalize_email(user_email)}

    # Get SSO role from result and include if valid
    sso_role = getattr(result, "user_role", None)
    if sso_role is not None:
        # Convert enum to string if needed
        sso_role_str = (
            sso_role.value if isinstance(sso_role, LitellmUserRoles) else sso_role
        )

        # Only include if it's a valid LiteLLM role
        if _should_use_role_from_sso_response(sso_role_str):
            update_data["user_role"] = sso_role_str
            verbose_proxy_logger.info(
                f"Updating user {user_id} role from SSO: {sso_role_str}"
            )

    return update_data


async def _sync_user_role_from_jwt_role_map(
    jwt_handler: Optional[JWTHandler],
    received_response: Optional[dict],
    user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]],
    prisma_client: PrismaClient,
    user_api_key_cache: UserApiKeyCache,
    user_defined_values: Optional[SSOUserDefinedValues],
) -> None:
    """
    Apply jwt_litellm_role_map during SSO login.

    When jwt_litellm_role_map is configured with sync_user_role_and_teams=True,
    this ensures SSO users get the same role mapping as API/JWT users. Without
    this, the SSO path falls back to INTERNAL_USER_VIEW_ONLY for roles that
    don't directly match LitellmUserRoles enum values.
    """
    if jwt_handler is None or received_response is None:
        return
    if not jwt_handler.litellm_jwtauth.sync_user_role_and_teams:
        return
    if not jwt_handler.litellm_jwtauth.jwt_litellm_role_map:
        return

    mapped_role = jwt_handler.map_jwt_role_to_litellm_role(received_response)
    if mapped_role is None:
        return

    verbose_proxy_logger.info(
        f"SSO jwt_litellm_role_map matched role: {mapped_role.value}"
    )

    # Update user_defined_values so downstream code uses the mapped role
    if user_defined_values is not None:
        user_defined_values["user_role"] = mapped_role.value

    # Update existing DB record if role differs
    if user_info is not None and user_info.user_role != mapped_role.value:
        await UserRepository(prisma_client).table.update(
            where={"user_id": user_info.user_id},
            data={"user_role": mapped_role.value},
        )
        user_info.user_role = mapped_role.value
        await user_api_key_cache.async_set_cache(
            key=user_info.user_id,
            value=user_info,
            model_type=LiteLLM_UserTable,
        )


def apply_user_info_values_to_sso_user_defined_values(
    user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]],
    user_defined_values: Optional[SSOUserDefinedValues],
) -> Optional[SSOUserDefinedValues]:
    if user_defined_values is None:
        return None
    if user_info is not None and user_info.user_id is not None:
        user_defined_values["user_id"] = user_info.user_id

    # SSO role takes precedence - only use DB role if SSO didn't provide one
    # This ensures SSO is the authoritative source for user roles
    sso_role = user_defined_values.get("user_role")
    db_role = user_info.user_role if user_info else None

    if _should_use_role_from_sso_response(sso_role):
        # SSO provided a valid role, keep it and log that we're using it
        verbose_proxy_logger.info(
            f"Using SSO role: {sso_role} (DB role was: {db_role})"
        )
    else:
        # SSO didn't provide a valid role, fall back to DB role or default
        if user_info is None or user_info.user_role is None:
            user_defined_values["user_role"] = (
                LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
            )
            verbose_proxy_logger.debug(
                "No SSO or DB role found, using default: INTERNAL_USER_VIEW_ONLY"
            )
        else:
            user_defined_values["user_role"] = user_info.user_role
            verbose_proxy_logger.debug(f"Using DB role: {user_info.user_role}")

    # Preserve the user's existing models from the database
    if user_info is not None and hasattr(user_info, "models") and user_info.models:
        user_defined_values["models"] = user_info.models

    return user_defined_values


async def check_and_update_if_proxy_admin_id(
    user_role: str, user_id: str, prisma_client: Optional[PrismaClient]
):
    """
    - Check if user role in DB is admin
    - If not, update user role in DB to admin role
    """
    proxy_admin_id = os.getenv("PROXY_ADMIN_ID")
    if proxy_admin_id is not None and proxy_admin_id == user_id:
        if user_role and user_role == LitellmUserRoles.PROXY_ADMIN.value:
            return user_role

        if prisma_client:
            await UserRepository(prisma_client).table.update(
                where={"user_id": user_id},
                data={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
            )

        user_role = LitellmUserRoles.PROXY_ADMIN.value

    return user_role


@router.get("/sso/callback", tags=["experimental"], include_in_schema=False)
async def auth_callback(request: Request, state: Optional[str] = None):
    """Verify login"""
    verbose_proxy_logger.info(f"Starting SSO callback with state: {state}")

    # Check if this is a CLI login (state starts with our CLI prefix)
    from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX
    from litellm.proxy._types import LiteLLM_JWTAuth
    from litellm.proxy.auth.handle_jwt import JWTHandler
    from litellm.proxy.proxy_server import (
        general_settings,
        jwt_handler,
        master_key,
        prisma_client,
        user_api_key_cache,
    )

    if prisma_client is None:
        raise HTTPException(
            status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
        )

    sso_jwt_handler: Optional[JWTHandler] = None
    ui_access_mode = general_settings.get("ui_access_mode", None)
    if ui_access_mode is not None and isinstance(ui_access_mode, dict):
        sso_jwt_handler = JWTHandler()
        sso_jwt_handler.update_environment(
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            litellm_jwtauth=LiteLLM_JWTAuth(
                team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get(
                    "sso_group_jwt_field", None
                ),
            ),
            leeway=0,
        )

    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
    generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)
    received_response: Optional[dict] = None
    access_token_payload: Optional[dict] = None
    # get url from request
    if master_key is None:
        raise ProxyException(
            message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml.  https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.",
            type=ProxyErrorTypes.auth_error,
            param="master_key",
            code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        )
    redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
        request=request, sso_callback_route="sso/callback"
    )

    verbose_proxy_logger.info(f"Redirecting to {redirect_url}")
    result = None
    if google_client_id is not None:
        result = await GoogleSSOHandler.get_google_callback_response(
            request=request,
            google_client_id=google_client_id,
            redirect_url=redirect_url,
        )
    elif microsoft_client_id is not None:
        result = await MicrosoftSSOHandler.get_microsoft_callback_response(
            request=request,
            microsoft_client_id=microsoft_client_id,
            redirect_url=redirect_url,
        )

    elif generic_client_id is not None:
        (
            result,
            received_response,
            access_token_payload,
        ) = await get_generic_sso_response(
            request=request,
            jwt_handler=jwt_handler,
            generic_client_id=generic_client_id,
            redirect_url=redirect_url,
            sso_jwt_handler=sso_jwt_handler,
        )

    if result is None:
        raise HTTPException(
            status_code=401,
            detail="Result not returned by SSO provider.",
        )

    if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"):
        # State format: {PREFIX}:{login_id}[:{user_code}]
        state_parts = state.split(":", 2)
        key_id = state_parts[1] if len(state_parts) > 1 else None
        prefill_user_code = state_parts[2] if len(state_parts) > 2 else None

        verbose_proxy_logger.info("CLI SSO callback detected")
        return await cli_sso_callback(
            request=request,
            key=key_id,
            prefill_user_code=prefill_user_code,
            result=result,
            received_response=received_response,
        )

    # Control-plane cross-origin: read return_to from cookie.
    # Starlette's cookie_parser already handles RFC 2109 unquoting.
    cp_return_to: Optional[str] = request.cookies.get("litellm_cp_return_to")

    return await SSOAuthenticationHandler.get_redirect_response_from_openid(
        result=result,
        request=request,
        received_response=received_response,
        generic_client_id=generic_client_id,
        ui_access_mode=ui_access_mode,
        access_token_payload=access_token_payload,
        jwt_handler=jwt_handler,
        return_to=cp_return_to,
    )


async def _build_cli_sso_user_defined_values(
    result: Union[OpenID, dict],
    parsed_openid_result: ParsedOpenIDResult,
) -> Optional[SSOUserDefinedValues]:
    from litellm.proxy.proxy_server import user_custom_sso

    user_id = parsed_openid_result.get("user_id")
    if user_custom_sso is not None:
        if inspect.iscoroutinefunction(user_custom_sso):
            return await user_custom_sso(result)  # type: ignore
        raise ValueError("user_custom_sso must be a coroutine function")
    if user_id is None:
        return None
    return SSOUserDefinedValues(
        models=[],
        user_id=user_id,
        user_email=parsed_openid_result.get("user_email"),
        max_budget=litellm.max_internal_user_budget,
        user_role=parsed_openid_result.get("user_role"),
        budget_duration=litellm.internal_user_budget_duration,
    )


async def _fetch_cli_sso_team_details(
    prisma_client: PrismaClient,
    teams: List[str],
) -> List[Dict[str, Any]]:
    team_details: List[Dict[str, Any]] = []
    try:
        if teams:
            prisma_teams = await TeamRepository(prisma_client).table.find_many(
                where={"team_id": {"in": teams}}
            )
            for team_row in prisma_teams:
                team_dict = team_row.model_dump()
                team_details.append(
                    {
                        "team_id": team_dict.get("team_id"),
                        "team_alias": team_dict.get("team_alias"),
                    }
                )
    except Exception as e:
        verbose_proxy_logger.error(
            f"Error fetching team details for CLI SSO session: {e}"
        )
    return team_details


async def _complete_cli_sso_callback_session(
    *,
    request: Request,
    key: str,
    flow: dict,
    result: Union[OpenID, dict],
    parsed_openid_result: ParsedOpenIDResult,
    user_defined_values: Optional[SSOUserDefinedValues],
    prisma_client: PrismaClient,
    user_api_key_cache: UserApiKeyCache,
    proxy_logging_obj: ProxyLogging,
    prefill_user_code: str | None = None,
):
    from fastapi.responses import HTMLResponse

    user_id = parsed_openid_result.get("user_id")
    user_email = parsed_openid_result.get("user_email")
    user_info = await get_user_info_from_db(
        result=result,
        prisma_client=prisma_client,
        user_api_key_cache=user_api_key_cache,
        proxy_logging_obj=proxy_logging_obj,
        user_email=user_email,
        user_defined_values=user_defined_values,
        alternate_user_id=user_id,
    )
    if user_info is None:
        raise HTTPException(
            status_code=500, detail="Failed to retrieve user information from SSO"
        )
    if not user_info.user_id:
        raise HTTPException(
            status_code=500, detail="Failed to retrieve user information from SSO"
        )

    teams: List[str] = []
    if hasattr(user_info, "teams") and user_info.teams:
        teams = user_info.teams if isinstance(user_info.teams, list) else []

    team_details = await _fetch_cli_sso_team_details(
        prisma_client=prisma_client, teams=teams
    )
    attribution_metadata = build_cli_sso_attribution_metadata(result=result)
    if attribution_metadata:
        await _persist_cli_sso_user_metadata(
            prisma_client=prisma_client,
            user_id=cast(str, user_info.user_id),
            attribution_metadata=attribution_metadata,
        )

    flow["session_data"] = {
        "user_id": cast(str, user_info.user_id),
        "user_role": user_info.user_role,
        "models": user_info.models if hasattr(user_info, "models") else [],
        "user_email": user_email,
        "teams": teams,
        "team_details": team_details,
        "attribution_metadata": attribution_metadata,
    }
    flow["sso_complete"] = True
    browser_complete_token = secrets.token_urlsafe(32)
    flow["browser_complete_token_hash"] = _hash_cli_sso_secret(browser_complete_token)
    _set_cli_sso_flow(login_id=key, cache=user_api_key_cache, flow=flow)

    verbose_proxy_logger.info(
        f"Stored CLI SSO session for user: {user_info.user_id}, teams: {teams}, num_teams: {len(teams)}"
    )
    verify_url = get_custom_url(
        request_base_url=str(request.base_url),
        route=f"sso/cli/complete/{key}",
    )
    return HTMLResponse(
        content=_render_cli_sso_verification_page(
            verify_url=verify_url,
            browser_complete_token=browser_complete_token,
            prefill_user_code=prefill_user_code,
        ),
        status_code=200,
    )


async def cli_sso_callback(
    request: Request,
    key: Optional[str] = None,
    result: Optional[Union[OpenID, dict]] = None,
    received_response: Optional[dict] = None,
    prefill_user_code: str | None = None,
):
    """CLI SSO callback - stores session info for JWT generation on polling"""
    verbose_proxy_logger.info("CLI SSO callback")

    from litellm.proxy.proxy_server import (
        general_settings,
        prisma_client,
        proxy_logging_obj,
        user_api_key_cache,
    )

    flow = _get_cli_sso_flow_or_raise(login_id=key, cache=user_api_key_cache)

    if prisma_client is None:
        raise HTTPException(
            status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
        )

    if result is None:
        raise HTTPException(
            status_code=500,
            detail="SSO authentication failed - no result returned from provider",
        )

    # After None check, cast to non-None type for type checker
    result_non_none: Union[OpenID, dict] = cast(Union[OpenID, dict], result)

    try:
        parsed_openid_result = (
            SSOAuthenticationHandler._get_user_email_and_id_from_result(
                result=result_non_none,
                generic_client_id=os.getenv("GENERIC_CLIENT_ID", None),
            )
        )
        verbose_proxy_logger.debug(f"parsed_openid_result: {parsed_openid_result}")
        user_defined_values = await _build_cli_sso_user_defined_values(
            result=result_non_none,
            parsed_openid_result=parsed_openid_result,
        )

        SSOAuthenticationHandler.verify_user_in_restricted_sso_group(
            general_settings=general_settings,
            result=result_non_none,
            received_response=received_response,
        )

        return await _complete_cli_sso_callback_session(
            request=request,
            key=cast(str, key),
            flow=flow,
            result=result_non_none,
            parsed_openid_result=parsed_openid_result,
            user_defined_values=user_defined_values,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            proxy_logging_obj=proxy_logging_obj,
            prefill_user_code=prefill_user_code,
        )
    except ProxyException:
        raise
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.error(f"Error with CLI SSO callback: {e}")
        raise HTTPException(
            status_code=500, detail=f"Failed to process CLI SSO: {str(e)}"
        )


@router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False)
async def cli_poll_key(
    key_id: str,
    team_id: Optional[str] = None,
    x_litellm_cli_poll_secret: Optional[str] = Header(default=None),
):
    """
    CLI polling endpoint - retrieves session from cache and generates JWT.

    Flow:
    1. First poll (no team_id): Returns teams list without generating JWT
    2. Second poll (with team_id): Generates JWT with selected team and deletes session

    Args:
        key_id: The CLI login session ID
        team_id: Optional team ID to assign to the JWT. If provided, must be one of user's teams.
    """
    from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken
    from litellm.proxy.proxy_server import user_api_key_cache

    try:
        flow = _get_cli_sso_flow_or_raise(login_id=key_id, cache=user_api_key_cache)
        if not _verify_cli_sso_poll_secret(
            flow=flow, poll_secret=x_litellm_cli_poll_secret
        ):
            raise HTTPException(status_code=403, detail="Invalid CLI polling secret")

        if not flow.get("sso_complete") or not flow.get("user_code_verified"):
            return {"status": "pending"}

        session_data = flow.get("session_data")

        if isinstance(session_data, dict):
            user_teams = session_data.get("teams", [])
            user_team_details = session_data.get("team_details")
            user_id = session_data["user_id"]

            verbose_proxy_logger.info(
                f"CLI poll: user={user_id}, team_id={team_id}, user_teams={user_teams}, num_teams={len(user_teams)}"
            )

            # If no team_id provided and user has teams, return teams list for selection
            # Don't generate JWT yet - let CLI select a team first. For newer
            # clients we return rich team details (id + alias); older clients
            # can continue to rely on the simple "teams" list.
            if team_id is None and len(user_teams) > 1:
                verbose_proxy_logger.info(
                    f"Returning teams list for user {user_id} to select from: {user_teams}"
                )
                # Best-effort construction of team_details if it wasn't
                # already cached for some reason.
                team_details_response: Optional[List[Dict[str, Any]]] = None
                if isinstance(user_team_details, list) and user_team_details:
                    team_details_response = user_team_details
                elif user_teams:
                    team_details_response = [
                        {"team_id": t, "team_alias": None} for t in user_teams
                    ]
                poll_response: Dict[str, Any] = {
                    "status": "ready",
                    "user_id": user_id,
                    "teams": user_teams,
                    "team_details": team_details_response,
                    "requires_team_selection": True,
                }
                attribution_metadata = _cli_poll_attribution_metadata_from_session(
                    session_data
                )
                if attribution_metadata:
                    poll_response["attribution_metadata"] = attribution_metadata
                return poll_response

            # Validate team_id if provided
            if team_id is not None:
                if team_id not in user_teams:
                    raise HTTPException(
                        status_code=403,
                        detail=f"User does not belong to team: {team_id}. Available teams: {user_teams}",
                    )
            else:
                # If no team_id provided and user has 0 or 1 team, use first team (or None)
                team_id = user_teams[0] if len(user_teams) > 0 else None

            team_alias = None
            if team_id and isinstance(user_team_details, list):
                team_alias = next(
                    (
                        team.get("team_alias")
                        for team in user_team_details
                        if team.get("team_id") == team_id
                    ),
                    None,
                )

            # Create user object for JWT generation
            user_info = LiteLLM_UserTable(
                user_id=user_id,
                user_role=session_data["user_role"],
                models=session_data.get("models", []),
                max_budget=litellm.max_ui_session_budget,
            )

            # Generate CLI JWT on-demand (expiration configurable via LITELLM_CLI_JWT_EXPIRATION_HOURS)
            # Pass selected team_id to ensure JWT has correct team
            jwt_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(
                user_info=user_info, team_id=team_id, team_alias=team_alias
            )

            # Delete cache entry (single-use)
            user_api_key_cache.delete_cache(key=_get_cli_sso_flow_cache_key(key_id))

            verbose_proxy_logger.info(
                f"CLI JWT generated for user: {user_id}, team: {team_id}"
            )
            poll_response = {
                "status": "ready",
                "key": jwt_token,
                "user_id": user_id,
                "team_id": team_id,
                "teams": user_teams,
                # Echo back any team details we have so clients can
                # present nicer information if needed.
                "team_details": user_team_details,
            }
            attribution_metadata = _cli_poll_attribution_metadata_from_session(
                session_data
            )
            if attribution_metadata:
                poll_response["attribution_metadata"] = attribution_metadata
            return poll_response
        else:
            return {"status": "pending"}

    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.error(f"Error polling for CLI JWT: {e}")
        raise HTTPException(
            status_code=500, detail=f"Error checking session status: {str(e)}"
        )


async def insert_sso_user(
    result_openid: Optional[Union[OpenID, dict]],
    user_defined_values: Optional[SSOUserDefinedValues] = None,
) -> NewUserResponse:
    """
    Helper function to create a New User in LiteLLM DB after a successful SSO login

    Args:
        result_openid (OpenID): User information in OpenID format if the login was successful.
        user_defined_values (Optional[SSOUserDefinedValues], optional): LiteLLM SSOValues / fields that were read

    Returns:
        Tuple[str, str]: User ID and User Role
    """
    verbose_proxy_logger.debug(
        f"Inserting SSO user into DB. User values: {user_defined_values}"
    )
    if result_openid is None:
        raise ValueError("result_openid is None")
    if isinstance(result_openid, dict):
        result_openid = OpenID(**result_openid)

    if user_defined_values is None:
        raise ValueError("user_defined_values is None")

    # Apply default_internal_user_params
    if litellm.default_internal_user_params:
        # Preserve the SSO-extracted role if it's a valid LiteLLM role,
        # regardless of how it was determined (role_mappings, Microsoft app_roles,
        # GENERIC_USER_ROLE_ATTRIBUTE, custom SSO handler, etc.)
        sso_role = user_defined_values.get("user_role")
        if _should_use_role_from_sso_response(sso_role):
            # Preserve the SSO-extracted role, but apply other defaults
            preserved_role = sso_role
            user_defined_values.update(litellm.default_internal_user_params)  # type: ignore
            user_defined_values["user_role"] = preserved_role  # Restore preserved role
            verbose_proxy_logger.debug(
                f"Preserved SSO-extracted role '{preserved_role}'"
            )
        else:
            # SSO didn't provide a valid role, apply all defaults including role
            user_defined_values.update(litellm.default_internal_user_params)  # type: ignore

    # Set budget for internal users
    if user_defined_values.get("user_role") == LitellmUserRoles.INTERNAL_USER.value:
        if user_defined_values.get("max_budget") is None:
            user_defined_values["max_budget"] = litellm.max_internal_user_budget
        if user_defined_values.get("budget_duration") is None:
            user_defined_values["budget_duration"] = (
                litellm.internal_user_budget_duration
            )

    if user_defined_values["user_role"] is None:
        user_defined_values["user_role"] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY

    new_user_request = NewUserRequest(
        user_id=user_defined_values["user_id"],
        user_email=normalize_email(user_defined_values["user_email"]),
        user_role=user_defined_values["user_role"],  # type: ignore
        max_budget=user_defined_values["max_budget"],
        budget_duration=user_defined_values["budget_duration"],
        sso_user_id=user_defined_values["user_id"],
        auto_create_key=False,
    )

    if result_openid and hasattr(result_openid, "provider"):
        new_user_request.metadata = {
            "auth_provider": getattr(result_openid, "provider")
        }

    response = await new_user(
        data=new_user_request,
        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
    )

    return response


@router.get(
    "/sso/get/ui_settings",
    tags=["experimental"],
    include_in_schema=False,
    dependencies=[Depends(user_api_key_auth)],
)
async def get_ui_settings(request: Request):
    from litellm.proxy.proxy_server import general_settings, proxy_state

    _proxy_base_url = os.getenv("PROXY_BASE_URL", None)
    _logout_url = os.getenv("PROXY_LOGOUT_URL", None)
    _api_doc_base_url = os.getenv("LITELLM_UI_API_DOC_BASE_URL", None)
    _is_sso_enabled = _has_user_setup_sso()
    disable_expensive_db_queries = (
        proxy_state.get_proxy_state_variable("spend_logs_row_count")
        > MAX_SPENDLOG_ROWS_TO_QUERY
    )
    default_team_disabled = general_settings.get("default_team_disabled", False)
    if "PROXY_DEFAULT_TEAM_DISABLED" in os.environ:
        if os.environ["PROXY_DEFAULT_TEAM_DISABLED"].lower() == "true":
            default_team_disabled = True

    return {
        "PROXY_BASE_URL": _proxy_base_url,
        "PROXY_LOGOUT_URL": _logout_url,
        "LITELLM_UI_API_DOC_BASE_URL": _api_doc_base_url,
        "DEFAULT_TEAM_DISABLED": default_team_disabled,
        "SSO_ENABLED": _is_sso_enabled,
        "NUM_SPEND_LOGS_ROWS": proxy_state.get_proxy_state_variable(
            "spend_logs_row_count"
        ),
        "DISABLE_EXPENSIVE_DB_QUERIES": disable_expensive_db_queries,
    }


@router.get(
    "/sso/readiness",
    tags=["experimental"],
    dependencies=[Depends(user_api_key_auth)],
)
async def sso_readiness():
    """
    Health endpoint for checking SSO readiness.
    Checks if the configured SSO provider has all required environment variables set in memory.
    """
    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
    generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)

    # Determine which SSO provider is configured
    configured_provider = None
    if google_client_id is not None:
        configured_provider = "google"
    elif microsoft_client_id is not None:
        configured_provider = "microsoft"
    elif generic_client_id is not None:
        configured_provider = "generic"

    # If no SSO is configured, return healthy (SSO is optional)
    if configured_provider is None:
        return {
            "status": "healthy",
            "sso_configured": False,
            "message": "No SSO provider configured",
        }

    # Check required environment variables for the configured provider
    missing_vars = []

    if configured_provider == "google":
        google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
        if google_client_secret is None:
            missing_vars.append("GOOGLE_CLIENT_SECRET")

    elif configured_provider == "microsoft":
        microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
        microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
        if microsoft_client_secret is None:
            missing_vars.append("MICROSOFT_CLIENT_SECRET")
        if microsoft_tenant is None:
            missing_vars.append("MICROSOFT_TENANT")

    elif configured_provider == "generic":
        generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
        generic_authorization_endpoint = os.getenv(
            "GENERIC_AUTHORIZATION_ENDPOINT", None
        )
        generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
        generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
        if generic_client_secret is None:
            missing_vars.append("GENERIC_CLIENT_SECRET")
        if generic_authorization_endpoint is None:
            missing_vars.append("GENERIC_AUTHORIZATION_ENDPOINT")
        if generic_token_endpoint is None:
            missing_vars.append("GENERIC_TOKEN_ENDPOINT")
        if generic_userinfo_endpoint is None:
            missing_vars.append("GENERIC_USERINFO_ENDPOINT")

    # If all required variables are present, return healthy
    if len(missing_vars) == 0:
        return {
            "status": "healthy",
            "sso_configured": True,
            "provider": configured_provider,
            "message": f"{configured_provider.capitalize()} SSO is properly configured",
        }

    # If some variables are missing, return unhealthy
    raise HTTPException(
        status_code=503,
        detail={
            "status": "unhealthy",
            "sso_configured": True,
            "provider": configured_provider,
            "missing_environment_variables": missing_vars,
            "message": f"{configured_provider.capitalize()} SSO is configured but missing required environment variables: {', '.join(missing_vars)}",
        },
    )


class SSOAuthenticationHandler:
    """
    Handler for SSO Authentication across all SSO providers
    """

    @staticmethod
    def _validate_return_to(return_to: str) -> bool:
        """
        Validate that return_to matches the configured control_plane_url origin.

        Returns True if return_to is valid and should be used.
        Returns False if control_plane_url is not configured (return_to is ignored).
        Raises HTTPException(400) if return_to origin does not match control_plane_url origin.
        """
        from litellm.proxy.proxy_server import general_settings

        control_plane_url = general_settings.get("control_plane_url")
        if control_plane_url is None:
            return False

        def _origin(url: str) -> tuple:
            parsed = urlparse(url)
            scheme = (parsed.scheme or "").lower()
            hostname = (parsed.hostname or "").lower()
            default_port = 443 if scheme == "https" else 80
            port = parsed.port if parsed.port is not None else default_port
            return (scheme, hostname, port)

        if _origin(return_to) != _origin(control_plane_url):
            raise HTTPException(
                status_code=400,
                detail="return_to does not match the configured control_plane_url",
            )

        return True

    @staticmethod
    async def get_sso_login_redirect(
        redirect_url: str,
        google_client_id: Optional[str] = None,
        microsoft_client_id: Optional[str] = None,
        generic_client_id: Optional[str] = None,
        state: Optional[str] = None,
        request: Optional[Request] = None,
    ) -> Optional[RedirectResponse]:
        """
        Step 1. Call Get Login Redirect for the SSO provider. Send the redirect response to `redirect_url`

        Args:
            redirect_url (str): The URL to redirect the user to after login
            google_client_id (Optional[str], optional): The Google Client ID. Defaults to None.
            microsoft_client_id (Optional[str], optional): The Microsoft Client ID. Defaults to None.
            generic_client_id (Optional[str], optional): The Generic Client ID. Defaults to None.
            request: Optional FastAPI request, used to drive the ``Secure``
                attribute on the ``litellm_oauth_state`` CSRF cookie.

        Returns:
            RedirectResponse: The redirect response from the SSO provider.
        """
        # Google SSO Auth
        if google_client_id is not None:
            from fastapi_sso.sso.google import GoogleSSO

            google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
            if google_client_secret is None:
                raise ProxyException(
                    message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="GOOGLE_CLIENT_SECRET",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            google_sso = GoogleSSO(
                client_id=google_client_id,
                client_secret=google_client_secret,
                redirect_uri=redirect_url,
            )
            verbose_proxy_logger.info(
                f"In /google-login/key/generate, \nGOOGLE_REDIRECT_URI: {redirect_url}\nGOOGLE_CLIENT_ID: {google_client_id}"
            )
            with google_sso:
                return await google_sso.get_login_redirect(state=state)
        # Microsoft SSO Auth
        elif microsoft_client_id is not None:
            microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
            microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
            if microsoft_client_secret is None:
                raise ProxyException(
                    message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="MICROSOFT_CLIENT_SECRET",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            microsoft_sso = CustomMicrosoftSSO(
                client_id=microsoft_client_id,
                client_secret=microsoft_client_secret,
                tenant=microsoft_tenant,
                redirect_uri=redirect_url,
                allow_insecure_http=True,
            )
            with microsoft_sso:
                return await microsoft_sso.get_login_redirect(state=state)
        elif generic_client_id is not None:
            from fastapi_sso.sso.base import DiscoveryDocument
            from fastapi_sso.sso.generic import create_provider

            generic_client_secret = os.getenv("GENERIC_CLIENT_SECRET", None)
            generic_scope = os.getenv("GENERIC_SCOPE", "openid email profile").split(
                " "
            )
            generic_authorization_endpoint = os.getenv(
                "GENERIC_AUTHORIZATION_ENDPOINT", None
            )
            generic_token_endpoint = os.getenv("GENERIC_TOKEN_ENDPOINT", None)
            generic_userinfo_endpoint = os.getenv("GENERIC_USERINFO_ENDPOINT", None)
            if generic_client_secret is None:
                raise ProxyException(
                    message="GENERIC_CLIENT_SECRET not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_CLIENT_SECRET",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            if generic_authorization_endpoint is None:
                raise ProxyException(
                    message="GENERIC_AUTHORIZATION_ENDPOINT not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_AUTHORIZATION_ENDPOINT",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            if generic_token_endpoint is None:
                raise ProxyException(
                    message="GENERIC_TOKEN_ENDPOINT not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_TOKEN_ENDPOINT",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            if generic_userinfo_endpoint is None:
                raise ProxyException(
                    message="GENERIC_USERINFO_ENDPOINT not set. Set it in .env file",
                    type=ProxyErrorTypes.auth_error,
                    param="GENERIC_USERINFO_ENDPOINT",
                    code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                )
            verbose_proxy_logger.debug(
                f"authorization_endpoint: {generic_authorization_endpoint}\ntoken_endpoint: {generic_token_endpoint}\nuserinfo_endpoint: {generic_userinfo_endpoint}"
            )
            verbose_proxy_logger.debug(
                f"GENERIC_REDIRECT_URI: {redirect_url}\nGENERIC_CLIENT_ID: {generic_client_id}\n"
            )
            discovery = DiscoveryDocument(
                authorization_endpoint=generic_authorization_endpoint,
                token_endpoint=generic_token_endpoint,
                userinfo_endpoint=generic_userinfo_endpoint,
            )
            SSOProvider = create_provider(name="oidc", discovery_document=discovery)
            generic_sso = SSOProvider(
                client_id=generic_client_id,
                client_secret=generic_client_secret,
                redirect_uri=redirect_url,
                allow_insecure_http=True,
                scope=generic_scope,
            )
            return await SSOAuthenticationHandler.get_generic_sso_redirect_response(
                generic_sso=generic_sso,
                state=state,
                generic_authorization_endpoint=generic_authorization_endpoint,
                request=request,
            )
        raise ValueError(
            "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso"
        )

    @staticmethod
    async def get_generic_sso_redirect_response(
        generic_sso: Any,
        state: Optional[str] = None,
        generic_authorization_endpoint: Optional[str] = None,
        request: Optional[Request] = None,
    ) -> Optional[RedirectResponse]:
        """
        Get the redirect response for Generic SSO
        """
        from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

        from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache

        with generic_sso:
            # State is bound to the caller's browser via a ``litellm_oauth_state``
            # HttpOnly cookie set on the redirect response below; the SSO
            # callback validates the URL ``state`` against that cookie before
            # completing the PKCE token exchange.  Without this binding, an
            # attacker who pre-mints a state + a cached PKCE verifier can hand
            # the link to a victim and capture the resulting access token
            # (Login CSRF / token theft).
            (
                redirect_params,
                code_verifier,
            ) = SSOAuthenticationHandler._get_generic_sso_redirect_params(
                state=state,
                generic_authorization_endpoint=generic_authorization_endpoint,
            )

            # Separate PKCE params from state params (fastapi-sso doesn't accept code_challenge)
            pkce_params = {}
            state_only_params = {}
            for key, value in redirect_params.items():
                if key in ("code_challenge", "code_challenge_method"):
                    pkce_params[key] = value
                else:
                    state_only_params[key] = value

            # Get the redirect response from fastapi-sso with only state param
            redirect_response = await generic_sso.get_login_redirect(**state_only_params)  # type: ignore

            # If PKCE is enabled, add PKCE parameters to the redirect URL
            if code_verifier and "state" in redirect_params:
                # Store code_verifier in cache (10 min TTL). Wrap in dict for proper
                # JSON serialization in Redis. Use Redis when available so callbacks
                # landing on another pod can retrieve it (multi-pod SSO).
                cache_key = f"pkce_verifier:{redirect_params['state']}"
                if redis_usage_cache is not None:
                    await redis_usage_cache.async_set_cache(
                        key=cache_key,
                        value={"code_verifier": code_verifier},
                        ttl=600,
                    )
                else:
                    await user_api_key_cache.async_set_cache(
                        key=cache_key,
                        value={"code_verifier": code_verifier},
                        ttl=600,
                    )
                verbose_proxy_logger.debug(
                    "PKCE code_verifier stored in cache (TTL: 600s)"
                )

                # Add PKCE parameters to the authorization URL
                if pkce_params:
                    parsed_url = urlparse(str(redirect_response.headers["location"]))
                    query_params = parse_qs(parsed_url.query)

                    # Add PKCE parameters
                    for key, value in pkce_params.items():
                        query_params[key] = [value]

                    # Reconstruct the URL with PKCE parameters
                    new_query = urlencode(query_params, doseq=True)
                    new_url = urlunparse(
                        (
                            parsed_url.scheme,
                            parsed_url.netloc,
                            parsed_url.path,
                            parsed_url.params,
                            new_query,
                            parsed_url.fragment,
                        )
                    )

                    # Update the redirect response
                    redirect_response.headers["location"] = new_url

                # Bind state to the user's browser session.  The /callback
                # handler validates the URL ``state`` against this cookie via
                # ``secrets.compare_digest`` before exchanging the PKCE
                # code_verifier.  Only set the cookie when PKCE is in use
                # (i.e. inside this ``code_verifier`` branch) so two
                # concurrent SSO sessions — one PKCE, one plain — cannot
                # overwrite each other's state cookie.
                state_value = redirect_params.get("state")
                if state_value and redirect_response is not None:
                    # Production-safe default: require HTTPS for the
                    # CSRF-protection cookie unless we can prove the
                    # incoming request is HTTP (local dev).  Without
                    # ``Secure`` the cookie is sent over plain HTTP,
                    # letting a network observer read and replay the
                    # state value and bypass this protection.
                    secure_flag = request is None or request.url.scheme == "https"
                    redirect_response.set_cookie(
                        key="litellm_oauth_state",
                        value=state_value,
                        max_age=600,
                        httponly=True,
                        samesite="lax",
                        secure=secure_flag,
                    )
            return redirect_response

    @staticmethod
    def _get_generic_sso_redirect_params(
        state: Optional[str] = None,
        generic_authorization_endpoint: Optional[str] = None,
    ) -> Tuple[dict, Optional[str]]:
        """
        Get redirect parameters for Generic SSO with proper state priority handling.
        Optionally generates PKCE parameters if GENERIC_CLIENT_USE_PKCE is enabled.

        Priority order:
        1. CLI state (if provided)
        2. GENERIC_CLIENT_STATE environment variable
        3. Generated UUID (required by Okta and most OAuth providers)


        Args:
            state: Optional state parameter (e.g., CLI state)
            generic_authorization_endpoint: Authorization endpoint URL

        Returns:
            Tuple[dict, Optional[str]]:
                - Redirect parameters for SSO login (may include PKCE params)
                - code_verifier (if PKCE is enabled, None otherwise)
        """
        redirect_params = {}
        code_verifier: Optional[str] = None

        if state:
            # CLI state takes priority
            # the litellm proxy cli sends the "state" parameter to the proxy server for auth. We should maintain the state parameter for the cli if it is provided
            redirect_params["state"] = state
        else:
            generic_client_state = os.getenv("GENERIC_CLIENT_STATE", None)
            if generic_client_state:
                redirect_params["state"] = generic_client_state
            else:
                redirect_params["state"] = uuid.uuid4().hex

        # Handle PKCE (Proof Key for Code Exchange) if enabled
        # Set GENERIC_CLIENT_USE_PKCE=true to enable PKCE for enhanced OAuth security
        use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"

        if use_pkce:
            (
                code_verifier,
                code_challenge,
            ) = SSOAuthenticationHandler.generate_pkce_params()
            redirect_params["code_challenge"] = code_challenge
            redirect_params["code_challenge_method"] = "S256"
            verbose_proxy_logger.debug("PKCE enabled for authorization request")

        return redirect_params, code_verifier

    @staticmethod
    def should_use_sso_handler(
        google_client_id: Optional[str] = None,
        microsoft_client_id: Optional[str] = None,
        generic_client_id: Optional[str] = None,
    ) -> bool:
        if (
            google_client_id is not None
            or microsoft_client_id is not None
            or generic_client_id is not None
        ):
            return True
        return False

    @staticmethod
    def get_redirect_url_for_sso(
        request: Request,
        sso_callback_route: str,
        existing_key: Optional[str] = None,
    ) -> str:
        """
        Get the redirect URL for SSO

        Note: existing_key is not added to the URL to avoid changing the callback URL.
        It should be passed via the state parameter instead.
        """
        from litellm.proxy.utils import get_custom_url

        redirect_url = get_custom_url(request_base_url=str(request.base_url))
        if redirect_url.endswith("/"):
            redirect_url += sso_callback_route
        else:
            redirect_url += "/" + sso_callback_route

        return redirect_url

    @staticmethod
    async def upsert_sso_user(
        result: Optional[Union[CustomOpenID, OpenID, dict]],
        user_info: Optional[Union[NewUserResponse, LiteLLM_UserTable]],
        user_email: Optional[str],
        user_defined_values: Optional[SSOUserDefinedValues],
        prisma_client: PrismaClient,
    ):
        """
        Connects the SSO Users to the User Table in LiteLLM DB

        - If user on LiteLLM DB, update the user_email and user_role (if SSO provides valid role) with the SSO values
        - If user not on LiteLLM DB, insert the user into LiteLLM DB
        """
        try:
            if user_info is not None:
                user_id = user_info.user_id
                update_data = _build_sso_user_update_data(
                    result=result,
                    user_email=user_email,
                    user_id=user_id,
                )

                await UserRepository(prisma_client).table.update_many(
                    where={"user_id": user_id}, data=update_data
                )
            else:
                verbose_proxy_logger.info(
                    "user not in DB, inserting user into LiteLLM DB"
                )
                # user not in DB, insert User into LiteLLM DB
                user_info = await insert_sso_user(
                    result_openid=result,
                    user_defined_values=user_defined_values,
                )
            return user_info
        except Exception as e:
            verbose_proxy_logger.exception(
                f"Error upserting SSO user into LiteLLM DB: {e}"
            )
            return user_info

    @staticmethod
    async def add_user_to_teams_from_sso_response(
        result: Optional[Union[CustomOpenID, OpenID, dict]],
        user_info: Optional[Union[NewUserResponse, LiteLLM_UserTable]],
    ):
        """
        Adds the user as a team member to the teams specified in the SSO responses `team_ids` field


        The `team_ids` field is populated by litellm after processing the SSO response
        """
        if user_info is None:
            verbose_proxy_logger.debug(
                "User not found in LiteLLM DB, skipping team member addition"
            )
            return
        sso_teams = getattr(result, "team_ids", [])
        await add_missing_team_member(user_info=user_info, sso_teams=sso_teams)

    @staticmethod
    def verify_user_in_restricted_sso_group(
        general_settings: Dict,
        result: Optional[Union[CustomOpenID, OpenID, dict]],
        received_response: Optional[dict],
    ) -> Literal[True]:
        """
        when ui_access_mode.type == "restricted_sso_group":

        - result.team_ids should contain the restricted_sso_group
        - if not, raise a ProxyException
        - if so, return True
        - if result.team_ids is None, return False
        - if result.team_ids is an empty list, return False
        - if result.team_ids is a list, return True if the restricted_sso_group is in the list, otherwise return False
        """

        ui_access_mode = cast(
            Optional[Union[Dict, str]], general_settings.get("ui_access_mode")
        )

        if ui_access_mode is None:
            return True
        if isinstance(ui_access_mode, str):
            return True
        team_ids = getattr(result, "team_ids", [])

        if ui_access_mode.get("type") == "restricted_sso_group":
            restricted_sso_group = ui_access_mode.get("restricted_sso_group")
            if restricted_sso_group not in team_ids:
                raise ProxyException(
                    message=f"User is not in the restricted SSO group: {restricted_sso_group}. User groups: {team_ids}. Received SSO response: {received_response}",
                    type=ProxyErrorTypes.auth_error,
                    param="restricted_sso_group",
                    code=status.HTTP_403_FORBIDDEN,
                )
        return True

    @staticmethod
    async def create_litellm_team_from_sso_group(
        litellm_team_id: str,
        litellm_team_name: Optional[str] = None,
    ):
        """
        Creates a Litellm Team from a SSO Group ID

        Your SSO provider might have groups that should be created on LiteLLM

        Use this helper to create a Litellm Team from a SSO Group ID

        Args:
            litellm_team_id (str): The ID of the Litellm Team
            litellm_team_name (Optional[str]): The name of the Litellm Team
        """
        from litellm.proxy.proxy_server import prisma_client

        if prisma_client is None:
            raise ProxyException(
                message="Prisma client not found. Set it in the proxy_server.py file",
                type=ProxyErrorTypes.auth_error,
                param="prisma_client",
                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )
        try:
            team_obj = await TeamRepository(prisma_client).table.find_first(
                where={"team_id": litellm_team_id}
            )
            verbose_proxy_logger.debug(f"Team object: {team_obj}")

            # only create a new team if it doesn't exist
            if team_obj:
                verbose_proxy_logger.debug(
                    f"Team already exists: {litellm_team_id} - {litellm_team_name}"
                )
                return

            team_request: NewTeamRequest = NewTeamRequest(
                team_id=litellm_team_id,
                team_alias=litellm_team_name,
            )
            if litellm.default_team_params:
                team_request = SSOAuthenticationHandler._cast_and_deepcopy_litellm_default_team_params(
                    default_team_params=litellm.default_team_params,
                    litellm_team_id=litellm_team_id,
                    litellm_team_name=litellm_team_name,
                    team_request=team_request,
                )

            await new_team(
                data=team_request,
                # params used for Audit Logging
                http_request=Request(scope={"type": "http", "method": "POST"}),
                user_api_key_dict=UserAPIKeyAuth(
                    token="",
                    key_alias=f"litellm.{MicrosoftSSOHandler.__name__}",
                ),
            )
        except Exception as e:
            verbose_proxy_logger.exception(f"Error creating Litellm Team: {e}")

    @staticmethod
    def _cast_and_deepcopy_litellm_default_team_params(
        default_team_params: Union[DefaultTeamSSOParams, Dict],
        team_request: NewTeamRequest,
        litellm_team_id: str,
        litellm_team_name: Optional[str] = None,
    ) -> NewTeamRequest:
        """
        Casts and deepcopies the litellm.default_team_params to a NewTeamRequest object

        - Ensures we create a new DefaultTeamSSOParams object
        - Handle the case where litellm.default_team_params is a dict or a DefaultTeamSSOParams object
        - Adds the litellm_team_id and litellm_team_name to the DefaultTeamSSOParams object
        """
        if isinstance(default_team_params, dict):
            _team_request = deepcopy(default_team_params)
            _team_request["team_id"] = litellm_team_id
            _team_request["team_alias"] = litellm_team_name
            team_request = NewTeamRequest(**_team_request)
        elif isinstance(litellm.default_team_params, DefaultTeamSSOParams):
            _default_team_params = deepcopy(litellm.default_team_params)
            _new_team_request = team_request.model_dump()
            _new_team_request.update(_default_team_params)
            team_request = NewTeamRequest(**_new_team_request)
        return team_request

    @staticmethod
    def _get_cli_state(
        source: str | None,
        key: str | None,
        existing_key: str | None = None,
        user_code: str | None = None,
    ) -> Optional[str]:
        """
        Checks the request 'source' if a cli state token was passed in

        This is used to authenticate through the CLI login flow.

        The state parameter format is: {PREFIX}:{login_id}[:{user_code}]
        - The state parameter is used to pass data through the OAuth flow without changing the callback URL
        - user_code is appended only for the opt-in verification_uri_complete flow so the verify page can pre-fill it
        """
        from litellm.constants import (
            LITELLM_CLI_SESSION_TOKEN_PREFIX,
        )

        if source == LITELLM_CLI_SOURCE_IDENTIFIER and key:
            if _is_valid_cli_sso_user_code(user_code):
                return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}:{user_code}"
            return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}"
        else:
            return None

    @staticmethod
    def _get_user_email_and_id_from_result(
        result: Optional[Union[OpenID, dict]],
        generic_client_id: Optional[str] = None,
    ) -> ParsedOpenIDResult:
        """
        Gets the user email and id from the OpenID result after validating the email domain
        """
        user_email: Optional[str] = normalize_email(getattr(result, "email", None))
        user_id: Optional[str] = (
            getattr(result, "id", None) if result is not None else None
        )
        user_role: Optional[str] = None

        if user_email is not None and os.getenv("ALLOWED_EMAIL_DOMAINS") is not None:
            email_domain = user_email.split("@")[1]
            allowed_domains = os.getenv("ALLOWED_EMAIL_DOMAINS").split(",")  # type: ignore
            if email_domain not in allowed_domains:
                raise HTTPException(
                    status_code=401,
                    detail={
                        "message": "The email domain={}, is not an allowed email domain={}. Contact your admin to change this.".format(
                            email_domain, allowed_domains
                        )
                    },
                )

        # Extract user_role from result (works for all SSO providers)
        if result is not None:
            _user_role = getattr(result, "user_role", None)
            if _user_role is not None:
                # Convert enum to string if needed
                user_role = (
                    _user_role.value
                    if isinstance(_user_role, LitellmUserRoles)
                    else _user_role
                )
                verbose_proxy_logger.debug(
                    f"Extracted user_role from SSO result: {user_role}"
                )

        # generic client id - override with custom attribute name if specified
        if generic_client_id is not None and result is not None:
            generic_user_role_attribute_name = os.getenv(
                "GENERIC_USER_ROLE_ATTRIBUTE", "role"
            )
            user_id = getattr(result, "id", None)
            user_email = normalize_email(getattr(result, "email", None))
            if user_role is None:
                _role_from_attr = getattr(result, generic_user_role_attribute_name, None)  # type: ignore
                if _role_from_attr is not None:
                    # Convert enum to string if needed
                    user_role = (
                        _role_from_attr.value
                        if isinstance(_role_from_attr, LitellmUserRoles)
                        else _role_from_attr
                    )

        if user_id is None and result is not None:
            _first_name = getattr(result, "first_name", "") or ""
            _last_name = getattr(result, "last_name", "") or ""
            user_id = _first_name + _last_name

        if user_email is not None and (user_id is None or len(user_id) == 0):
            user_id = user_email

        return ParsedOpenIDResult(
            user_email=user_email,
            user_id=user_id,
            user_role=user_role,
        )

    @staticmethod
    async def get_redirect_response_from_openid(
        result: Union[OpenID, dict, CustomOpenID],
        request: Request,
        received_response: Optional[dict] = None,
        generic_client_id: Optional[str] = None,
        ui_access_mode: Optional[Dict] = None,
        access_token_payload: Optional[dict] = None,
        jwt_handler: Optional[JWTHandler] = None,
        return_to: Optional[str] = None,
    ) -> RedirectResponse:
        import jwt

        from litellm.proxy.proxy_server import (
            general_settings,
            generate_key_helper_fn,
            master_key,
            premium_user,
            proxy_logging_obj,
            redis_usage_cache,
            user_api_key_cache,
            user_custom_sso,
        )
        from litellm.proxy.utils import get_prisma_client_or_throw
        from litellm.types.proxy.ui_sso import ReturnedUITokenObject

        prisma_client = get_prisma_client_or_throw(
            "Prisma client is None, connect a database to your proxy"
        )

        # User is Authe'd in - generate key for the UI to access Proxy
        parsed_openid_result = (
            SSOAuthenticationHandler._get_user_email_and_id_from_result(
                result=result, generic_client_id=generic_client_id
            )
        )
        user_email = parsed_openid_result.get("user_email")
        user_id = parsed_openid_result.get("user_id")
        user_role = parsed_openid_result.get("user_role")
        verbose_proxy_logger.info(f"SSO callback result: {result}")

        user_info = None
        user_id_models: List = []
        max_internal_user_budget = litellm.max_internal_user_budget
        internal_user_budget_duration = litellm.internal_user_budget_duration

        # User might not be already created on first generation of key
        # But if it is, we want their models preferences
        default_ui_key_values: Dict[str, Any] = {
            "duration": LITELLM_UI_SESSION_DURATION,
            "key_max_budget": litellm.max_ui_session_budget,
            "aliases": {},
            "config": {},
            "spend": 0,
            "team_id": "litellm-dashboard",
        }
        user_defined_values: Optional[SSOUserDefinedValues] = None

        if user_custom_sso is not None:
            if inspect.iscoroutinefunction(user_custom_sso):
                user_defined_values = await user_custom_sso(result)  # type: ignore
            else:
                raise ValueError("user_custom_sso must be a coroutine function")
        elif user_id is not None:
            user_defined_values = SSOUserDefinedValues(
                models=user_id_models,
                user_id=user_id,
                user_email=user_email,
                max_budget=max_internal_user_budget,
                user_role=user_role,
                budget_duration=internal_user_budget_duration,
            )

        # (IF SET) Verify user is in restricted SSO group
        SSOAuthenticationHandler.verify_user_in_restricted_sso_group(
            general_settings=general_settings,
            result=result,
            received_response=received_response,
        )

        user_info = await get_user_info_from_db(
            result=result,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            proxy_logging_obj=proxy_logging_obj,
            user_email=user_email,
            user_defined_values=user_defined_values,
            alternate_user_id=user_id,
        )

        # Sync user role from JWT claims via jwt_litellm_role_map (if configured).
        # This ensures SSO users get the same role mapping as API/JWT users.
        # Use the decoded access_token_payload (not received_response) because
        # custom role claims (e.g. custom_roles) are encoded inside the JWT
        # access token, which is stripped from received_response.
        await _sync_user_role_from_jwt_role_map(
            jwt_handler=jwt_handler,
            received_response=access_token_payload or received_response,
            user_info=user_info,
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            user_defined_values=user_defined_values,
        )

        user_defined_values = apply_user_info_values_to_sso_user_defined_values(
            user_info=user_info, user_defined_values=user_defined_values
        )

        if user_defined_values is None:
            raise Exception(
                "Unable to map user identity to known values. 'user_defined_values' is None. File an issue - https://github.com/BerriAI/litellm/issues"
            )

        verbose_proxy_logger.info(
            f"user_defined_values for creating ui key: {user_defined_values}"
        )

        default_ui_key_values.update(user_defined_values)
        default_ui_key_values["request_type"] = "key"
        response = await generate_key_helper_fn(
            **default_ui_key_values,  # type: ignore
            table_name="key",
        )

        key = response["token"]  # type: ignore
        user_id = response["user_id"]  # type: ignore

        user_role = (
            user_defined_values["user_role"]
            or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
        )
        if user_id and isinstance(user_id, str):
            user_role = await check_and_update_if_proxy_admin_id(
                user_role=user_role, user_id=user_id, prisma_client=prisma_client
            )

        verbose_proxy_logger.debug(
            f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
        )
        ## CHECK IF ROLE ALLOWED TO USE PROXY ##
        is_admin_only_access = check_is_admin_only_access(ui_access_mode or {})
        if is_admin_only_access:
            has_access = has_admin_ui_access(user_role or "")
            if not has_access:
                raise HTTPException(
                    status_code=401,
                    detail={
                        "error": f"User not allowed to access proxy. User role={user_role}, proxy mode={ui_access_mode}"
                    },
                )

        disabled_non_admin_personal_key_creation = (
            get_disabled_non_admin_personal_key_creation()
        )
        litellm_dashboard_ui = get_custom_url(
            request_base_url=str(request.base_url), route="ui/"
        )

        if get_secret_bool("EXPERIMENTAL_UI_LOGIN"):
            _user_info: Optional[LiteLLM_UserTable] = None
            if (
                user_defined_values is not None
                and user_defined_values["user_id"] is not None
            ):
                _user_info = LiteLLM_UserTable(
                    user_id=user_defined_values["user_id"],
                    user_role=user_defined_values["user_role"] or user_role,
                    models=[],
                    max_budget=litellm.max_ui_session_budget,
                )
            if _user_info is None:
                raise HTTPException(
                    status_code=401,
                    detail={
                        "error": "User Information is required for experimental UI login"
                    },
                )

            key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(
                _user_info
            )

        returned_ui_token_object = ReturnedUITokenObject(
            user_id=cast(str, user_id),
            key=key,
            user_email=user_email,
            user_role=user_role or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value,
            login_method="sso",
            premium_user=premium_user,
            auth_header_name=general_settings.get(
                "litellm_key_header_name", "Authorization"
            ),
            disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation,
            server_root_path=get_server_root_path(),
        )

        jwt_token = jwt.encode(
            cast(dict, returned_ui_token_object),
            master_key or "",
            algorithm="HS256",
        )

        # Control-plane cross-origin: store JWT behind a single-use opaque
        # code (60s TTL) so the token never appears in browser history / logs.
        # The control plane redeems it via POST /v3/login/exchange.
        if return_to is not None and SSOAuthenticationHandler._validate_return_to(
            return_to
        ):
            code = secrets.token_urlsafe(32)
            cache_key = f"login_code:{code}"
            cache_value = {"token": jwt_token, "redirect_url": return_to}
            if redis_usage_cache is not None:
                await redis_usage_cache.async_set_cache(
                    key=cache_key, value=cache_value, ttl=60
                )
            else:
                await user_api_key_cache.async_set_cache(
                    key=cache_key, value=cache_value, ttl=60
                )

            separator = "&" if "?" in return_to else "?"
            redirect_url = (
                return_to + separator + urlencode({"login": "success", "code": code})
            )
            verbose_proxy_logger.info(
                "Cross-origin SSO: redirecting to control plane with login code"
            )
            redirect_response = RedirectResponse(url=redirect_url, status_code=303)
            redirect_response.delete_cookie("litellm_cp_return_to")
            return redirect_response

        if user_id is not None and isinstance(user_id, str):
            litellm_dashboard_ui += "?login=success"
        verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}")
        redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
        redirect_response.set_cookie(key="token", value=jwt_token)
        return redirect_response

    @staticmethod
    async def prepare_token_exchange_parameters(
        request: Request,
        generic_include_client_id: bool,
    ) -> dict:
        """
        Prepare token exchange parameters for Generic SSO.

        Args:
            request: Request object
            generic_include_client_id: Generic OAuth Client ID

        Returns:
            dict: Token exchange parameters
        """
        # Prepare token exchange parameters (may add code_verifier: str later)
        token_params: Dict[str, Any] = {"include_client_id": generic_include_client_id}

        # Retrieve PKCE code_verifier if PKCE was used in authorization.
        # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip
        # on every non-PKCE SSO callback.
        query_params = dict(request.query_params)
        state = query_params.get("state")

        use_pkce = os.getenv("GENERIC_CLIENT_USE_PKCE", "false").lower() == "true"

        if use_pkce and not state:
            verbose_proxy_logger.warning(
                "PKCE is enabled (GENERIC_CLIENT_USE_PKCE=true) but no 'state' parameter "
                "was found in the callback. The PKCE verifier cannot be retrieved without "
                "a state value — the token exchange will proceed without code_verifier, "
                "which the provider may reject. Ensure your OAuth provider returns 'state' "
                "in the callback redirect."
            )

        if state and use_pkce:
            from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache

            cache_key = f"pkce_verifier:{state}"
            if redis_usage_cache is not None:
                cached_data = await redis_usage_cache.async_get_cache(key=cache_key)
            else:
                cached_data = await user_api_key_cache.async_get_cache(key=cache_key)

            code_verifier = None
            # Track why code_verifier is absent for accurate strict-mode diagnostics.
            _empty_value_in_dict = False  # dict format correct but value is empty/null

            if cached_data:
                # Extract code_verifier from dict (stored as dict for JSON serialization)
                if isinstance(cached_data, dict) and "code_verifier" in cached_data:
                    code_verifier = cached_data["code_verifier"]
                    if not code_verifier:
                        # Dict format is correct but value is empty or null.  This is
                        # a distinct case from an unrecognized format — the entry exists
                        # but was stored with an empty/null verifier (data integrity issue).
                        _empty_value_in_dict = True
                        verbose_proxy_logger.warning(
                            "PKCE verifier dict for state '%s' has an empty/null code_verifier "
                            "value — may indicate a storage bug. Treating as a cache miss.",
                            state,
                        )
                    else:
                        verbose_proxy_logger.debug(
                            "PKCE code_verifier retrieved from cache"
                        )
                elif isinstance(cached_data, str):
                    # Handle legacy format (plain string) for backward compatibility
                    code_verifier = cached_data
                    verbose_proxy_logger.warning(
                        "Retrieved code_verifier in legacy plain-string format. "
                        "Future storage will use dict format."
                    )
                else:
                    # Defer the detailed ERROR log to the strict-mode branch below
                    # (which includes state and a diagnostic message).  Log at DEBUG
                    # here to avoid duplicate ERROR entries in the same request.
                    verbose_proxy_logger.debug(
                        "Unexpected PKCE verifier cache format (type=%s); skipping.",
                        type(cached_data).__name__,
                    )

            if code_verifier:
                # Add code_verifier to token exchange parameters.
                token_params["code_verifier"] = code_verifier
                # Return the cache key so the caller can delete it *after* a
                # successful token exchange (avoids losing the verifier on retry
                # if the exchange fails partway through).
                token_params["_pkce_cache_key"] = cache_key
            else:
                await SSOAuthenticationHandler._handle_missing_pkce_verifier(
                    state=state,
                    cache_key=cache_key,
                    cached_data=cached_data,
                    empty_value_in_dict=_empty_value_in_dict,
                    redis_usage_cache=redis_usage_cache,
                    user_api_key_cache=user_api_key_cache,
                )
        return token_params

    @staticmethod
    async def _handle_missing_pkce_verifier(
        state: Optional[str],
        cache_key: str,
        cached_data: object,
        empty_value_in_dict: bool,
        redis_usage_cache: object,
        user_api_key_cache: object,
    ) -> None:
        """Handle the case where PKCE verifier could not be extracted from cache.

        In strict mode (PKCE_STRICT_CACHE_MISS=true) raises ProxyException.
        Otherwise logs a warning and returns (token exchange proceeds without verifier).
        """
        active_cache = (
            redis_usage_cache if redis_usage_cache is not None else user_api_key_cache
        )
        strict_cache_miss = (
            os.getenv("PKCE_STRICT_CACHE_MISS", "false").lower() == "true"
        )
        if strict_cache_miss:
            if empty_value_in_dict:
                await SSOAuthenticationHandler._delete_pkce_verifier(cache_key)
                raise ProxyException(
                    message=(
                        f"PKCE verifier for state '{state}' was found in cache but "
                        f"has an empty or null code_verifier value — possible storage bug."
                    ),
                    type=ProxyErrorTypes.auth_error,
                    param="PKCE_CACHE_MISS",
                    code=status.HTTP_401_UNAUTHORIZED,
                )
            elif cached_data is not None:
                await SSOAuthenticationHandler._delete_pkce_verifier(cache_key)
                verbose_proxy_logger.error(
                    "PKCE verifier for state '%s' has an unrecognized format (type=%s); "
                    "treating as a cache miss. Investigate the cached value — it may be "
                    "a corrupt or stale entry.",
                    state,
                    type(cached_data).__name__,
                )
                raise ProxyException(
                    message=(
                        f"PKCE verifier for state '{state}' has an unrecognized format "
                        f"(type={type(cached_data).__name__}). The cached entry may be corrupt."
                    ),
                    type=ProxyErrorTypes.auth_error,
                    param="PKCE_CACHE_MISS",
                    code=status.HTTP_401_UNAUTHORIZED,
                )
            else:
                if redis_usage_cache is not None:
                    cause = (
                        "The authorization and callback were likely handled by different "
                        "instances — the verifier was stored on one pod but not found on another."
                    )
                else:
                    cause = (
                        "The verifier may have expired (TTL), been lost on a pod restart, "
                        "or the PKCE authorization step was never completed. "
                        "Configure Redis so all proxy instances share the PKCE verifier."
                    )
                verbose_proxy_logger.error(
                    "PKCE is enabled but no verifier found in cache for state '%s'. "
                    "%s Cache type: %s.",
                    state,
                    cause,
                    type(active_cache).__name__,
                )
                raise ProxyException(
                    message=f"PKCE verifier not found in cache for state '{state}'. {cause}",
                    type=ProxyErrorTypes.auth_error,
                    param="PKCE_CACHE_MISS",
                    code=status.HTTP_401_UNAUTHORIZED,
                )
        else:
            if cached_data is not None:
                await SSOAuthenticationHandler._delete_pkce_verifier(cache_key)
            verbose_proxy_logger.warning(
                "PKCE is enabled but verifier not found in cache for state '%s' "
                "(cache type: %s, raw data present: %s). "
                "Continuing without code_verifier — set PKCE_STRICT_CACHE_MISS=true to fail fast instead.",
                state,
                type(active_cache).__name__,
                cached_data is not None,
            )

    @staticmethod
    async def _delete_pkce_verifier(cache_key: str) -> None:
        """Delete a single-use PKCE verifier from cache after a successful exchange.

        Failure is non-fatal: a leftover verifier is a minor security concern
        (unused key in cache) but not worth aborting an otherwise-successful login.
        """
        from litellm.proxy.proxy_server import redis_usage_cache, user_api_key_cache

        try:
            if redis_usage_cache is not None:
                await redis_usage_cache.async_delete_cache(key=cache_key)
            else:
                await user_api_key_cache.async_delete_cache(key=cache_key)
        except Exception as exc:
            verbose_proxy_logger.warning(
                "PKCE: failed to delete verifier cache key '%s' (best-effort cleanup): %s",
                cache_key,
                exc,
            )

    @staticmethod
    def generate_pkce_params() -> Tuple[str, str]:
        """
        Generate PKCE (Proof Key for Code Exchange) parameters for OAuth 2.0.

        Returns:
            Tuple[str, str]: (code_verifier, code_challenge)
            - code_verifier: Random 43-128 character string (we use 43 for efficiency)
            - code_challenge: Base64-URL-encoded SHA256 hash of the code_verifier

        Reference: https://datatracker.ietf.org/doc/html/rfc7636
        """
        # Generate a cryptographically random code_verifier (43 characters)
        # Using 32 random bytes which becomes 43 characters when base64-url-encoded
        code_verifier = (
            base64.urlsafe_b64encode(secrets.token_bytes(32))
            .decode("utf-8")
            .rstrip("=")
        )

        # Generate code_challenge using S256 method (SHA256)
        code_challenge_bytes = hashlib.sha256(code_verifier.encode("utf-8")).digest()
        code_challenge = (
            base64.urlsafe_b64encode(code_challenge_bytes).decode("utf-8").rstrip("=")
        )

        return code_verifier, code_challenge

    @staticmethod
    def _validate_token_response(response: "httpx.Response") -> dict:
        """
        Parse and validate the token endpoint response.

        Ensures the response is valid JSON, a dict, and contains a non-null
        access_token string. Raises ProxyException on any validation failure.
        """
        try:
            token_response_raw = response.json()
        except Exception as json_err:
            verbose_proxy_logger.error(
                "Failed to parse token response as JSON: %s. Body: %s",
                json_err,
                response.text[:500],
            )
            raise ProxyException(
                message=f"Token endpoint returned invalid JSON: {json_err}",
                type=ProxyErrorTypes.auth_error,
                param="token_exchange",
                code=status.HTTP_401_UNAUTHORIZED,
            )

        if not isinstance(token_response_raw, dict):
            verbose_proxy_logger.error(
                "Token endpoint returned non-dict JSON (type=%s). Body: %s",
                type(token_response_raw).__name__,
                response.text[:500],
            )
            raise ProxyException(
                message=(
                    f"Token endpoint returned unexpected response format "
                    f"(expected JSON object, got {type(token_response_raw).__name__})"
                ),
                type=ProxyErrorTypes.auth_error,
                param="token_exchange",
                code=status.HTTP_401_UNAUTHORIZED,
            )
        token_response: dict = token_response_raw

        access_token_val = token_response.get("access_token")
        if not isinstance(access_token_val, str) or not access_token_val:
            error = token_response.get("error")
            error_desc = token_response.get("error_description", "")
            if error:
                detail = f"{error} - {error_desc}" if error_desc else error
            else:
                detail = (
                    "token endpoint returned HTTP 200 but no access_token "
                    f"(response keys: {sorted(token_response.keys())})"
                )
            verbose_proxy_logger.error(
                "Token response missing or null access_token. detail=%s", detail
            )
            raise ProxyException(
                message=f"Token exchange failed: {detail}",
                type=ProxyErrorTypes.auth_error,
                param="token_exchange",
                code=status.HTTP_401_UNAUTHORIZED,
            )

        return token_response

    @staticmethod
    async def _pkce_token_exchange(
        authorization_code: str,
        code_verifier: str,
        client_id: str,
        client_secret: Optional[str],
        token_endpoint: str,
        userinfo_endpoint: Optional[str],
        include_client_id: bool,
        redirect_url: Optional[str],
        additional_headers: Dict[str, str],
    ) -> dict:
        """
        Performs a direct OAuth token exchange including the PKCE code_verifier.

        fastapi-sso does not forward code_verifier, so when PKCE is enabled we
        bypass it and call the token endpoint ourselves, then fetch user info.

        Returns a combined dict of the token response and user info, suitable
        for passing to a response_convertor.
        """
        verbose_proxy_logger.debug(
            "PKCE: performing direct token exchange (code_verifier length=%d)",
            len(code_verifier),
        )

        token_data: Dict[str, str] = {
            "grant_type": "authorization_code",
            "code": authorization_code,
            "code_verifier": code_verifier,
        }
        # Only include redirect_uri when set — omitting it avoids sending the
        # literal string "None" to the provider if the env var is missing.
        if redirect_url:
            token_data["redirect_uri"] = redirect_url

        request_headers = {
            **additional_headers,
            "Content-Type": "application/x-www-form-urlencoded",  # must not be overridden
            "Accept": "application/json",
        }

        if not include_client_id:
            # Use Basic Auth only when a secret is available; public PKCE clients omit it.
            if client_secret:
                credentials = base64.b64encode(
                    f"{client_id}:{client_secret}".encode()
                ).decode()
                request_headers["Authorization"] = f"Basic {credentials}"
            else:
                token_data["client_id"] = client_id
        else:
            token_data["client_id"] = client_id
            if client_secret:
                token_data["client_secret"] = client_secret

        http_client = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.SSO_HANDLER
        )
        try:
            response = await http_client.post(
                url=token_endpoint,
                data=token_data,
                headers=request_headers,
                timeout=30.0,
            )
        except Exception as exc:
            # Catch network-level errors (SSL, DNS, TCP, timeout, etc.) and
            # wrap them as a clean ProxyException rather than leaking raw
            # httpx or OS exceptions to callers.
            verbose_proxy_logger.error("PKCE token endpoint unreachable: %s", exc)
            raise ProxyException(
                message=f"Token endpoint request failed: {exc}",
                type=ProxyErrorTypes.auth_error,
                param="token_exchange",
                code=status.HTTP_401_UNAUTHORIZED,
            ) from exc
        if response.status_code != 200:
            verbose_proxy_logger.error(
                "PKCE token exchange failed. status=%s body=%s",
                response.status_code,
                response.text[:500],
            )
            raise ProxyException(
                message=f"Token exchange failed: {response.status_code} - {response.text[:500]}",
                type=ProxyErrorTypes.auth_error,
                param="token_exchange",
                code=status.HTTP_401_UNAUTHORIZED,
            )

        token_response = SSOAuthenticationHandler._validate_token_response(response)

        verbose_proxy_logger.debug(
            "PKCE token exchange successful. id_token_present=%s",
            bool(token_response.get("id_token")),
        )
        # Bearer credentials (access_token, id_token, refresh_token) are always sourced
        # from token_response — not from userinfo — in the merge step below.
        userinfo = await SSOAuthenticationHandler._get_pkce_userinfo(
            access_token=token_response["access_token"],
            id_token=token_response.get("id_token"),
            userinfo_endpoint=userinfo_endpoint,
            additional_headers=additional_headers,
        )

        # Merge: userinfo takes precedence for identity claims (sub, email, name, …) per
        # the OpenID Connect spec (userinfo is the authoritative source for identity).
        # Bearer credentials (access_token, id_token, refresh_token) from the token endpoint
        # take precedence over same-named fields in userinfo — non-standard providers sometimes
        # include token fields in userinfo, which must not shadow the real bearer token.
        # If a bearer field is absent from the token response, any userinfo-provided value
        # is preserved as a fallback (useful for non-standard providers that omit id_token
        # from the token response but include it in userinfo).
        #
        # Three-way merge semantics for each bearer-credential field:
        #   1. token_response has a non-null value → use it (token endpoint is authoritative)
        #   2. token_response explicitly sent null  → remove the key so callers get a clean
        #      absence signal; the null from the token endpoint overrides userinfo too
        #   3. field absent from token_response     → leave whatever userinfo provided as-is
        #      (e.g. userinfo-provided id_token from a non-standard provider)
        merged = {**token_response, **userinfo}
        for field in _OAUTH_TOKEN_FIELDS:
            if token_response.get(field) is not None:
                # Case 1: non-null in token_response — restore authoritative value.
                merged[field] = token_response[field]
            elif field in token_response:
                # Case 2: key exists but value is explicitly null — remove from merged.
                merged.pop(field, None)
            # Case 3: field absent from token_response — leave userinfo value as-is.
        return merged

    @staticmethod
    async def _get_pkce_userinfo(
        access_token: str,
        id_token: Optional[str],
        userinfo_endpoint: Optional[str],
        additional_headers: Dict[str, str],
    ) -> dict:
        """
        Fetches user info from the userinfo endpoint.
        Falls back to decoding the id_token if the endpoint is unavailable.
        """
        # None = request not yet attempted, failed, or returned empty/null (treated as failure
        # so the id_token fallback can be attempted instead of returning a session with no claims).
        userinfo: Optional[dict] = None

        if userinfo_endpoint:
            try:
                client = get_async_httpx_client(
                    llm_provider=httpxSpecialProvider.SSO_HANDLER
                )
                resp = await client.get(
                    url=userinfo_endpoint,
                    headers={
                        **additional_headers,
                        "Authorization": f"Bearer {access_token}",  # must not be overridden
                    },
                )
                if resp.status_code == 200:
                    try:
                        userinfo_raw = resp.json()
                        if not userinfo_raw:
                            # JSON null (None) or empty dict ({}) — no identity claims.
                            # Treat as failure so id_token fallback can be attempted.
                            verbose_proxy_logger.warning(
                                "Userinfo endpoint returned an empty or null response "
                                "(type=%s); treating as failure and attempting id_token fallback. "
                                "Check your provider's userinfo endpoint configuration.",
                                type(userinfo_raw).__name__,
                            )
                            userinfo = None
                        else:
                            userinfo = userinfo_raw
                    except Exception as json_err:
                        verbose_proxy_logger.warning(
                            "Userinfo endpoint returned non-JSON response (status 200): %s",
                            json_err,
                        )
                else:
                    verbose_proxy_logger.warning(
                        "Userinfo endpoint returned %s (body: %s), falling back to id_token",
                        resp.status_code,
                        resp.text[:500],
                    )
            except Exception as e:
                verbose_proxy_logger.warning(
                    "Userinfo endpoint error: %s, falling back to id_token", e
                )

        # Only fall back to id_token when the userinfo request failed (None).
        # Empty dict ({}) and JSON null are both treated as failure (set to None above) since
        # they contain no identity claims — id_token fallback is attempted in that case too.
        # Explicitly check for a non-empty string to avoid attempting JWT decode on
        # a blank or non-string id_token field from a misbehaving provider.
        if userinfo is None and isinstance(id_token, str) and id_token:
            try:
                userinfo = jwt.decode(id_token, options={"verify_signature": False})
                if not userinfo:
                    # jwt.decode returned an empty dict (payload-free JWT or provider bug).
                    # Treat this the same as a missing userinfo — the session would have no
                    # identity claims, which is equivalent to a broken session.
                    verbose_proxy_logger.warning(
                        "id_token decoded to an empty payload — treating as failure."
                    )
                    userinfo = None
            except Exception as decode_err:
                verbose_proxy_logger.error("Failed to decode id_token: %s", decode_err)
                raise ProxyException(
                    message=f"Failed to decode id_token JWT: {decode_err}",
                    type=ProxyErrorTypes.auth_error,
                    param="userinfo",
                    code=status.HTTP_401_UNAUTHORIZED,
                )

        if userinfo is None:
            id_token_attempted = isinstance(id_token, str) and bool(id_token)
            if userinfo_endpoint:
                if id_token_attempted:
                    detail = (
                        "userinfo endpoint failed and id_token was present but "
                        "decoded to an empty payload — no identity claims available"
                    )
                else:
                    detail = "userinfo endpoint failed and no id_token was present in the token response"
            else:
                if id_token_attempted:
                    detail = (
                        "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) "
                        "and id_token decoded to an empty payload — no identity claims available"
                    )
                else:
                    detail = "no userinfo endpoint is configured (GENERIC_USERINFO_ENDPOINT) and no id_token was present"
            raise ProxyException(
                message=f"SSO user info unavailable: {detail}.",
                type=ProxyErrorTypes.auth_error,
                param="userinfo",
                code=status.HTTP_401_UNAUTHORIZED,
            )

        return userinfo


class MicrosoftSSOHandler:
    """
    Handles Microsoft SSO callback response and returns a CustomOpenID object
    """

    graph_api_base_url = "https://graph.microsoft.com/v1.0"
    graph_api_user_groups_endpoint = f"{graph_api_base_url}/me/memberOf"

    """
    Constants
    """
    MAX_GRAPH_API_PAGES = 200

    # used for debugging to show the user groups litellm found from Graph API
    GRAPH_API_RESPONSE_KEY = "graph_api_user_groups"

    @staticmethod
    async def get_microsoft_callback_response(
        request: Request,
        microsoft_client_id: str,
        redirect_url: str,
        return_raw_sso_response: bool = False,
    ) -> Union[CustomOpenID, OpenID, dict]:
        """
        Get the Microsoft SSO callback response

        Args:
            return_raw_sso_response: If True, return the raw SSO response
        """
        microsoft_client_secret = os.getenv("MICROSOFT_CLIENT_SECRET", None)
        microsoft_tenant = os.getenv("MICROSOFT_TENANT", None)
        if microsoft_client_secret is None:
            raise ProxyException(
                message="MICROSOFT_CLIENT_SECRET not set. Set it in .env file",
                type=ProxyErrorTypes.auth_error,
                param="MICROSOFT_CLIENT_SECRET",
                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )
        if microsoft_tenant is None:
            raise ProxyException(
                message="MICROSOFT_TENANT not set. Set it in .env file",
                type=ProxyErrorTypes.auth_error,
                param="MICROSOFT_TENANT",
                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )
        microsoft_sso = CustomMicrosoftSSO(
            client_id=microsoft_client_id,
            client_secret=microsoft_client_secret,
            tenant=microsoft_tenant,
            redirect_uri=redirect_url,
            allow_insecure_http=True,
        )
        original_msft_result = (
            await microsoft_sso.verify_and_process(
                request=request,
                convert_response=False,  # type: ignore
            )
            or {}
        )

        user_team_ids = await MicrosoftSSOHandler.get_user_groups_from_graph_api(
            access_token=microsoft_sso.access_token
        )

        # Extract app roles from the id_token JWT
        app_roles = MicrosoftSSOHandler.get_app_roles_from_id_token(
            id_token=microsoft_sso.id_token
        )
        verbose_proxy_logger.debug(f"Extracted app roles from id_token: {app_roles}")

        # Combine groups and app roles
        user_role: Optional[LitellmUserRoles] = None
        if app_roles:
            # Check if any app role is a valid LitellmUserRoles
            for role_str in app_roles:
                role = get_litellm_user_role(role_str)
                if role is not None:
                    user_role = role
                    verbose_proxy_logger.debug(
                        f"Found valid LitellmUserRoles '{role.value}' in app_roles"
                    )
                    break

        verbose_proxy_logger.debug(
            f"Combined team_ids (groups + app roles): {user_team_ids}"
        )

        # if user is trying to get the raw sso response for debugging, return the raw sso response
        if return_raw_sso_response:
            original_msft_result[MicrosoftSSOHandler.GRAPH_API_RESPONSE_KEY] = (
                user_team_ids
            )
            original_msft_result["app_roles"] = app_roles
            return original_msft_result or {}

        result = MicrosoftSSOHandler.openid_from_response(
            response=original_msft_result,
            team_ids=user_team_ids,
            user_role=user_role,
        )
        return result

    @staticmethod
    def openid_from_response(
        response: Optional[dict],
        team_ids: List[str],
        user_role: Optional[LitellmUserRoles],
    ) -> CustomOpenID:
        response = response or {}
        verbose_proxy_logger.debug(f"Microsoft SSO Callback Response: {response}")
        openid_response = CustomOpenID(
            email=normalize_email(
                response.get(MICROSOFT_USER_EMAIL_ATTRIBUTE) or response.get("mail")
            ),
            display_name=response.get(MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE),
            provider="microsoft",
            id=response.get(MICROSOFT_USER_ID_ATTRIBUTE),
            first_name=response.get(MICROSOFT_USER_FIRST_NAME_ATTRIBUTE),
            last_name=response.get(MICROSOFT_USER_LAST_NAME_ATTRIBUTE),
            team_ids=team_ids,
            user_role=user_role,
        )
        verbose_proxy_logger.debug(f"Microsoft SSO OpenID Response: {openid_response}")
        return openid_response

    @staticmethod
    def get_app_roles_from_id_token(id_token: Optional[str]) -> List[str]:
        """
        Extract app roles from the Microsoft Entra ID (Azure AD) id_token JWT.

        App roles are assigned in the Azure AD Enterprise Application and appear
        in the 'app_roles' claim of the id_token.

        Args:
            id_token (Optional[str]): The JWT id_token from Microsoft SSO

        Returns:
            List[str]: List of app role names assigned to the user
        """
        if not id_token:
            verbose_proxy_logger.debug("No id_token provided for app role extraction")
            return []

        try:
            import jwt

            # Decode the JWT without signature verification
            # (signature is already verified by fastapi_sso)
            decoded_token = jwt.decode(id_token, options={"verify_signature": False})

            # Extract app_roles claim from the token
            ## check for both 'roles' and 'app_roles' claims
            roles = decoded_token.get("app_roles", []) or decoded_token.get("roles", [])

            if roles and isinstance(roles, list):
                verbose_proxy_logger.debug(
                    f"Found {len(roles)} app role(s) in id_token: {roles}"
                )
                return roles
            else:
                verbose_proxy_logger.debug(
                    "No app roles found in id_token or roles claim is not a list"
                )
                return []

        except Exception as e:
            verbose_proxy_logger.error(f"Error extracting app roles from id_token: {e}")
            return []

    @staticmethod
    async def get_user_groups_from_graph_api(
        access_token: Optional[str] = None,
    ) -> List[str]:
        """
        Returns a list of `team_ids` the user belongs to from the Microsoft Graph API

        Args:
            access_token (Optional[str]): Microsoft Graph API access token

        Returns:
            List[str]: List of group IDs the user belongs to
        """
        try:
            async_client = get_async_httpx_client(
                llm_provider=httpxSpecialProvider.SSO_HANDLER
            )

            # Handle MSFT Enterprise Application Groups
            service_principal_id = os.getenv("MICROSOFT_SERVICE_PRINCIPAL_ID", None)
            service_principal_group_ids: Optional[List[str]] = []
            service_principal_teams: Optional[List[MicrosoftServicePrincipalTeam]] = []
            if service_principal_id:
                (
                    service_principal_group_ids,
                    service_principal_teams,
                ) = await MicrosoftSSOHandler.get_group_ids_from_service_principal(
                    service_principal_id=service_principal_id,
                    async_client=async_client,
                    access_token=access_token,
                )
                verbose_proxy_logger.debug(
                    f"Service principal group IDs: {service_principal_group_ids}"
                )
                if len(service_principal_group_ids) > 0:
                    await MicrosoftSSOHandler.create_litellm_teams_from_service_principal_team_ids(
                        service_principal_teams=service_principal_teams,
                    )

            # Fetch user membership from Microsoft Graph API
            all_group_ids = []
            next_link: Optional[str] = (
                MicrosoftSSOHandler.graph_api_user_groups_endpoint
            )
            auth_headers = {"Authorization": f"Bearer {access_token}"}
            page_count = 0

            while (
                next_link is not None
                and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES
            ):
                group_ids, next_link = await MicrosoftSSOHandler.fetch_and_parse_groups(
                    url=next_link, headers=auth_headers, async_client=async_client
                )
                all_group_ids.extend(group_ids)
                page_count += 1

            if (
                next_link is not None
                and page_count >= MicrosoftSSOHandler.MAX_GRAPH_API_PAGES
            ):
                verbose_proxy_logger.warning(
                    f"Reached maximum page limit of {MicrosoftSSOHandler.MAX_GRAPH_API_PAGES}. Some groups may not be included."
                )

            # If service_principal_group_ids is not empty, only return group_ids that are in both all_group_ids and service_principal_group_ids
            if service_principal_group_ids and len(service_principal_group_ids) > 0:
                all_group_ids = [
                    group_id
                    for group_id in all_group_ids
                    if group_id in service_principal_group_ids
                ]

            return all_group_ids

        except Exception as e:
            verbose_proxy_logger.error(
                f"Error getting user groups from Microsoft Graph API: {e}"
            )
            return []

    @staticmethod
    async def fetch_and_parse_groups(
        url: str, headers: dict, async_client: AsyncHTTPHandler
    ) -> Tuple[List[str], Optional[str]]:
        """Helper function to fetch and parse group data from a URL"""
        response = await async_client.get(url, headers=headers)
        response_json = response.json()
        response_typed = await MicrosoftSSOHandler._cast_graph_api_response_dict(
            response=response_json
        )
        group_ids = MicrosoftSSOHandler._get_group_ids_from_graph_api_response(
            response=response_typed
        )
        return group_ids, response_typed.get("odata_nextLink")

    @staticmethod
    def _get_group_ids_from_graph_api_response(
        response: MicrosoftGraphAPIUserGroupResponse,
    ) -> List[str]:
        group_ids = []
        for _object in response.get("value", []) or []:
            _group_id = _object.get("id")
            if _group_id is not None:
                group_ids.append(_group_id)
        return group_ids

    @staticmethod
    async def _cast_graph_api_response_dict(
        response: dict,
    ) -> MicrosoftGraphAPIUserGroupResponse:
        directory_objects: List[MicrosoftGraphAPIUserGroupDirectoryObject] = []
        for _object in response.get("value", []):
            directory_objects.append(
                MicrosoftGraphAPIUserGroupDirectoryObject(
                    odata_type=_object.get("@odata.type"),
                    id=_object.get("id"),
                    deletedDateTime=_object.get("deletedDateTime"),
                    description=_object.get("description"),
                    displayName=_object.get("displayName"),
                    roleTemplateId=_object.get("roleTemplateId"),
                )
            )
        return MicrosoftGraphAPIUserGroupResponse(
            odata_context=response.get("@odata.context"),
            odata_nextLink=response.get("@odata.nextLink"),
            value=directory_objects,
        )

    @staticmethod
    async def get_group_ids_from_service_principal(
        service_principal_id: str,
        async_client: AsyncHTTPHandler,
        access_token: Optional[str] = None,
    ) -> Tuple[List[str], List[MicrosoftServicePrincipalTeam]]:
        """
        Gets the groups belonging to the Service Principal Application

        Service Principal Id is an `Enterprise Application` in Azure AD

        Users use Enterprise Applications to manage Groups and Users on Microsoft Entra ID
        """
        base_url = "https://graph.microsoft.com/v1.0"
        # Endpoint to get app role assignments for the given service principal
        endpoint = f"/servicePrincipals/{service_principal_id}/appRoleAssignedTo"
        url = base_url + endpoint

        headers = {
            "Authorization": f"Bearer {access_token}",
            "Content-Type": "application/json",
        }

        response = await async_client.get(url, headers=headers)
        response_json = response.json()
        verbose_proxy_logger.debug(
            f"Response from service principal app role assigned to: {response_json}"
        )
        group_ids: List[str] = []
        service_principal_teams: List[MicrosoftServicePrincipalTeam] = []

        for _object in response_json.get("value", []):
            if _object.get("principalType") == "Group":
                # Append the group ID to the list
                group_ids.append(_object.get("principalId"))
                # Append the service principal team to the list
                service_principal_teams.append(
                    MicrosoftServicePrincipalTeam(
                        principalDisplayName=_object.get("principalDisplayName"),
                        principalId=_object.get("principalId"),
                    )
                )

        return group_ids, service_principal_teams

    @staticmethod
    async def create_litellm_teams_from_service_principal_team_ids(
        service_principal_teams: List[MicrosoftServicePrincipalTeam],
    ):
        """
        Creates Litellm Teams from the Service Principal Group IDs

        When a user sets a `SERVICE_PRINCIPAL_ID` in the env, litellm will fetch groups under that service principal and create Litellm Teams from them
        """
        verbose_proxy_logger.debug(
            f"Creating Litellm Teams from Service Principal Teams: {service_principal_teams}"
        )
        for service_principal_team in service_principal_teams:
            litellm_team_id: Optional[str] = service_principal_team.get("principalId")
            litellm_team_name: Optional[str] = service_principal_team.get(
                "principalDisplayName"
            )
            if not litellm_team_id:
                verbose_proxy_logger.debug(
                    f"Skipping team creation for {litellm_team_name} because it has no principalId"
                )
                continue

            await SSOAuthenticationHandler.create_litellm_team_from_sso_group(
                litellm_team_id=litellm_team_id,
                litellm_team_name=litellm_team_name,
            )


class GoogleSSOHandler:
    """
    Handles Google SSO callback response and returns a CustomOpenID object
    """

    @staticmethod
    async def get_google_callback_response(
        request: Request,
        google_client_id: str,
        redirect_url: str,
        return_raw_sso_response: bool = False,
    ) -> Union[OpenID, dict]:
        """
        Get the Google SSO callback response

        Args:
            return_raw_sso_response: If True, return the raw SSO response
        """
        from fastapi_sso.sso.google import GoogleSSO

        google_client_secret = os.getenv("GOOGLE_CLIENT_SECRET", None)
        if google_client_secret is None:
            raise ProxyException(
                message="GOOGLE_CLIENT_SECRET not set. Set it in .env file",
                type=ProxyErrorTypes.auth_error,
                param="GOOGLE_CLIENT_SECRET",
                code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )
        google_sso = GoogleSSO(
            client_id=google_client_id,
            redirect_uri=redirect_url,
            client_secret=google_client_secret,
        )

        # if user is trying to get the raw sso response for debugging, return the raw sso response
        if return_raw_sso_response:
            return (
                await google_sso.verify_and_process(
                    request=request,
                    convert_response=False,  # type: ignore
                )
                or {}
            )

        result = await google_sso.verify_and_process(request)
        return result or {}


@router.get("/sso/debug/login", tags=["experimental"], include_in_schema=False)
async def debug_sso_login(request: Request):
    """
    Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env
    PROXY_BASE_URL should be the your deployed proxy endpoint, e.g. PROXY_BASE_URL="https://litellm-production-7002.up.railway.app/"
    Example:
    """
    from litellm.proxy.proxy_server import premium_user

    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
    generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)

    ####### Check if user is a Enterprise / Premium User #######
    if (
        microsoft_client_id is not None
        or google_client_id is not None
        or generic_client_id is not None
    ):
        if premium_user is not True:
            raise ProxyException(
                message="You must be a LiteLLM Enterprise user to use SSO. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://enterprise.litellm.ai/demo You are seeing this error message because You set one of `MICROSOFT_CLIENT_ID`, `GOOGLE_CLIENT_ID`, or `GENERIC_CLIENT_ID` in your env. Please unset this",
                type=ProxyErrorTypes.auth_error,
                param="premium_user",
                code=status.HTTP_403_FORBIDDEN,
            )

    # get url from request
    redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
        request=request,
        sso_callback_route="sso/debug/callback",
    )

    # Check if we should use SSO handler
    if (
        SSOAuthenticationHandler.should_use_sso_handler(
            microsoft_client_id=microsoft_client_id,
            google_client_id=google_client_id,
            generic_client_id=generic_client_id,
        )
        is True
    ):
        return await SSOAuthenticationHandler.get_sso_login_redirect(
            redirect_url=redirect_url,
            microsoft_client_id=microsoft_client_id,
            google_client_id=google_client_id,
            generic_client_id=generic_client_id,
            request=request,
        )


@router.get("/sso/debug/callback", tags=["experimental"], include_in_schema=False)
async def debug_sso_callback(request: Request):
    """
    Returns the OpenID object returned by the SSO provider
    """
    import json

    from fastapi.responses import HTMLResponse

    from litellm.proxy._types import LiteLLM_JWTAuth
    from litellm.proxy.auth.handle_jwt import JWTHandler
    from litellm.proxy.proxy_server import (
        general_settings,
        jwt_handler,
        prisma_client,
        user_api_key_cache,
    )

    sso_jwt_handler: Optional[JWTHandler] = None
    ui_access_mode = general_settings.get("ui_access_mode", None)
    if ui_access_mode is not None and isinstance(ui_access_mode, dict):
        sso_jwt_handler = JWTHandler()
        sso_jwt_handler.update_environment(
            prisma_client=prisma_client,
            user_api_key_cache=user_api_key_cache,
            litellm_jwtauth=LiteLLM_JWTAuth(
                team_ids_jwt_field=general_settings.get("ui_access_mode", {}).get(
                    "sso_group_jwt_field", None
                ),
            ),
            leeway=0,
        )

    microsoft_client_id = os.getenv("MICROSOFT_CLIENT_ID", None)
    google_client_id = os.getenv("GOOGLE_CLIENT_ID", None)
    generic_client_id = os.getenv("GENERIC_CLIENT_ID", None)

    redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
    if redirect_url.endswith("/"):
        redirect_url += "sso/debug/callback"
    else:
        redirect_url += "/sso/debug/callback"

    result = None
    received_response: Optional[dict] = None
    access_token_payload: Optional[dict] = None
    if google_client_id is not None:
        result = await GoogleSSOHandler.get_google_callback_response(
            request=request,
            google_client_id=google_client_id,
            redirect_url=redirect_url,
            return_raw_sso_response=True,
        )
    elif microsoft_client_id is not None:
        result = await MicrosoftSSOHandler.get_microsoft_callback_response(
            request=request,
            microsoft_client_id=microsoft_client_id,
            redirect_url=redirect_url,
            return_raw_sso_response=True,
        )

    elif generic_client_id is not None:
        result, received_response, access_token_payload = (
            await get_generic_sso_response(
                request=request,
                jwt_handler=jwt_handler,
                generic_client_id=generic_client_id,
                redirect_url=redirect_url,
                sso_jwt_handler=sso_jwt_handler,
            )
        )

    # If result is None, return a basic error message
    if result is None:
        return HTMLResponse(
            content="<h1>SSO Authentication Failed</h1><p>No data was returned from the SSO provider.</p>",
            status_code=400,
        )

    # Convert the OpenID object to a dictionary
    if hasattr(result, "__dict__"):
        result_dict = result.__dict__
    else:
        result_dict = dict(result)

    # Filter out any None values and convert to JSON serializable format
    filtered_result = {}
    for key, value in result_dict.items():
        if value is not None and not key.startswith("_"):
            if isinstance(value, (str, int, float, bool)) or value is None:
                filtered_result[key] = value
            else:
                try:
                    # Try to convert to string or another JSON serializable format
                    filtered_result[key] = str(value)
                except Exception as e:
                    filtered_result[key] = f"Complex value (not displayable): {str(e)}"

    # Defense-in-depth: ensure no bearer tokens leak into the rendered HTML even if
    # a non-conforming IdP places them in its userinfo response.
    safe_raw_claims = {
        k: v
        for k, v in (received_response or {}).items()
        if k not in _OAUTH_TOKEN_FIELDS
    }
    safe_access_token_claims = {
        k: v
        for k, v in (access_token_payload or {}).items()
        if k not in _OAUTH_TOKEN_FIELDS
    }

    sso_payload = {
        "parsed_by_proxy": filtered_result,
        "raw_claims": safe_raw_claims,
        "access_token_claims": safe_access_token_claims,
    }

    # Replace the placeholder in the template with the actual data
    sso_payload_json = json.dumps(sso_payload, indent=2, default=str).replace(
        "</", "<\\/"
    )
    html_content = jwt_display_template.replace(
        "const ssoData = SSO_DATA;",
        f"const ssoData = {sso_payload_json};",
    )

    return HTMLResponse(content=html_content)
