# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/services/oauth_manager.py
Copyright 2026
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

OAuth 2.0 Manager for ContextForge.

This module handles OAuth 2.0 authentication flows including:
- Client Credentials (Machine-to-Machine)
- Authorization Code (User Delegation)
"""

# Standard
import asyncio
import base64
from datetime import datetime, timedelta, timezone
import hashlib
import logging
import re
import secrets
from typing import Any, Dict, Optional
from urllib.parse import parse_qsl, quote, urlparse

# Third-Party
import httpx
import orjson
from requests_oauthlib import OAuth2Session

# First-Party
from mcpgateway.common.validators import SecurityValidator, validate_core_url
from mcpgateway.config import get_settings
from mcpgateway.services.encryption_service import decrypt_oauth_config_for_runtime, get_encryption_service
from mcpgateway.services.http_client_service import get_http_client
from mcpgateway.utils.redis_client import get_redis_client as _get_shared_redis_client
from mcpgateway.utils.ssl_context_cache import get_cached_ssl_context

logger = logging.getLogger(__name__)

# In-memory storage for OAuth states with expiration (fallback for single-process)
# Format: {state_key: {"state": state, "gateway_id": gateway_id, "expires_at": datetime}}
_oauth_states: Dict[str, Dict[str, Any]] = {}
# Reverse lookup for callback handlers that only receive state.
# Format: {state: gateway_id}
_oauth_state_lookup: Dict[str, str] = {}
# Lock for thread-safe state operations
_state_lock = asyncio.Lock()

# State TTL in seconds (5 minutes)
STATE_TTL_SECONDS = 300

# Redis client for distributed state storage (uses shared factory)
_redis_client: Optional[Any] = None
_REDIS_INITIALIZED = False


async def _get_redis_client():
    """Get shared Redis client for distributed state storage.

    Uses the centralized Redis client factory for consistent configuration.

    Returns:
        Redis client instance or None if unavailable
    """
    global _redis_client, _REDIS_INITIALIZED  # pylint: disable=global-statement

    if _REDIS_INITIALIZED:
        return _redis_client

    settings = get_settings()
    if settings.cache_type == "redis" and settings.redis_url:
        try:
            _redis_client = await _get_shared_redis_client()
            if _redis_client:
                logger.info("OAuth manager connected to shared Redis client")
        except Exception as e:
            logger.warning(f"Failed to get Redis client, falling back to in-memory storage: {e}")
            _redis_client = None
    else:
        _redis_client = None

    _REDIS_INITIALIZED = True
    return _redis_client


class OAuthManager:
    """Manages OAuth 2.0 authentication flows.

    Examples:
        >>> manager = OAuthManager(request_timeout=30, max_retries=3)
        >>> manager.request_timeout
        30
        >>> manager.max_retries
        3
        >>> manager.token_storage is None
        True
        >>>
        >>> # Test grant type validation
        >>> grant_type = "client_credentials"
        >>> grant_type in ["client_credentials", "authorization_code"]
        True
        >>> grant_type = "invalid_grant"
        >>> grant_type in ["client_credentials", "authorization_code"]
        False
        >>>
        >>> # Test encrypted secret detection heuristic
        >>> short_secret = "secret123"
        >>> len(short_secret) > 50
        False
        >>> encrypted_secret = "gAAAAABh" + "x" * 60  # Simulated encrypted secret
        >>> len(encrypted_secret) > 50
        True
        >>>
        >>> # Test scope list handling
        >>> scopes = ["read", "write"]
        >>> " ".join(scopes)
        'read write'
        >>> empty_scopes = []
        >>> " ".join(empty_scopes)
        ''
    """

    # Known Microsoft Entra login hosts (global + sovereign clouds).
    _ENTRA_HOSTS: frozenset[str] = frozenset(
        {
            "login.microsoftonline.com",
            "login.microsoftonline.us",
            "login.microsoftonline.de",
            "login.partner.microsoftonline.cn",
        }
    )

    def __init__(self, request_timeout: int = 30, max_retries: int = 3, token_storage: Optional[Any] = None):
        """Initialize OAuth Manager.

        Args:
            request_timeout: Timeout for OAuth requests in seconds
            max_retries: Maximum number of retry attempts for token requests
            token_storage: Optional TokenStorageService for storing tokens
        """
        self.request_timeout = request_timeout
        self.max_retries = max_retries
        self.token_storage = token_storage
        self.settings = get_settings()

    async def _get_client(self) -> httpx.AsyncClient:
        """Get the shared singleton HTTP client.

        Returns:
            Shared httpx.AsyncClient instance with connection pooling
        """
        return await get_http_client()

    def _generate_pkce_params(self) -> Dict[str, str]:
        """Generate PKCE parameters for OAuth Authorization Code flow (RFC 7636).

        Returns:
            Dict containing code_verifier, code_challenge, and code_challenge_method
        """
        # Generate code_verifier: 43-128 character random string
        code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).decode("utf-8").rstrip("=")

        # Generate code_challenge: base64url(SHA256(code_verifier))
        code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode("utf-8")).digest()).decode("utf-8").rstrip("=")

        return {"code_verifier": code_verifier, "code_challenge": code_challenge, "code_challenge_method": "S256"}

    async def get_access_token(self, credentials: Dict[str, Any], ca_certificate: Optional[str] = None, client_cert: Optional[str] = None, client_key: Optional[str] = None) -> str:
        """Get access token based on grant type.

        Args:
            credentials: OAuth configuration containing grant_type and other params
            ca_certificate: Optional custom CA certificate for SSL verification (PEM format)
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Access token string

        Raises:
            ValueError: If grant type is unsupported
            OAuthError: If token acquisition fails

        Examples:
            Client credentials flow:
            >>> import asyncio
            >>> class TestMgr(OAuthManager):
            ...     async def _client_credentials_flow(self, credentials, ca_certificate=None, client_cert=None, client_key=None):
            ...         return 'tok'
            >>> mgr = TestMgr()
            >>> asyncio.run(mgr.get_access_token({'grant_type': 'client_credentials'}))
            'tok'

            Authorization code flow requires interactive completion:
            >>> def _auth_code_requires_consent():
            ...     try:
            ...         asyncio.run(mgr.get_access_token({'grant_type': 'authorization_code'}))
            ...     except OAuthError:
            ...         return True
            >>> _auth_code_requires_consent()
            True

            Unsupported grant type raises ValueError:
            >>> def _unsupported():
            ...     try:
            ...         asyncio.run(mgr.get_access_token({'grant_type': 'bad'}))
            ...     except ValueError:
            ...         return True
            >>> _unsupported()
            True
        """
        grant_type = credentials.get("grant_type")
        logger.debug(f"Getting access token for grant type: {grant_type}")

        if grant_type == "client_credentials":
            return await self._client_credentials_flow(credentials, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key)
        if grant_type == "password":
            return await self._password_flow(credentials, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key)
        if grant_type == "authorization_code":
            raise OAuthError("Authorization code flow requires user consent via /oauth/authorize and does not support client_credentials fallback")
        raise ValueError(f"Unsupported grant type: {grant_type}")

    @staticmethod
    async def _prepare_runtime_credentials(credentials: Dict[str, Any], flow_name: str) -> Dict[str, Any]:
        """Return runtime-ready oauth credentials with sensitive fields decrypted.

        Args:
            credentials: Stored oauth_config payload.
            flow_name: Flow label for diagnostic logging.

        Returns:
            Dict[str, Any]: Runtime-ready credentials.
        """
        try:
            settings = get_settings()
            encryption = get_encryption_service(settings.auth_encryption_secret)
            runtime_credentials = await decrypt_oauth_config_for_runtime(credentials, encryption=encryption)
        except Exception as exc:
            logger.warning("Failed to decrypt runtime OAuth credentials for %s flow; falling back to stored values: %s", flow_name, exc)
            return credentials

        if not isinstance(runtime_credentials, dict):
            raise OAuthError(f"Invalid runtime OAuth configuration for {flow_name} flow")

        token_url = runtime_credentials.get("token_url")
        if isinstance(token_url, str) and token_url:
            runtime_credentials["token_url"] = validate_core_url(token_url, "OAuth config token_url")

        auth_server = runtime_credentials.get("authorization_server")
        if isinstance(auth_server, str) and auth_server:
            runtime_credentials["authorization_server"] = validate_core_url(auth_server, "OAuth config authorization_server")

        issuer = runtime_credentials.get("issuer")
        if isinstance(issuer, str) and issuer:
            runtime_credentials["issuer"] = validate_core_url(issuer, "OAuth config issuer")

        auth_servers = runtime_credentials.get("authorization_servers")
        if auth_servers not in (None, ""):
            if not isinstance(auth_servers, list):
                raise OAuthError("OAuth configuration authorization_servers must be a list")
            validated_servers = []
            for idx, server_url in enumerate(auth_servers):
                if not isinstance(server_url, str):
                    raise OAuthError(f"OAuth configuration authorization_servers[{idx}] must be a string URL")
                if server_url:
                    validated_servers.append(validate_core_url(server_url, f"OAuth config authorization_servers[{idx}]"))
            runtime_credentials["authorization_servers"] = validated_servers

        return runtime_credentials

    async def _post_token_request(
        self, url: str, data: Any, ca_certificate: Optional[str] = None, client_cert: Optional[str] = None, client_key: Optional[str] = None, headers: Optional[Dict[str, str]] = None
    ) -> httpx.Response:
        """POST to a token endpoint, using a custom SSL context when CA certs or mTLS are provided.

        When ``ca_certificate``, ``client_cert``, or ``client_key`` is supplied,
        an isolated ``httpx.AsyncClient`` is created with the corresponding SSL
        context so that OAuth token exchange works against self-signed or
        custom-CA upstream servers and/or presents client certificates for mTLS.
        Otherwise the shared HTTP client (which respects the global
        ``SKIP_SSL_VERIFY`` setting) is used.

        Note:
            When only ``client_cert``/``client_key`` are provided (no custom CA),
            the isolated client uses the system's default trust store and does
            not honour ``SKIP_SSL_VERIFY``.

        Args:
            url: Token endpoint URL.
            data: Form-encoded request body (dict or list of tuples for RFC 8707).
            ca_certificate: Optional PEM-encoded CA certificate.
            client_cert: Optional client certificate for mTLS.
            client_key: Optional client private key for mTLS.
            headers: Optional HTTP headers (e.g., Authorization for Basic Auth).

        Returns:
            The HTTP response from the token endpoint.
        """
        if ca_certificate or client_cert or client_key:
            ssl_context = get_cached_ssl_context(ca_certificate, client_cert=client_cert, client_key=client_key)
            async with httpx.AsyncClient(verify=ssl_context) as client:
                return await client.post(url, data=data, headers=headers, timeout=self.request_timeout)
        client = await self._get_client()
        return await client.post(url, data=data, headers=headers, timeout=self.request_timeout)

    # Keys whose values must never be echoed in error messages or logs.
    _SENSITIVE_TOKEN_KEYS = frozenset({"access_token", "refresh_token", "id_token", "client_secret", "password"})

    # Cap on raw_response excerpts and any other string values surfaced via
    # OAuthError / logs (defense-in-depth against unbounded provider bodies).
    _MAX_RAW_RESPONSE_LEN = 256

    # OAuth parameter names per RFC 6749 are token-shaped (alphanumerics plus
    # a few separators). A parsed key outside this shape means parse_qsl picked
    # garbage out of an HTML body (e.g. <meta charset="utf-8">).
    _OAUTH_KEY_RE = re.compile(r"^[A-Za-z0-9_.\-]+$")

    # Scrub URL/form-style ``key=value`` leaks inside arbitrary strings so that
    # secrets embedded in HTML error pages or stack traces don't survive the
    # length cap (the value can fit entirely inside the truncation window).
    _LEAKY_PARAM_RE = re.compile(
        r"(?i)\b(access_token|refresh_token|id_token|token|code|secret|key|password|api[_-]?key)=[^&\s\"'<>]+",
    )

    @staticmethod
    def _safe_response_text(response: httpx.Response) -> str:
        """Return ``response.text`` or a placeholder if the body is undecodable.

        ``httpx.Response.text`` raises ``UnicodeDecodeError`` (or ``LookupError``
        for an unknown charset) when the body bytes don't match the declared
        encoding. The caller wants a string for diagnostics, not a crash.

        Args:
            response: HTTP response whose body we want as text.

        Returns:
            Decoded body, or a ``"<undecodable body, N bytes>"`` placeholder.
        """
        try:
            return response.text
        except (ValueError, LookupError):
            return f"<undecodable body, {len(response.content)} bytes>"

    @staticmethod
    def _parse_token_response(response: httpx.Response) -> Dict[str, Any]:
        """Parse an OAuth token response that may be JSON or form-encoded.

        Per RFC 7231 §3.1.1.1, media type tokens are case-insensitive.
        Form-encoded values are URL-decoded via ``urllib.parse.parse_qsl``.
        Failures fall back to ``{"raw_response": <text>}`` so operators see
        what the provider actually sent, in three cases: a JSON parse error
        (``ValueError`` covering ``json.JSONDecodeError`` and
        ``UnicodeDecodeError``), ``parse_qsl`` returning ``{}`` from a
        non-empty body (e.g. an HTML error page served with a form-encoded
        content-type), and ``response.text`` failing to decode the body
        bytes.

        Args:
            response: HTTP response from the token endpoint.

        Returns:
            Parsed token payload, or ``{"raw_response": <text>}`` when the
            body is neither valid JSON nor parseable as form-encoded.
        """
        raw_content_type = response.headers.get("content-type", "")
        content_type = raw_content_type.lower()

        if "application/x-www-form-urlencoded" in content_type:
            text = OAuthManager._safe_response_text(response)
            # parse_qsl drops malformed pairs (no "="); we deliberately do not
            # set keep_blank_values=True so that garbage like an HTML error page
            # parses to {} and falls through to the raw_response capture below.
            parsed = dict(parse_qsl(text))
            # An HTML body that happens to contain "=" (e.g. <meta charset="utf-8">)
            # parses to a non-empty dict with garbage keys. Reject anything whose
            # keys aren't OAuth parameter shaped so the leak is bounded by the
            # raw_response truncation in _redact_token_response.
            if parsed and not all(OAuthManager._OAUTH_KEY_RE.match(k) for k in parsed):
                return {"raw_response": text}
            if not parsed and text:
                return {"raw_response": text}
            return parsed

        try:
            return response.json()
        except ValueError as exc:
            # ValueError covers json.JSONDecodeError (malformed JSON) and
            # UnicodeDecodeError (bad charset). Narrower than bare Exception,
            # which would swallow httpx.ResponseNotRead, MemoryError, etc.
            text = OAuthManager._safe_response_text(response)
            logger.warning(
                "Failed to parse OAuth token response as JSON: %s (status=%s, content-type=%r, body_bytes=%d)",
                exc,
                response.status_code,
                raw_content_type,
                len(response.content),
                exc_info=True,
            )
            return {"raw_response": text}

    @staticmethod
    def _redact_token_response(token_response: Dict[str, Any]) -> Dict[str, Any]:
        """Return a log/error-safe copy of a token response.

        Three layers of protection so that misbehaving providers, HTML error
        pages, and verbose stack traces don't leak secrets via OAuthError or
        log lines:

        1. Replace values for known credential-bearing keys with
           ``"[REDACTED]"``.
        2. Scrub URL/form-style ``<key>=<secret>`` patterns inside any string
           value (HTML hrefs, form actions, stack traces).
        3. Cap any string value at ``_MAX_RAW_RESPONSE_LEN`` chars with a
           ``... [truncated, N chars total]`` marker.

        Args:
            token_response: Parsed token payload (possibly containing tokens
                or a captured raw body).

        Returns:
            New dict safe to interpolate into log lines and exception messages.
        """
        cap = OAuthManager._MAX_RAW_RESPONSE_LEN
        redacted: Dict[str, Any] = {}
        for key, value in token_response.items():
            if key in OAuthManager._SENSITIVE_TOKEN_KEYS:
                redacted[key] = "[REDACTED]"
                continue
            if isinstance(value, str):
                # Scrub URL/form-style "<key>=<secret>" patterns first (HTML
                # bodies often carry tokens in href / form action attributes
                # that fit entirely inside the truncation window), then cap.
                scrubbed = OAuthManager._LEAKY_PARAM_RE.sub(lambda m: f"{m.group(1)}=[REDACTED]", value)
                if len(scrubbed) > cap:
                    redacted[key] = f"{scrubbed[:cap]}... [truncated, {len(value)} chars total]"
                else:
                    redacted[key] = scrubbed
            else:
                redacted[key] = value
        return redacted

    @staticmethod
    def _build_basic_auth_header(client_id: str, client_secret: str) -> str:
        """Build an RFC 6749 Section 2.3.1 Basic Auth header value.

        Client credentials are first URL-encoded per RFC 6749 Appendix B,
        then combined as ``client_id:client_secret`` and base64-encoded.
        """
        credentials_str = f"{quote(client_id, safe='')}:{quote(client_secret, safe='')}"
        encoded_credentials = base64.b64encode(credentials_str.encode("utf-8")).decode("utf-8")
        return f"Basic {encoded_credentials}"

    async def _client_credentials_flow(
        self,
        credentials: Dict[str, Any],
        ca_certificate: Optional[str] = None,
        client_cert: Optional[str] = None,
        client_key: Optional[str] = None,
    ) -> str:
        """Machine-to-machine authentication using client credentials.

        Args:
            credentials: OAuth configuration with client_id, client_secret, token_url
            ca_certificate: Optional custom CA certificate for SSL verification (PEM format)
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Access token string

        Raises:
            OAuthError: If token acquisition fails after all retries
        """
        runtime_credentials = await self._prepare_runtime_credentials(credentials, "client_credentials")
        client_id = runtime_credentials["client_id"]
        client_secret = runtime_credentials["client_secret"]
        token_url = runtime_credentials["token_url"]
        if not isinstance(token_url, str):
            raise OAuthError("OAuth configuration missing valid token_url")
        scopes = runtime_credentials.get("scopes", [])

        # Check if provider requires Basic Auth for client authentication (RFC 6749 Section 2.3.1)
        # Default to form-based auth for backward compatibility
        use_basic_auth = runtime_credentials.get("token_endpoint_auth_method", "client_secret_post") == "client_secret_basic"

        # Prepare token request data and headers
        token_data = {"grant_type": "client_credentials"}
        headers = {}

        if use_basic_auth:
            headers["Authorization"] = self._build_basic_auth_header(client_id, client_secret)
            logger.debug("Using HTTP Basic Auth for token endpoint authentication")
        else:
            # Default: client credentials in POST body (client_secret_post)
            token_data["client_id"] = client_id
            token_data["client_secret"] = client_secret
            logger.debug("Using POST body for token endpoint authentication")

        if scopes:
            token_data["scope"] = " ".join(scopes) if isinstance(scopes, list) else scopes

        # Fetch token with retries
        for attempt in range(self.max_retries):
            try:
                response = await self._post_token_request(token_url, token_data, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key, headers=headers)
                response.raise_for_status()

                token_response = self._parse_token_response(response)

                if "access_token" not in token_response:
                    raise OAuthError("OAuth token endpoint response did not contain access_token")

                logger.info("""Successfully obtained access token via client credentials""")
                return token_response["access_token"]

            except httpx.HTTPError as e:
                logger.warning(f"Token request attempt {attempt + 1} failed: {str(e)}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Failed to obtain access token after {self.max_retries} attempts: {str(e)}")
                await asyncio.sleep(2**attempt)  # Exponential backoff

        # This should never be reached due to the exception above, but needed for type safety
        raise OAuthError("Failed to obtain access token after all retry attempts")

    async def _password_flow(self, credentials: Dict[str, Any], ca_certificate: Optional[str] = None, client_cert: Optional[str] = None, client_key: Optional[str] = None) -> str:
        """Resource Owner Password Credentials flow (RFC 6749 Section 4.3).

        This flow is used when the application can directly handle the user's credentials,
        such as with trusted first-party applications or legacy integrations like Keycloak.

        Args:
            credentials: OAuth configuration with client_id, optional client_secret, token_url, username, password
            ca_certificate: Optional custom CA certificate for SSL verification (PEM format)
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Access token string

        Raises:
            OAuthError: If token acquisition fails after all retries
        """
        runtime_credentials = await self._prepare_runtime_credentials(credentials, "password")
        client_id = runtime_credentials.get("client_id")
        client_secret = runtime_credentials.get("client_secret")
        token_url = runtime_credentials["token_url"]
        username = runtime_credentials.get("username")
        password = runtime_credentials.get("password")
        scopes = runtime_credentials.get("scopes", [])

        if not username or not password:
            raise OAuthError("Username and password are required for password grant type")

        # Prepare token request data
        token_data = {
            "grant_type": "password",
            "username": username,
            "password": password,
        }

        # Add client_id (required by most providers including Keycloak)
        if client_id:
            token_data["client_id"] = client_id

        # Add client_secret if present (some providers require it, others don't)
        if client_secret:
            token_data["client_secret"] = client_secret

        if scopes:
            token_data["scope"] = " ".join(scopes) if isinstance(scopes, list) else scopes

        # Fetch token with retries
        for attempt in range(self.max_retries):
            try:
                response = await self._post_token_request(token_url, token_data, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key)
                response.raise_for_status()

                token_response = self._parse_token_response(response)

                if "access_token" not in token_response:
                    raise OAuthError(f"No access_token in response: {self._redact_token_response(token_response)}")

                logger.info("Successfully obtained access token via password grant")
                return token_response["access_token"]

            except httpx.HTTPError as e:
                logger.warning(f"Token request attempt {attempt + 1} failed: {str(e)}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Failed to obtain access token after {self.max_retries} attempts: {str(e)}")
                await asyncio.sleep(2**attempt)  # Exponential backoff

        # This should never be reached due to the exception above, but needed for type safety
        raise OAuthError("Failed to obtain access token after all retry attempts")

    async def get_authorization_url(self, credentials: Dict[str, Any]) -> Dict[str, str]:
        """Get authorization URL for user delegation flow.

        Args:
            credentials: OAuth configuration with client_id, authorization_url, etc.

        Returns:
            Dict containing authorization_url and state
        """
        client_id = credentials["client_id"]
        redirect_uri = credentials["redirect_uri"]
        authorization_url = credentials["authorization_url"]
        scopes = credentials.get("scopes", [])

        # Create OAuth2 session
        oauth = OAuth2Session(client_id, redirect_uri=redirect_uri, scope=scopes)

        # Generate authorization URL with state for CSRF protection
        auth_url, state = oauth.authorization_url(authorization_url)

        logger.info(f"Generated authorization URL for client {client_id}")

        return {"authorization_url": auth_url, "state": state}

    async def exchange_code_for_token(self, credentials: Dict[str, Any], code: str, state: str) -> str:  # pylint: disable=unused-argument
        """Exchange authorization code for access token.

        Args:
            credentials: OAuth configuration
            code: Authorization code from callback
            state: State parameter for CSRF validation

        Returns:
            Access token string

        Raises:
            OAuthError: If token exchange fails
        """
        runtime_credentials = await self._prepare_runtime_credentials(credentials, "authorization_code_exchange")
        client_id = runtime_credentials["client_id"]
        client_secret = runtime_credentials.get("client_secret")  # Optional for public clients (PKCE-only)
        token_url = runtime_credentials["token_url"]
        redirect_uri = runtime_credentials["redirect_uri"]

        # Prepare token exchange data
        token_data = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
            "client_id": client_id,
        }

        # Only include client_secret if present (public clients don't have secrets)
        if client_secret:
            token_data["client_secret"] = client_secret

        # Exchange code for token with retries
        for attempt in range(self.max_retries):
            try:
                client = await self._get_client()
                response = await client.post(token_url, data=token_data, timeout=self.request_timeout)
                response.raise_for_status()

                token_response = self._parse_token_response(response)

                if "access_token" not in token_response:
                    raise OAuthError(f"No access_token in response: {self._redact_token_response(token_response)}")

                logger.info("""Successfully exchanged authorization code for access token""")
                return token_response["access_token"]

            except httpx.HTTPError as e:
                logger.warning(f"Token exchange attempt {attempt + 1} failed: {str(e)}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Failed to exchange code for token after {self.max_retries} attempts: {str(e)}")
                await asyncio.sleep(2**attempt)  # Exponential backoff

        # This should never be reached due to the exception above, but needed for type safety
        raise OAuthError("Failed to exchange code for token after all retry attempts")

    async def token_exchange(
        self,
        token_url: str,
        subject_token: str,
        client_id: str,
        client_secret: str,
        audience: Optional[str] = None,
        scope: Optional[str] = None,
        requested_token_type: str = "urn:ietf:params:oauth:token-type:access_token",
    ) -> Dict[str, Any]:
        """RFC 8693 token exchange for on-behalf-of flows.

        Exchanges a subject token (e.g. the gateway's JWT) for a new token
        scoped to a downstream service, enabling the downstream to act on
        behalf of the original user.

        Args:
            token_url: Token endpoint of the authorization server.
            subject_token: The original user's access token.
            client_id: Client ID for the gateway.
            client_secret: Client secret for the gateway.
            audience: Intended audience for the exchanged token.
            scope: Requested scope for the exchanged token.
            requested_token_type: The type of token being requested.

        Returns:
            Dict with ``access_token``, ``token_type``, and optionally
            ``expires_in`` and ``scope``.

        Raises:
            OAuthError: If token exchange fails after all retries.
        """
        # Decrypt client secret if encrypted
        if client_secret:
            try:
                settings = get_settings()
                encryption = get_encryption_service(settings.auth_encryption_secret)
                if encryption.is_encrypted(client_secret):
                    decrypted = await encryption.decrypt_secret_async(client_secret)
                    if decrypted:
                        client_secret = decrypted
            except Exception as e:
                logger.warning(f"Failed to decrypt client secret for token exchange: {e}")

        token_data: Dict[str, str] = {
            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
            "subject_token": subject_token,
            "subject_token_type": "urn:ietf:params:oauth:token-type:access_token",  # nosec
            "requested_token_type": requested_token_type,
            "client_id": client_id,
            "client_secret": client_secret,
        }
        if audience:
            token_data["audience"] = audience
        if scope:
            token_data["scope"] = scope

        for attempt in range(self.max_retries):
            try:
                client = await self._get_client()
                response = await client.post(token_url, data=token_data, timeout=self.request_timeout)
                response.raise_for_status()

                token_response = response.json()
                if "access_token" not in token_response:
                    raise OAuthError(f"No access_token in token exchange response: {token_response}")

                logger.info("Successfully performed RFC 8693 token exchange")
                return token_response

            except httpx.HTTPError as e:
                logger.warning(f"Token exchange attempt {attempt + 1} failed: {e}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Token exchange failed after {self.max_retries} attempts: {e}")
                await asyncio.sleep(2**attempt)

        raise OAuthError("Token exchange failed after all retry attempts")

    async def initiate_authorization_code_flow(self, gateway_id: str, credentials: Dict[str, Any], app_user_email: str = None) -> Dict[str, str]:
        """Initiate Authorization Code flow with PKCE and return authorization URL.

        Args:
            gateway_id: ID of the gateway being configured
            credentials: OAuth configuration with client_id, authorization_url, etc.
            app_user_email: ContextForge user email to associate with tokens

        Returns:
            Dict containing authorization_url and state
        """

        # Generate PKCE parameters (RFC 7636)
        pkce_params = self._generate_pkce_params()

        # Generate state parameter with user context for CSRF protection
        state = self._generate_state(gateway_id, app_user_email)

        # Store state with code_verifier in session/cache for validation
        if self.token_storage:
            await self._store_authorization_state(
                gateway_id,
                state,
                code_verifier=pkce_params["code_verifier"],
                app_user_email=app_user_email,
            )

        # Generate authorization URL with PKCE
        auth_url = self._create_authorization_url_with_pkce(credentials, state, pkce_params["code_challenge"], pkce_params["code_challenge_method"])

        logger.info(f"Generated authorization URL with PKCE for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")

        return {"authorization_url": auth_url, "state": state, "gateway_id": gateway_id}

    async def complete_authorization_code_flow(
        self, gateway_id: str, code: str, state: str, credentials: Dict[str, Any], ca_certificate: Optional[str] = None, client_cert: Optional[str] = None, client_key: Optional[str] = None
    ) -> Dict[str, Any]:
        """Complete Authorization Code flow with PKCE and store tokens.

        Args:
            gateway_id: ID of the gateway
            code: Authorization code from callback
            state: State parameter for CSRF validation
            credentials: OAuth configuration
            ca_certificate: Optional custom CA certificate for SSL verification (PEM format)
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Dict containing success status, user_id, and expiration info

        Raises:
            OAuthError: If state validation fails or token exchange fails
        """
        # Validate state and retrieve code_verifier
        state_data = await self._validate_and_retrieve_state(gateway_id, state)
        if not state_data:
            raise OAuthError("Invalid or expired state parameter - possible replay attack")

        code_verifier = state_data.get("code_verifier")
        app_user_email = state_data.get("app_user_email")

        # Defence-in-depth: if app_user_email is absent from server-side
        # state (e.g. state stored by an older code path), attempt a
        # gateway-mismatch check via the legacy state parser but NEVER
        # extract identity fields from unsigned payloads (CWE-345).
        # Note: the /oauth/callback router rejects pure legacy states
        # before reaching here (allow_legacy_fallback=False), so this
        # block only fires for server-stored states that lack the email.
        if not app_user_email:
            legacy_state_payload = self._extract_legacy_state_payload(state)
            if legacy_state_payload:
                legacy_gateway_id = legacy_state_payload.get("gateway_id")
                if legacy_gateway_id and legacy_gateway_id != gateway_id:
                    raise OAuthError("State parameter gateway mismatch")
            if self.token_storage:
                logger.error("User context (app_user_email) missing from OAuth state; refusing to bind tokens (CWE-287). gateway_id=%s", gateway_id)
                raise OAuthError("User context required for OAuth token storage")
            logger.warning("User context (app_user_email) missing from OAuth state; no token_storage configured — proceeding without binding. gateway_id=%s", gateway_id)

        # Exchange code for tokens with PKCE code_verifier
        token_response = await self._exchange_code_for_tokens(credentials, code, code_verifier=code_verifier, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key)

        # Extract user information from token response
        user_id = self._extract_user_id(token_response, credentials)

        # Extract audience from token (best-effort) for caller to persist as resource.
        # This enables audience learning for IdPs that map resource to a different aud.
        token_aud = self._extract_token_audience(token_response.get("access_token", ""))

        # Store tokens if storage service is available
        if self.token_storage:
            # Handle scope as either string or list (OAuth providers vary)
            scope_value = token_response.get("scope", "")
            if isinstance(scope_value, list):
                scopes_list = [s for s in scope_value if isinstance(s, str)]
            elif isinstance(scope_value, str):
                scopes_list = scope_value.split() if scope_value else []
            else:
                scopes_list = []

            token_record = await self.token_storage.store_tokens(
                gateway_id=gateway_id,
                user_id=user_id,
                app_user_email=app_user_email,  # User from state
                access_token=token_response["access_token"],
                refresh_token=token_response.get("refresh_token"),
                expires_in=parse_expires_in(token_response),
                scopes=scopes_list,
            )

            return {"success": True, "user_id": user_id, "expires_at": token_record.expires_at.isoformat() if token_record.expires_at else None, "token_aud": token_aud}
        return {"success": True, "user_id": user_id, "expires_at": None, "token_aud": token_aud}

    async def get_access_token_for_user(self, gateway_id: str, app_user_email: str) -> Optional[str]:
        """Get valid access token for a specific user.

        Args:
            gateway_id: ID of the gateway
            app_user_email: ContextForge user email

        Returns:
            Valid access token or None if not available
        """
        if self.token_storage:
            return await self.token_storage.get_user_token(gateway_id, app_user_email)
        return None

    def _generate_state(self, _gateway_id: str, _app_user_email: str = None) -> str:
        """Generate an opaque state token for CSRF protection.

        Args:
            _gateway_id: Gateway identifier (reserved for compatibility with
                prior embedded-state call sites).
            _app_user_email: ContextForge user email (reserved for
                compatibility with prior embedded-state call sites).

        Returns:
            Opaque random state token
        """
        return secrets.token_urlsafe(48)

    @staticmethod
    def _extract_legacy_state_payload(state: str) -> Optional[Dict[str, Any]]:
        """Best-effort decode of legacy state payloads used before opaque states.

        Legacy formats supported:
        - base64url(payload || signature) where payload is JSON
        - gateway_id_random suffix format

        Security: Legacy payloads lack signature verification, so only
        ``gateway_id`` is returned — never identity-sensitive fields like
        ``app_user_email`` which could be forged (CWE-345).

        Args:
            state: Callback state token to decode.

        Returns:
            Dict containing only ``gateway_id`` when format is recognized;
            otherwise ``None``.
        """
        safe_legacy_fields = {"gateway_id"}

        try:
            state_raw = base64.urlsafe_b64decode(state.encode())
            if len(state_raw) <= 32:
                return None

            payload_bytes = state_raw[:-32]
            payload = orjson.loads(payload_bytes)
            if isinstance(payload, dict):
                # Only return gateway_id — unsigned payloads must not
                # carry identity claims.
                safe = {k: v for k, v in payload.items() if k in safe_legacy_fields}
                return safe if safe else None
        except Exception:
            # Fall back to legacy gateway_id_random format
            if "_" in state:
                gateway_id = state.split("_", 1)[0]
                if gateway_id:
                    return {"gateway_id": gateway_id}
        return None

    async def resolve_gateway_id_from_state(self, state: str, allow_legacy_fallback: bool = True) -> Optional[str]:
        """Resolve gateway ID for a callback state token without consuming it.

        Args:
            state: OAuth callback state parameter
            allow_legacy_fallback: Whether to decode legacy callback state formats.

        Returns:
            Gateway ID when resolvable, otherwise ``None``.
        """
        settings = get_settings()

        if settings.cache_type == "redis":
            redis = await _get_redis_client()
            if redis:
                try:
                    lookup_key = f"oauth:state_lookup:{state}"
                    gateway_id = await redis.get(lookup_key)
                    if gateway_id:
                        if isinstance(gateway_id, bytes):
                            gateway_id = gateway_id.decode("utf-8")
                        return gateway_id
                except Exception as e:
                    logger.warning(f"Failed to resolve state gateway in Redis: {e}")

        if settings.cache_type == "database":
            try:
                # First-Party
                from mcpgateway.db import get_db, OAuthState  # pylint: disable=import-outside-toplevel

                db_gen = get_db()
                db = next(db_gen)
                try:
                    oauth_state = db.query(OAuthState).filter(OAuthState.state == state).first()
                    if oauth_state:
                        return oauth_state.gateway_id
                finally:
                    db_gen.close()
            except Exception as e:
                logger.warning(f"Failed to resolve state gateway in database: {e}")

        async with _state_lock:
            now = datetime.now(timezone.utc)
            expired_keys = [key for key, data in _oauth_states.items() if datetime.fromisoformat(data["expires_at"]) < now]
            for key in expired_keys:
                expired_state = _oauth_states[key].get("state")
                del _oauth_states[key]
                if expired_state:
                    _oauth_state_lookup.pop(expired_state, None)
            gateway_id = _oauth_state_lookup.get(state)
            if gateway_id:
                return gateway_id

        if allow_legacy_fallback:
            legacy_payload = self._extract_legacy_state_payload(state)
            if legacy_payload:
                return legacy_payload.get("gateway_id")
        return None

    async def _store_authorization_state(
        self,
        gateway_id: str,
        state: str,
        code_verifier: str = None,
        app_user_email: str = None,
    ) -> None:
        """Store authorization state for validation with TTL.

        Args:
            gateway_id: ID of the gateway
            state: State parameter to store
            code_verifier: Optional PKCE code verifier (RFC 7636)
            app_user_email: Requesting user email for token association
        """
        expires_at = datetime.now(timezone.utc) + timedelta(seconds=STATE_TTL_SECONDS)
        settings = get_settings()

        # Try Redis first for distributed storage
        if settings.cache_type == "redis":
            redis = await _get_redis_client()
            if redis:
                try:
                    state_key = f"oauth:state:{gateway_id}:{state}"
                    lookup_key = f"oauth:state_lookup:{state}"
                    state_data = {
                        "state": state,
                        "gateway_id": gateway_id,
                        "code_verifier": code_verifier,
                        "app_user_email": app_user_email,
                        "expires_at": expires_at.isoformat(),
                        "used": False,
                    }
                    # Store in Redis with TTL
                    await redis.setex(state_key, STATE_TTL_SECONDS, orjson.dumps(state_data))
                    await redis.setex(lookup_key, STATE_TTL_SECONDS, gateway_id)
                    logger.debug(f"Stored OAuth state in Redis for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                    return
                except Exception as e:
                    logger.warning(f"Failed to store state in Redis: {e}, falling back")

        # Try database storage for multi-worker deployments
        if settings.cache_type == "database":
            try:
                # First-Party
                from mcpgateway.db import get_db, OAuthState  # pylint: disable=import-outside-toplevel

                db_gen = get_db()
                db = next(db_gen)
                try:
                    # Clean up expired states first
                    db.query(OAuthState).filter(OAuthState.expires_at < datetime.now(timezone.utc)).delete()

                    # Store new state with code_verifier
                    oauth_state_kwargs = {
                        "gateway_id": gateway_id,
                        "state": state,
                        "code_verifier": code_verifier,
                        "expires_at": expires_at,
                        "used": False,
                    }
                    if hasattr(OAuthState, "app_user_email"):
                        oauth_state_kwargs["app_user_email"] = app_user_email

                    oauth_state = OAuthState(**oauth_state_kwargs)
                    db.add(oauth_state)
                    db.commit()
                    logger.debug(f"Stored OAuth state in database for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                    return
                finally:
                    db_gen.close()
            except Exception as e:
                logger.warning(f"Failed to store state in database: {e}, falling back to memory")

        # Fallback to in-memory storage for development
        async with _state_lock:
            # Clean up expired states first
            now = datetime.now(timezone.utc)
            state_key = f"oauth:state:{gateway_id}:{state}"
            state_data = {
                "state": state,
                "gateway_id": gateway_id,
                "code_verifier": code_verifier,
                "app_user_email": app_user_email,
                "expires_at": expires_at.isoformat(),
                "used": False,
            }
            expired_states = [key for key, data in _oauth_states.items() if datetime.fromisoformat(data["expires_at"]) < now]
            for key in expired_states:
                expired_state_value = _oauth_states[key].get("state")
                del _oauth_states[key]
                if expired_state_value:
                    _oauth_state_lookup.pop(expired_state_value, None)
                logger.debug(f"Cleaned up expired state: {key[:20]}...")

            # Store the new state with expiration
            _oauth_states[state_key] = state_data
            _oauth_state_lookup[state] = gateway_id
            logger.debug(f"Stored OAuth state in memory for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")

    async def _validate_authorization_state(self, gateway_id: str, state: str) -> bool:
        """Validate authorization state parameter and mark as used.

        Args:
            gateway_id: ID of the gateway
            state: State parameter to validate

        Returns:
            True if state is valid and not yet used, False otherwise
        """
        settings = get_settings()

        # Try Redis first for distributed storage
        if settings.cache_type == "redis":
            redis = await _get_redis_client()
            if redis:
                try:
                    state_key = f"oauth:state:{gateway_id}:{state}"
                    lookup_key = f"oauth:state_lookup:{state}"
                    # Get and delete state atomically (single-use)
                    state_json = await redis.getdel(state_key)
                    await redis.delete(lookup_key)
                    if not state_json:
                        logger.warning(f"State not found in Redis for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                        return False

                    state_data = orjson.loads(state_json)

                    # Parse expires_at as timezone-aware datetime. If the stored value
                    # is naive, assume UTC for compatibility.
                    try:
                        expires_at = datetime.fromisoformat(state_data["expires_at"])
                    except Exception:
                        # Fallback: try parsing without microseconds/offsets
                        expires_at = datetime.strptime(state_data["expires_at"], "%Y-%m-%dT%H:%M:%S")

                    if expires_at.tzinfo is None:
                        # Assume UTC for naive timestamps
                        expires_at = expires_at.replace(tzinfo=timezone.utc)

                    # Check if state has expired
                    if expires_at < datetime.now(timezone.utc):
                        logger.warning(f"State has expired for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                        return False

                    # Check if state was already used (should not happen with getdel)
                    if state_data.get("used", False):
                        logger.warning(f"State was already used for gateway {SecurityValidator.sanitize_log_message(gateway_id)} - possible replay attack")
                        return False

                    logger.debug(f"Successfully validated OAuth state from Redis for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                    return True
                except Exception as e:
                    logger.warning(f"Failed to validate state in Redis: {e}, falling back")

        # Try database storage for multi-worker deployments
        if settings.cache_type == "database":
            try:
                # First-Party
                from mcpgateway.db import get_db, OAuthState  # pylint: disable=import-outside-toplevel

                db_gen = get_db()
                db = next(db_gen)
                try:
                    # Find the state
                    oauth_state = db.query(OAuthState).filter(OAuthState.gateway_id == gateway_id, OAuthState.state == state).first()

                    if not oauth_state:
                        logger.warning(f"State not found in database for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                        return False

                    # Check if state has expired
                    # Ensure oauth_state.expires_at is timezone-aware. If naive, assume UTC.
                    expires_at = oauth_state.expires_at
                    if expires_at.tzinfo is None:
                        expires_at = expires_at.replace(tzinfo=timezone.utc)

                    if expires_at < datetime.now(timezone.utc):
                        logger.warning(f"State has expired for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                        db.delete(oauth_state)
                        db.commit()
                        return False

                    # Check if state was already used
                    if oauth_state.used:
                        logger.warning(f"State has already been used for gateway {SecurityValidator.sanitize_log_message(gateway_id)} - possible replay attack")
                        return False

                    # Mark as used and delete (single-use)
                    db.delete(oauth_state)
                    db.commit()
                    logger.debug(f"Successfully validated OAuth state from database for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                    return True
                finally:
                    db_gen.close()
            except Exception as e:
                logger.warning(f"Failed to validate state in database: {e}, falling back to memory")

        # Fallback to in-memory storage for development
        state_key = f"oauth:state:{gateway_id}:{state}"
        async with _state_lock:
            state_data = _oauth_states.get(state_key)

            # Check if state exists
            if not state_data:
                logger.warning(f"State not found in memory for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                return False

            # Parse and normalize expires_at to timezone-aware datetime
            expires_at = datetime.fromisoformat(state_data["expires_at"])
            if expires_at.tzinfo is None:
                expires_at = expires_at.replace(tzinfo=timezone.utc)

            if expires_at < datetime.now(timezone.utc):
                logger.warning(f"State has expired for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
                del _oauth_states[state_key]  # Clean up expired state
                _oauth_state_lookup.pop(state, None)
                return False

            # Check if state has already been used (prevent replay)
            if state_data.get("used", False):
                logger.warning(f"State has already been used for gateway {SecurityValidator.sanitize_log_message(gateway_id)} - possible replay attack")
                return False

            # Mark state as used and remove it (single-use)
            del _oauth_states[state_key]
            _oauth_state_lookup.pop(state, None)
            logger.debug(f"Successfully validated OAuth state from memory for gateway {SecurityValidator.sanitize_log_message(gateway_id)}")
            return True

    async def _validate_and_retrieve_state(self, gateway_id: str, state: str) -> Optional[Dict[str, Any]]:
        """Validate state and return full state data including code_verifier.

        Args:
            gateway_id: ID of the gateway
            state: State parameter to validate

        Returns:
            Dict with state data including code_verifier, or None if invalid/expired
        """
        settings = get_settings()

        # Try Redis first
        if settings.cache_type == "redis":
            redis = await _get_redis_client()
            if redis:
                try:
                    state_key = f"oauth:state:{gateway_id}:{state}"
                    lookup_key = f"oauth:state_lookup:{state}"
                    state_json = await redis.getdel(state_key)  # Atomic get+delete
                    await redis.delete(lookup_key)
                    if not state_json:
                        return None

                    state_data = orjson.loads(state_json)

                    # Check expiration
                    try:
                        expires_at = datetime.fromisoformat(state_data["expires_at"])
                    except Exception:
                        expires_at = datetime.strptime(state_data["expires_at"], "%Y-%m-%dT%H:%M:%S")

                    if expires_at.tzinfo is None:
                        expires_at = expires_at.replace(tzinfo=timezone.utc)

                    if expires_at < datetime.now(timezone.utc):
                        return None

                    return state_data
                except Exception as e:
                    logger.warning(f"Failed to validate state in Redis: {e}, falling back")

        # Try database
        if settings.cache_type == "database":
            try:
                # First-Party
                from mcpgateway.db import get_db, OAuthState  # pylint: disable=import-outside-toplevel

                db_gen = get_db()
                db = next(db_gen)
                try:
                    oauth_state = db.query(OAuthState).filter(OAuthState.gateway_id == gateway_id, OAuthState.state == state).first()

                    if not oauth_state:
                        return None

                    # Check expiration
                    expires_at = oauth_state.expires_at
                    if expires_at.tzinfo is None:
                        expires_at = expires_at.replace(tzinfo=timezone.utc)

                    if expires_at < datetime.now(timezone.utc):
                        db.delete(oauth_state)
                        db.commit()
                        return None

                    # Check if already used
                    if oauth_state.used:
                        return None

                    # Build state data
                    state_data = {
                        "state": oauth_state.state,
                        "gateway_id": oauth_state.gateway_id,
                        "code_verifier": oauth_state.code_verifier,
                        "expires_at": oauth_state.expires_at.isoformat(),
                    }
                    if hasattr(oauth_state, "app_user_email"):
                        state_data["app_user_email"] = getattr(oauth_state, "app_user_email", None)

                    # Mark as used and delete
                    db.delete(oauth_state)
                    db.commit()

                    return state_data
                finally:
                    db_gen.close()
            except Exception as e:
                logger.warning(f"Failed to validate state in database: {e}")

        # Fallback to in-memory
        state_key = f"oauth:state:{gateway_id}:{state}"
        async with _state_lock:
            state_data = _oauth_states.get(state_key)
            if not state_data:
                return None

            # Check expiration
            expires_at = datetime.fromisoformat(state_data["expires_at"])
            if expires_at.tzinfo is None:
                expires_at = expires_at.replace(tzinfo=timezone.utc)

            if expires_at < datetime.now(timezone.utc):
                del _oauth_states[state_key]
                _oauth_state_lookup.pop(state, None)
                return None

            # Remove from memory (single-use)
            del _oauth_states[state_key]
            _oauth_state_lookup.pop(state, None)
            return state_data

    def _create_authorization_url(self, credentials: Dict[str, Any], state: str) -> tuple[str, str]:
        """Create authorization URL with state parameter.

        Args:
            credentials: OAuth configuration
            state: State parameter for CSRF protection

        Returns:
            Tuple of (authorization_url, state)
        """
        client_id = credentials["client_id"]
        redirect_uri = credentials["redirect_uri"]
        authorization_url = credentials["authorization_url"]
        scopes = credentials.get("scopes", [])

        # Create OAuth2 session
        oauth = OAuth2Session(client_id, redirect_uri=redirect_uri, scope=scopes)

        # Generate authorization URL with state for CSRF protection
        auth_url, state = oauth.authorization_url(authorization_url, state=state)

        return auth_url, state

    @staticmethod
    def _is_microsoft_entra_v2_endpoint(endpoint_url: Any) -> bool:
        """Return True when endpoint matches Microsoft Entra v2 login endpoints.

        Args:
            endpoint_url: OAuth endpoint URL to check

        Returns:
            True if the endpoint is a Microsoft Entra v2 OAuth endpoint
        """
        if not isinstance(endpoint_url, str) or not endpoint_url:
            return False

        parsed = urlparse(endpoint_url)
        host = (parsed.hostname or "").lower()
        path = parsed.path.lower()

        return host in OAuthManager._ENTRA_HOSTS and "/oauth2/v2.0/" in path

    @staticmethod
    def _is_enabled_flag(value: Any) -> bool:
        """Parse boolean-like config values from oauth_config.

        Args:
            value: Config value to interpret as boolean

        Returns:
            True if value represents an enabled/truthy setting
        """
        if isinstance(value, bool):
            return value
        if isinstance(value, str):
            return value.strip().lower() in {"1", "true", "yes", "on"}
        return False

    def _should_include_resource_parameter(self, credentials: Dict[str, Any], scopes: Any) -> bool:
        """Determine whether RFC 8707 resource should be sent for this request.

        Args:
            credentials: OAuth configuration containing resource and endpoint URLs
            scopes: OAuth scopes for the request

        Returns:
            True if the resource parameter should be included in the request
        """
        if not credentials.get("resource"):
            return False

        if self._is_enabled_flag(credentials.get("omit_resource")):
            return False

        # Microsoft Entra v2 does not accept legacy resource with v2 scope-based requests.
        if scopes and (self._is_microsoft_entra_v2_endpoint(credentials.get("authorization_url")) or self._is_microsoft_entra_v2_endpoint(credentials.get("token_url"))):
            logger.info("Omitting OAuth resource parameter for Microsoft Entra v2 scope-based flow")
            return False

        return True

    def _create_authorization_url_with_pkce(self, credentials: Dict[str, Any], state: str, code_challenge: str, code_challenge_method: str) -> str:
        """Create authorization URL with PKCE parameters (RFC 7636).

        Args:
            credentials: OAuth configuration
            state: State parameter for CSRF protection
            code_challenge: PKCE code challenge
            code_challenge_method: PKCE method (S256)

        Returns:
            Authorization URL string with PKCE parameters
        """
        # Standard
        from urllib.parse import urlencode  # pylint: disable=import-outside-toplevel

        client_id = credentials["client_id"]
        redirect_uri = credentials["redirect_uri"]
        authorization_url = credentials["authorization_url"]
        scopes = credentials.get("scopes", [])

        # Build authorization parameters
        params = {"response_type": "code", "client_id": client_id, "redirect_uri": redirect_uri, "state": state, "code_challenge": code_challenge, "code_challenge_method": code_challenge_method}

        # Add scopes if present
        if scopes:
            params["scope"] = " ".join(scopes) if isinstance(scopes, list) else scopes

        # Add resource parameter for JWT access token (RFC 8707)
        # The resource is the MCP server URL, set by oauth_router.py
        resource = credentials.get("resource")
        if self._should_include_resource_parameter(credentials, scopes):
            params["resource"] = resource  # urlencode with doseq=True handles lists

        # Build full URL (doseq=True handles list values like multiple resource params)
        query_string = urlencode(params, doseq=True)
        return f"{authorization_url}?{query_string}"

    async def _exchange_code_for_tokens(
        self, credentials: Dict[str, Any], code: str, code_verifier: str = None, ca_certificate: Optional[str] = None, client_cert: Optional[str] = None, client_key: Optional[str] = None
    ) -> Dict[str, Any]:
        """Exchange authorization code for tokens with PKCE support.

        Args:
            credentials: OAuth configuration
            code: Authorization code from callback
            code_verifier: Optional PKCE code verifier (RFC 7636)
            ca_certificate: Optional custom CA certificate for SSL verification (PEM format).
                When provided along with ``client_cert``/``client_key``, creates an
                isolated HTTP client with custom SSL context instead of using the
                shared client. This enables OAuth token exchange with self-signed or
                custom CA upstream OAuth servers and/or mTLS authentication.
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Token response dictionary

        Raises:
            OAuthError: If token exchange fails
        """
        runtime_credentials = await self._prepare_runtime_credentials(credentials, "authorization_code_exchange_with_pkce")
        client_id = runtime_credentials["client_id"]
        client_secret = runtime_credentials.get("client_secret")  # Optional for public clients (PKCE-only)
        token_url = runtime_credentials["token_url"]
        redirect_uri = runtime_credentials["redirect_uri"]

        # Check if provider requires Basic Auth for client authentication (RFC 6749 Section 2.3.1)
        # Default to form-based auth for backward compatibility
        use_basic_auth = runtime_credentials.get("token_endpoint_auth_method", "client_secret_post") == "client_secret_basic"

        # Prepare token exchange data and headers
        token_data = {
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
        }
        headers = {}

        if use_basic_auth and client_secret:
            headers["Authorization"] = self._build_basic_auth_header(client_id, client_secret)
            logger.debug("Using HTTP Basic Auth for token endpoint authentication")
        elif use_basic_auth and not client_secret:
            # Public PKCE clients can't use Basic Auth (no secret to encode)
            logger.warning("Basic Auth requested but client_secret is missing - falling back to POST body mode (public client)")
            token_data["client_id"] = client_id
            logger.debug("Using POST body for token endpoint authentication")
        else:
            # Default: client credentials in POST body (client_secret_post)
            token_data["client_id"] = client_id
            # Only include client_secret if present (public clients don't have secrets)
            if client_secret:
                token_data["client_secret"] = client_secret
            logger.debug("Using POST body for token endpoint authentication")

        # Add PKCE code_verifier if present (RFC 7636)
        if code_verifier:
            token_data["code_verifier"] = code_verifier

        # Add resource parameter to request JWT access token (RFC 8707)
        # The resource identifies the MCP server (resource server), not the OAuth server
        resource = runtime_credentials.get("resource")
        scopes = runtime_credentials.get("scopes", [])
        if self._should_include_resource_parameter(credentials, scopes):
            if isinstance(resource, list):
                # RFC 8707 allows multiple resource parameters - use list of tuples
                form_data: list[tuple[str, str]] = list(token_data.items())
                for r in resource:
                    if r:
                        form_data.append(("resource", r))
                token_data = form_data  # type: ignore[assignment]
            else:
                token_data["resource"] = resource

        # Exchange code for token with retries
        for attempt in range(self.max_retries):
            try:
                response = await self._post_token_request(token_url, token_data, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key, headers=headers)
                response.raise_for_status()

                token_response = self._parse_token_response(response)

                if "access_token" not in token_response:
                    raise OAuthError(f"No access_token in response: {self._redact_token_response(token_response)}")

                logger.info("""Successfully exchanged authorization code for tokens""")
                return token_response

            except httpx.HTTPError as e:
                logger.warning(f"Token exchange attempt {attempt + 1} failed: {str(e)}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Failed to exchange code for token after {self.max_retries} attempts: {str(e)}")
                await asyncio.sleep(2**attempt)  # Exponential backoff

        # This should never be reached due to the exception above, but needed for type safety
        raise OAuthError("Failed to exchange code for token after all retry attempts")

    async def refresh_token(
        self,
        refresh_token: str,
        credentials: Dict[str, Any],
        ca_certificate: Optional[str] = None,
        client_cert: Optional[str] = None,
        client_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Refresh an expired access token using a refresh token.

        Args:
            refresh_token: The refresh token to use
            credentials: OAuth configuration including client_id, client_secret, token_url
            ca_certificate: Optional custom CA certificate (PEM format or file path) to use
                instead of system trust store. When provided, creates an isolated HTTP client
                instead of using the shared client. This enables OAuth token refresh
                with self-signed or custom CA upstream OAuth servers.
            client_cert: Optional client certificate for mTLS (PEM format or file path)
            client_key: Optional client private key for mTLS (PEM format or file path)

        Returns:
            Dict containing new access_token, optional refresh_token, and expires_in

        Raises:
            OAuthError: If token refresh fails
        """
        if not refresh_token:
            raise OAuthError("No refresh token available")

        runtime_credentials = await self._prepare_runtime_credentials(credentials, "refresh_token")
        token_url = runtime_credentials.get("token_url")
        if not token_url:
            raise OAuthError("No token URL configured for OAuth provider")

        client_id = runtime_credentials.get("client_id")
        client_secret = runtime_credentials.get("client_secret")

        if not client_id:
            raise OAuthError("No client_id configured for OAuth provider")

        # Check if provider requires Basic Auth for client authentication (RFC 6749 Section 2.3.1)
        # Default to form-based auth for backward compatibility
        use_basic_auth = runtime_credentials.get("token_endpoint_auth_method", "client_secret_post") == "client_secret_basic"

        # Prepare token refresh request and headers
        token_data = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
        }
        headers = {}

        if use_basic_auth and client_secret:
            headers["Authorization"] = self._build_basic_auth_header(client_id, client_secret)
            logger.debug("Using HTTP Basic Auth for token endpoint authentication")
        elif use_basic_auth and not client_secret:
            # Misconfiguration: Basic Auth requested but no secret available
            logger.warning("Basic Auth requested but client_secret is missing - falling back to POST body mode")
            token_data["client_id"] = client_id
            logger.debug("Using POST body for token endpoint authentication")
        else:
            # Default: client credentials in POST body (client_secret_post)
            token_data["client_id"] = client_id
            # Add client_secret if available (some providers require it)
            if client_secret:
                token_data["client_secret"] = client_secret
            logger.debug("Using POST body for token endpoint authentication")

        # Add resource parameter for JWT access token (RFC 8707)
        # Must be included in refresh requests to maintain JWT token type
        resource = runtime_credentials.get("resource")
        scopes = runtime_credentials.get("scopes", [])
        if self._should_include_resource_parameter(credentials, scopes):
            if isinstance(resource, list):
                # RFC 8707 allows multiple resource parameters - use list of tuples
                form_data: list[tuple[str, str]] = list(token_data.items())
                for r in resource:
                    if r:
                        form_data.append(("resource", r))
                token_data = form_data  # type: ignore[assignment]
            else:
                token_data["resource"] = resource

        # Attempt token refresh with retries
        for attempt in range(self.max_retries):
            try:
                response = await self._post_token_request(token_url, token_data, ca_certificate=ca_certificate, client_cert=client_cert, client_key=client_key, headers=headers)
                if response.status_code == 200:
                    token_response = self._parse_token_response(response)

                    # Validate required fields
                    if "access_token" not in token_response:
                        raise OAuthError(f"No access_token in refresh response: {self._redact_token_response(token_response)}")

                    logger.info("Successfully refreshed OAuth token")
                    return token_response

                # Bound and redact the body before surfacing it. Some providers echo
                # request parameters (including refresh_token / client_secret) in error
                # responses, and HTML error pages can be unbounded — both leak via logs
                # and OAuthError messages without this scrub.
                error_payload = self._redact_token_response(self._parse_token_response(response))
                if response.status_code in [400, 401]:
                    raise OAuthError(f"Refresh token invalid or expired: {error_payload}")
                logger.warning("Token refresh failed with status %s: %s", response.status_code, error_payload)

            except httpx.HTTPError as e:
                logger.warning(f"Token refresh attempt {attempt + 1} failed: {str(e)}")
                if attempt == self.max_retries - 1:
                    raise OAuthError(f"Failed to refresh token after {self.max_retries} attempts: {str(e)}")
                await asyncio.sleep(2**attempt)  # Exponential backoff

        raise OAuthError("Failed to refresh token after all retry attempts")

    def _extract_user_id(self, token_response: Dict[str, Any], credentials: Dict[str, Any]) -> str:
        """Extract user ID from token response.

        Args:
            token_response: Response from token exchange
            credentials: OAuth configuration

        Returns:
            User ID string
        """
        # Try to extract user ID from various common fields in token response
        # Different OAuth providers use different field names

        # Check for 'sub' (subject) - JWT standard
        if "sub" in token_response:
            return token_response["sub"]

        # Check for 'user_id' - common in some OAuth responses
        if "user_id" in token_response:
            return token_response["user_id"]

        # Check for 'id' - also common
        if "id" in token_response:
            return token_response["id"]

        # Fallback to client_id if no user info is available
        if credentials.get("client_id"):
            return credentials["client_id"]

        # Final fallback
        return "unknown_user"

    @staticmethod
    def _extract_token_audience(access_token: str) -> Any:
        """Extract the ``aud`` claim from a JWT access token (best-effort).

        Returns the raw ``aud`` value (string or list) or ``None`` for opaque
        tokens or decode failures.  No signature verification is performed.

        Args:
            access_token: The raw access token string.

        Returns:
            The ``aud`` claim value, or None.
        """
        if not access_token:
            return None
        try:
            # Third-Party
            import jwt as pyjwt  # pylint: disable=import-outside-toplevel

            claims = pyjwt.decode(
                access_token,
                options={"verify_signature": False, "verify_aud": False, "verify_iss": False, "verify_exp": False},
                algorithms=["RS256", "RS384", "RS512", "ES256", "ES384", "ES512", "PS256", "PS384", "PS512", "HS256", "HS384", "HS512", "EdDSA"],
            )
            return claims.get("aud")
        except Exception:  # noqa: BLE001
            return None


