# 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

import asyncio
import logging
import os
import time

from google.adk.agents.callback_context import CallbackContext
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import HttpAuth
from google.adk.auth.auth_credential import HttpCredentials
from google.adk.auth.auth_credential import OAuth2Auth
from google.adk.flows.llm_flows.functions import REQUEST_EUC_FUNCTION_CALL_NAME
from google.api_core.client_options import ClientOptions

try:
  from google.cloud.iamconnectorcredentials_v1alpha import IAMConnectorCredentialsServiceClient as Client
  from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsMetadata
  from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsRequest
  from google.cloud.iamconnectorcredentials_v1alpha import RetrieveCredentialsResponse
except ImportError as e:
  raise ImportError(
      "Missing required dependencies for Agent Identity Auth Manager. "
      'Please install with: pip install "google-adk[agent-identity]"'
  ) from e
from google.longrunning.operations_pb2 import Operation

from .gcp_auth_provider_scheme import GcpAuthProviderScheme

# Notes on the current IAM Connector Credentials service implementation:
# 1. The service does not yet support LROs, so even though the
#    retrieve_credentials method returns an Operation object, the methods like
#    operation.done() and operation.result() will not work yet.
# 2. For API key flows, the returned Operation contains the credentials.
# 3. For 2-legged OAuth flows, the returned Operation contains pending status,
#    client needs to retry the request until response with credentials is
#    returned or timeout occurs.
# 4. For 3-legged OAuth flows, the returned Operation contains consent pending
#    status along with the authorization URI.

# TODO: Catch specific exceptions instead of generic ones.

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

NON_INTERACTIVE_TOKEN_POLL_INTERVAL_SEC: float = 1.0
NON_INTERACTIVE_TOKEN_POLL_TIMEOUT_SEC: float = 10.0


def _construct_auth_credential(
    response: RetrieveCredentialsResponse,
) -> AuthCredential:
  """Constructs a simplified HTTP auth credential from the header-token tuple returned by the upstream service."""
  if not response.header or not response.token:
    raise ValueError(
        "Received either empty header or token from IAM Connector Credentials"
        " service."
    )

  header_name, _, header_value = response.header.partition(":")
  if (
      header_name.strip().lower() == "authorization"
      and header_value.strip().lower().startswith("bearer")
  ):
    return AuthCredential(
        auth_type=AuthCredentialTypes.HTTP,
        http=HttpAuth(
            scheme="Bearer",
            credentials=HttpCredentials(token=response.token),
        ),
    )

  # Handle custom header.
  return AuthCredential(
      auth_type=AuthCredentialTypes.HTTP,
      http=HttpAuth(
          # For custom headers, scheme and credentials fields are not used.
          scheme="",
          credentials=HttpCredentials(),
          additional_headers={
              response.header: response.token,
              "X-GOOG-API-KEY": response.token,
          },
      ),
  )


