# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Sequence
import logging
import threading
from typing import Optional

from fastapi.openapi.models import OAuth2

from ..agents.callback_context import CallbackContext
from ..tools.openapi_tool.auth.credential_exchangers.service_account_exchanger import ServiceAccountCredentialExchanger
from ..utils.feature_decorator import experimental
from .auth_credential import AuthCredential
from .auth_credential import AuthCredentialTypes
from .auth_provider_registry import AuthProviderRegistry
from .auth_schemes import AuthSchemeType
from .auth_schemes import CustomAuthScheme
from .auth_schemes import ExtendedOAuth2
from .auth_schemes import OpenIdConnectWithConfig
from .auth_tool import AuthConfig
from .base_auth_provider import BaseAuthProvider
from .exchanger.base_credential_exchanger import BaseCredentialExchanger
from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
from .oauth2_discovery import OAuth2DiscoveryManager
from .refresher.credential_refresher_registry import CredentialRefresherRegistry

logger = logging.getLogger("google_adk." + __name__)


def _rehydrate_custom_scheme(
    scheme: CustomAuthScheme,
    supported_schemes: Sequence[type[CustomAuthScheme]],
) -> CustomAuthScheme:
  """Rehydrate a CustomAuthScheme into one of the given supported_schemes."""
  incoming_type = scheme.type_
  for scheme_class in supported_schemes:
    type_field = scheme_class.model_fields.get("type_")
    # Custom AuthScheme classes must define a `default` for their `type_` field
    # to be rehydrated correctly.
    if type_field and type_field.default == incoming_type:
      data = scheme.model_dump(by_alias=True)
      if scheme.model_extra:
        data.update(scheme.model_extra)
      return scheme_class.model_validate(data)
  raise ValueError(
      f"Cannot rehydrate: no registered scheme matches type '{incoming_type}'"
  )