class OAuthError(Exception):
    """OAuth-related errors.

    Examples:
        >>> try:
        ...     raise OAuthError("Token acquisition failed")
        ... except OAuthError as e:
        ...     str(e)
        'Token acquisition failed'
        >>> try:
        ...     raise OAuthError("Invalid grant type")
        ... except Exception as e:
        ...     isinstance(e, OAuthError)
        True
    """


class OAuthRequiredError(OAuthError):
    """Raised when a server requires OAuth but the caller is unauthenticated.

    Carries ``server_id`` so the middleware can identify which server
    triggered the rejection when constructing the ``WWW-Authenticate``
    header.

    Examples:
        >>> err = OAuthRequiredError("auth required", server_id="s1")
        >>> err.server_id
        's1'
        >>> isinstance(err, OAuthError)
        True
    """

    def __init__(self, message: str, *, server_id: str = "") -> None:
        """Initialize with message and optional server_id.

        Args:
            message: Human-readable error description.
            server_id: Virtual-server identifier that triggered the rejection.
        """
        super().__init__(message)
        self.server_id = server_id


class OAuthEnforcementUnavailableError(OAuthError):
    """Raised when OAuth enforcement cannot be performed due to infrastructure failure.

    Used when the database or other backing services needed to check a
    server's ``oauth_enabled`` flag are unavailable.  The middleware
    translates this into an HTTP 503 to avoid silently allowing
    unauthenticated access (fail-closed).

    Examples:
        >>> err = OAuthEnforcementUnavailableError("DB down", server_id="s1")
        >>> err.server_id
        's1'
        >>> isinstance(err, OAuthError)
        True
    """

    def __init__(self, message: str, *, server_id: str = "") -> None:
        """Initialize with message and optional server_id.

        Args:
            message: Human-readable error description.
            server_id: Virtual-server identifier that triggered the rejection.
        """
        super().__init__(message)
        self.server_id = server_id