class _IamConnectorCredentialsProvider:
  """Implementation for auth provider using IAM Connector credentials service."""

  _client: Client | None = None

  def __init__(self, client: Client | None = None):
    self._client = client

  def _get_client(self) -> Client:
    """Lazy loads the client to avoid unnecessary setup on startup."""
    if self._client is None:
      client_options = None
      if host := os.environ.get("IAM_CONNECTOR_CREDENTIALS_TARGET_HOST"):
        client_options = ClientOptions(api_endpoint=host)
      self._client = Client(client_options=client_options, transport="rest")
    return self._client

  async def _retrieve_credentials(
      self,
      user_id: str,
      auth_scheme: GcpAuthProviderScheme,
  ) -> Operation:
    request = RetrieveCredentialsRequest(
        connector=auth_scheme.name,
        user_id=user_id,
        scopes=auth_scheme.scopes,
        continue_uri=auth_scheme.continue_uri or "",
        force_refresh=False,
    )
    # TODO: Use async client once available. Temporarily using threading to
    # prevent blocking the event loop.
    operation = await asyncio.to_thread(
        self._get_client().retrieve_credentials, request
    )
    return operation.operation

  def _unpack_operation(
      self, operation: Operation
  ) -> tuple[
      RetrieveCredentialsResponse | None, RetrieveCredentialsMetadata | None
  ]:
    """Deserializes the response and metadata from the operation."""
    response = None
    metadata = None
    if operation.response:
      response = RetrieveCredentialsResponse.deserialize(
          operation.response.value
      )
    if operation.metadata:
      metadata = RetrieveCredentialsMetadata.deserialize(
          operation.metadata.value
      )
    return response, metadata

  async def _poll_credentials(
      self, user_id: str, auth_scheme: GcpAuthProviderScheme, timeout: float
  ) -> Operation:
    end_time = time.time() + timeout
    while time.time() < end_time:
      operation = await self._retrieve_credentials(user_id, auth_scheme)
      if operation.done:
        return operation
      await asyncio.sleep(NON_INTERACTIVE_TOKEN_POLL_INTERVAL_SEC)
    raise TimeoutError("Timeout waiting for credentials.")

  @staticmethod
  def _is_consent_completed(context: CallbackContext) -> bool:
    """Checks if the user consent flow is completed for the current function call."""
    if not context.function_call_id:
      return False

    if not context.session:
      return False

    events = context.session.events
    target_tool_call_id = context.function_call_id

    # Find all relevant function calls and responses
    euc_calls = {}
    euc_responses = {}

    for event in events:
      for call in event.get_function_calls():
        if call.name == REQUEST_EUC_FUNCTION_CALL_NAME:
          euc_calls[call.id] = call
      for response in event.get_function_responses():
        if response.name == REQUEST_EUC_FUNCTION_CALL_NAME:
          euc_responses[response.id] = response

    # Check for a response that matches a call for the current tool invocation
    for call_id, _ in euc_responses.items():
      if call_id in euc_calls:
        call = euc_calls[call_id]
        if call.args and call.args.get("functionCallId") == target_tool_call_id:
          return True
    return False

  async def get_auth_credential(
      self,
      auth_scheme: GcpAuthProviderScheme,
      context: CallbackContext | None = None,
  ) -> AuthCredential:
    """Retrieves credentials using the IAM Connector Credentials service.

    Args:
      auth_scheme: The GcpAuthProviderScheme.
      context: Optional context for the callback.

    Returns:
      An AuthCredential instance.

    Raises:
      RuntimeError: If credential retrieval or polling fails.
    """

    if context is None or context.user_id is None:
      raise ValueError(
          "GcpAuthProvider requires a context with a valid user_id."
      )

    user_id = context.user_id

    try:
      operation = await self._retrieve_credentials(user_id, auth_scheme)
    except Exception as e:
      raise RuntimeError(
          f"Failed to retrieve credential for user '{user_id}' on connector"
          f" '{auth_scheme.name}'."
      ) from e

    response, metadata = self._unpack_operation(operation)

    if operation.HasField("error"):
      raise RuntimeError(f"Operation failed: {operation.error.message}")

    if operation.done:
      logger.debug("Auth credential obtained immediately.")
      return _construct_auth_credential(response)

    if metadata is not None and "consent_pending" in metadata:
      # Get 2-legged OAuth token. Allow enough time for token exchange.
      try:
        operation = await self._poll_credentials(
            user_id,
            auth_scheme,
            timeout=NON_INTERACTIVE_TOKEN_POLL_TIMEOUT_SEC,
        )
        if operation.HasField("error"):
          raise RuntimeError(f"Operation failed: {operation.error.message}")
        if operation.done:
          logger.debug("Auth credential obtained after polling.")
          response, _ = self._unpack_operation(operation)
          return _construct_auth_credential(response)
      except Exception as e:
        raise RuntimeError(
            f"Failed to retrieve credential for user '{user_id}' on connector"
            f" '{auth_scheme.name}'."
        ) from e

    if metadata is not None and metadata.uri_consent_required:
      if self._is_consent_completed(context):
        raise RuntimeError("Failed to retrieve consent based credential.")

      # Return AuthCredential with only auth_uri to trigger user consent flow.
      return AuthCredential(
          auth_type=AuthCredentialTypes.OAUTH2,
          oauth2=OAuth2Auth(
              auth_uri=metadata.uri_consent_required.authorization_uri,
              nonce=metadata.uri_consent_required.consent_nonce,
          ),
      )