@experimental
class CredentialManager:
  """Manages authentication credentials through a structured workflow.

  The CredentialManager orchestrates the complete lifecycle of authentication
  credentials, from initial loading to final preparation for use. It provides
  a centralized interface for handling various credential types and authentication
  schemes while maintaining proper credential hygiene (refresh, exchange, caching).

  This class is only for use by Agent Development Kit.

  Args:
      auth_config: Configuration containing authentication scheme and credentials

  Example:
      ```python
      auth_config = AuthConfig(
          auth_scheme=oauth2_scheme,
          raw_auth_credential=service_account_credential
      )
      manager = CredentialManager(auth_config)

      # Register custom exchanger if needed
      manager.register_credential_exchanger(
          AuthCredentialTypes.CUSTOM_TYPE,
          CustomCredentialExchanger()
      )

      # Register custom refresher if needed
      manager.register_credential_refresher(
          AuthCredentialTypes.CUSTOM_TYPE,
          CustomCredentialRefresher()
      )

      # Load and prepare credential
      credential = await manager.load_auth_credential(tool_context)
      ```
  """

  _auth_provider_registry = AuthProviderRegistry()
  _registry_lock = threading.Lock()

  @classmethod
  def register_auth_provider(cls, provider: BaseAuthProvider) -> None:
    """Public API for developers to register custom auth providers."""
    with cls._registry_lock:
      for scheme_type in provider.supported_auth_schemes:
        existing_provider = cls._auth_provider_registry.get_provider(
            scheme_type
        )
        if existing_provider is not None:
          if existing_provider is not provider:
            logger.warning(
                "An auth provider is already registered for scheme %s. "
                "Ignoring the new provider.",
                scheme_type,
            )
            continue
        cls._auth_provider_registry.register(scheme_type, provider)

  def __init__(
      self,
      auth_config: AuthConfig,
  ):
    self._auth_config = auth_config
    self._exchanger_registry = CredentialExchangerRegistry()
    self._refresher_registry = CredentialRefresherRegistry()
    self._discovery_manager = OAuth2DiscoveryManager()

    # Register default exchangers and refreshers
    from .exchanger.oauth2_credential_exchanger import OAuth2CredentialExchanger
    from .refresher.oauth2_credential_refresher import OAuth2CredentialRefresher

    oauth2_exchanger = OAuth2CredentialExchanger()
    self._exchanger_registry.register(
        AuthCredentialTypes.OAUTH2, oauth2_exchanger
    )
    self._exchanger_registry.register(
        AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_exchanger
    )

    # TODO: Move ServiceAccountCredentialExchanger to the auth module
    self._exchanger_registry.register(
        AuthCredentialTypes.SERVICE_ACCOUNT,
        ServiceAccountCredentialExchanger(),
    )

    oauth2_refresher = OAuth2CredentialRefresher()
    self._refresher_registry.register(
        AuthCredentialTypes.OAUTH2, oauth2_refresher
    )
    self._refresher_registry.register(
        AuthCredentialTypes.OPEN_ID_CONNECT, oauth2_refresher
    )

  def register_credential_exchanger(
      self,
      credential_type: AuthCredentialTypes,
      exchanger_instance: BaseCredentialExchanger,
  ) -> None:
    """Register a credential exchanger for a credential type.

    Args:
        credential_type: The credential type to register for.
        exchanger_instance: The exchanger instance to register.
    """
    self._exchanger_registry.register(credential_type, exchanger_instance)

  async def request_credential(self, context: CallbackContext) -> None:
    if not hasattr(context, "request_credential"):
      raise TypeError(
          "request_credential requires a ToolContext with request_credential"
          " method, not a plain CallbackContext"
      )
    context.request_credential(self._auth_config)

  async def get_auth_credential(
      self, context: CallbackContext
  ) -> Optional[AuthCredential]:
    """Load and prepare authentication credential through a structured workflow."""

    # Step 0: Handle CustomAuthScheme if present
    if isinstance(self._auth_config.auth_scheme, CustomAuthScheme):
      # Pydantic may have deserialized an unknown scheme into a generic
      # CustomAuthScheme. If so, rehydrate it first into a specific subclass.
      # Note: Custom authentication scheme classes must have been imported into
      # the Python runtime before get_auth_credential is called for their
      # subclasses to be registered. This is fine as developer will anyway
      # import them while registering the auth providers.
      # Note: `__subclasses__()` only returns immediate subclasses, if there is
      # a subclass of a subclass of CustomAuthScheme then it will not be
      # returned.
      # pylint: disable=unidiomatic-typecheck Needs exact class matching.
      if type(self._auth_config.auth_scheme) is CustomAuthScheme:
        self._auth_config.auth_scheme = _rehydrate_custom_scheme(
            self._auth_config.auth_scheme,
            CustomAuthScheme.__subclasses__(),
        )

      provider = self._auth_provider_registry.get_provider(
          self._auth_config.auth_scheme
      )
      if provider is None:
        raise ValueError(
            "No auth provider registered for custom auth scheme "
            f"{self._auth_config.auth_scheme.type_!r}. "
            "Register it using `CredentialManager.register_auth_provider("
            "<YourAuthProviderInstance>)`."
        )
      provided_credential = await provider.get_auth_credential(
          self._auth_config, context
      )
      if not provided_credential:
        raise ValueError("AuthProvider did not return a credential.")
      # Handle special case for OAuth2 user consent flow.
      if (
          provided_credential.oauth2
          and not provided_credential.oauth2.access_token
          and provided_credential.oauth2.auth_uri
      ):
        # User consent is required. We save the auth uri and return None
        # to signal the need for user consent.
        self._auth_config.exchanged_auth_credential = provided_credential
        return None
      return provided_credential

    # Step 1: Validate credential configuration
    await self._validate_credential()

    # Step 2: Check if credential is already ready (no processing needed)
    raw_auth_credential = self._auth_config.raw_auth_credential
    if self._is_credential_ready() and raw_auth_credential is not None:
      # Return a copy to avoid leaking mutations across invocations/users when
      # tools share a long-lived AuthConfig instance.
      return raw_auth_credential.model_copy(deep=True)

    # Step 3: Try to load existing processed credential
    credential = await self._load_existing_credential(context)

    # Step 4: If no existing credential, load from auth response
    # TODO instead of load from auth response, we can store auth response in
    # credential service.
    was_from_auth_response = False
    if not credential:
      credential = await self._load_from_auth_response(context)
      was_from_auth_response = True

    # Step 5: If still no credential available, check if client credentials
    if not credential:
      # For client credentials flow, use raw credentials directly
      if self._is_client_credentials_flow():
        # Exchange/refresh steps may mutate the credential object in-place, so
        # do not operate on the shared tool config.
        credential = self._auth_config.raw_auth_credential.model_copy(deep=True)
      else:
        # For authorization code flow, return None to trigger user authorization
        return None

    # Step 6: Exchange credential if needed (e.g., service account to access token)
    credential, was_exchanged = await self._exchange_credential(credential)

    # Step 7: Refresh credential if expired
    was_refreshed = False
    if not was_exchanged:
      credential, was_refreshed = await self._refresh_credential(credential)

    # Step 8: Save credential if it was modified
    if was_from_auth_response or was_exchanged or was_refreshed:
      await self._save_credential(context, credential)

    return credential

  async def _load_existing_credential(
      self, context: CallbackContext
  ) -> Optional[AuthCredential]:
    """Load existing credential from credential service."""

    # Try loading from credential service first
    credential = await self._load_from_credential_service(context)
    if credential:
      return credential

    return None

  async def _load_from_credential_service(
      self, context: CallbackContext
  ) -> Optional[AuthCredential]:
    """Load credential from credential service if available."""
    credential_service = context._invocation_context.credential_service
    if credential_service:
      # Note: This should be made async in a future refactor
      # For now, assuming synchronous operation
      return await context.load_credential(self._auth_config)
    return None

  async def _load_from_auth_response(
      self, context: CallbackContext
  ) -> Optional[AuthCredential]:
    """Load credential from auth response in context."""
    return context.get_auth_response(self._auth_config)

  async def _exchange_credential(
      self, credential: AuthCredential
  ) -> tuple[AuthCredential, bool]:
    """Exchange credential if needed and return the credential and whether it was exchanged."""
    exchanger = self._exchanger_registry.get_exchanger(credential.auth_type)
    if not exchanger:
      return credential, False

    if isinstance(exchanger, ServiceAccountCredentialExchanger):
      return (
          exchanger.exchange_credential(
              self._auth_config.auth_scheme, credential
          ),
          True,
      )

    exchange_result = await exchanger.exchange(
        credential, self._auth_config.auth_scheme
    )
    return exchange_result.credential, exchange_result.was_exchanged

  async def _refresh_credential(
      self, credential: AuthCredential
  ) -> tuple[AuthCredential, bool]:
    """Refresh credential if expired and return the credential and whether it was refreshed."""
    refresher = self._refresher_registry.get_refresher(credential.auth_type)
    if not refresher:
      return credential, False

    if await refresher.is_refresh_needed(
        credential, self._auth_config.auth_scheme
    ):
      refreshed_credential = await refresher.refresh(
          credential, self._auth_config.auth_scheme
      )
      return refreshed_credential, True

    return credential, False

  def _is_credential_ready(self) -> bool:
    """Check if credential is ready to use without further processing."""
    raw_credential = self._auth_config.raw_auth_credential
    if not raw_credential:
      return False

    # Simple credentials that don't need exchange or refresh
    return raw_credential.auth_type in (
        AuthCredentialTypes.API_KEY,
        AuthCredentialTypes.HTTP,
        # Add other simple auth types as needed
    )

  async def _validate_credential(self) -> None:
    """Validate credential configuration and raise errors if invalid."""
    if not self._auth_config.raw_auth_credential:
      if self._auth_config.auth_scheme.type_ in (
          AuthSchemeType.oauth2,
          AuthSchemeType.openIdConnect,
      ):
        raise ValueError(
            "raw_auth_credential is required for auth_scheme type "
            f"{self._auth_config.auth_scheme.type_}"
        )

    raw_credential = self._auth_config.raw_auth_credential
    if raw_credential:
      if (
          raw_credential.auth_type
          in (
              AuthCredentialTypes.OAUTH2,
              AuthCredentialTypes.OPEN_ID_CONNECT,
          )
          and not raw_credential.oauth2
      ):
        raise ValueError(
            "auth_config.raw_credential.oauth2 required for credential type "
            f"{raw_credential.auth_type}"
        )

    if self._missing_oauth_info() and not await self._populate_auth_scheme():
      raise ValueError(
          "OAuth scheme info is missing, and auto-discovery has failed to fill"
          " them in."
      )

    # Additional validation can be added here

  async def _save_credential(
      self, context: CallbackContext, credential: AuthCredential
  ) -> None:
    """Save credential to credential service if available."""
    credential_service = context._invocation_context.credential_service
    if credential_service:
      auth_config_to_save = self._auth_config.model_copy(deep=True)
      auth_config_to_save.exchanged_auth_credential = credential
      await context.save_credential(auth_config_to_save)

  async def _populate_auth_scheme(self) -> bool:
    """Auto-discover server metadata and populate missing auth scheme info.

    Returns:
      True if auto-discovery was successful, False otherwise.
    """
    auth_scheme = self._auth_config.auth_scheme
    if (
        not isinstance(auth_scheme, ExtendedOAuth2)
        or not auth_scheme.issuer_url
    ):
      logger.warning("No issuer_url was provided for auto-discovery.")
      return False

    metadata = await self._discovery_manager.discover_auth_server_metadata(
        auth_scheme.issuer_url
    )
    if not metadata:
      logger.warning("Auto-discovery has failed to populate OAuth scheme info.")
      return False

    flows = auth_scheme.flows

    if flows.implicit and not flows.implicit.authorizationUrl:
      flows.implicit.authorizationUrl = metadata.authorization_endpoint
    if flows.password and not flows.password.tokenUrl:
      flows.password.tokenUrl = metadata.token_endpoint
    if flows.clientCredentials and not flows.clientCredentials.tokenUrl:
      flows.clientCredentials.tokenUrl = metadata.token_endpoint
    if flows.authorizationCode and not flows.authorizationCode.authorizationUrl:
      flows.authorizationCode.authorizationUrl = metadata.authorization_endpoint
    if flows.authorizationCode and not flows.authorizationCode.tokenUrl:
      flows.authorizationCode.tokenUrl = metadata.token_endpoint
    return True

  def _missing_oauth_info(self) -> bool:
    """Checks if we are missing auth/token URLs needed for OAuth."""
    auth_scheme = self._auth_config.auth_scheme
    if isinstance(auth_scheme, OAuth2):
      flows = auth_scheme.flows
      return bool(
          flows.implicit
          and not flows.implicit.authorizationUrl
          or flows.password
          and not flows.password.tokenUrl
          or flows.clientCredentials
          and not flows.clientCredentials.tokenUrl
          or flows.authorizationCode
          and not flows.authorizationCode.authorizationUrl
          or flows.authorizationCode
          and not flows.authorizationCode.tokenUrl
      )
    return False

  def _is_client_credentials_flow(self) -> bool:
    """Check if the auth scheme uses client credentials flow.

    Supports both OAuth2 and OIDC schemes.

    Returns:
      True if using client credentials flow, False otherwise.
    """
    auth_scheme = self._auth_config.auth_scheme

    # Check OAuth2 schemes
    if isinstance(auth_scheme, OAuth2) and auth_scheme.flows:
      return auth_scheme.flows.clientCredentials is not None

    # Check OIDC schemes
    if isinstance(auth_scheme, OpenIdConnectWithConfig):
      return (
          auth_scheme.grant_types_supported is not None
          and "client_credentials" in auth_scheme.grant_types_supported
      )

    return False