def parse_expires_in(token_response: Dict[str, Any]) -> Optional[int]:
    """Parse and validate the ``expires_in`` field from an OAuth token response.

    RFC 6749 §5.1 marks ``expires_in`` as RECOMMENDED (not REQUIRED). When the
    field is absent or null, the gateway records ``expires_at`` as ``None`` and
    the token is treated as having no known local expiry, subject to the
    stale-token cleanup policy in
    :meth:`mcpgateway.services.token_storage_service.TokenStorageService.cleanup_expired_tokens`.

    Args:
        token_response: Raw OAuth token response dict from the provider.

    Returns:
        ``int`` lifetime in seconds when the provider supplied a non-negative
        integer (or numeric string convertible to one), or ``None`` when the
        field is absent or explicitly null.

    Raises:
        OAuthError: If ``expires_in`` is present but malformed (negative,
            non-numeric, or a non-scalar type).

    Examples:
        >>> parse_expires_in({"expires_in": 3600})
        3600
        >>> parse_expires_in({"expires_in": "3600"})
        3600
        >>> parse_expires_in({"expires_in": 0})
        0
        >>> parse_expires_in({}) is None
        True
        >>> parse_expires_in({"expires_in": None}) is None
        True
        >>> try:
        ...     parse_expires_in({"expires_in": -1})
        ... except OAuthError as exc:
        ...     "negative" in str(exc)
        True
        >>> try:
        ...     parse_expires_in({"expires_in": "garbage"})
        ... except OAuthError as exc:
        ...     "Invalid expires_in" in str(exc)
        True
    """
    raw = token_response.get("expires_in")
    if raw is None:
        return None
    # Reject bools (True/False are int subclasses in Python) and any non-scalar types.
    if isinstance(raw, bool) or not isinstance(raw, (int, float, str)):
        raise OAuthError(f"Invalid expires_in from OAuth provider: {raw!r}")
    # Sign-check the original numeric BEFORE int() truncation, otherwise int(-0.5) == 0
    # would bypass the negative check.
    if isinstance(raw, (int, float)) and raw < 0:
        raise OAuthError(f"Invalid expires_in from OAuth provider (negative): {raw}")
    # RFC 6749 §5.1 specifies integer seconds; reject non-integral floats explicitly.
    if isinstance(raw, float) and not raw.is_integer():
        raise OAuthError(f"Invalid expires_in from OAuth provider (non-integer): {raw}")
    try:
        value = int(raw)
    except (TypeError, ValueError) as exc:
        raise OAuthError(f"Invalid expires_in from OAuth provider: {raw!r}") from exc
    if value < 0:
        raise OAuthError(f"Invalid expires_in from OAuth provider (negative): {value}")
    return value
