# 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.

import time
from typing import Optional
from unittest.mock import Mock
from unittest.mock import patch

from authlib.oauth2.rfc6749 import OAuth2Token
from fastapi.openapi.models import OAuth2
from fastapi.openapi.models import OAuthFlowAuthorizationCode
from fastapi.openapi.models import OAuthFlows
from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import OAuth2Auth
from google.adk.auth.auth_schemes import OpenIdConnectWithConfig
from google.adk.auth.oauth2_credential_util import create_oauth2_session
from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
import pytest


@pytest.fixture
def openid_connect_scheme() -> OpenIdConnectWithConfig:
  """Fixture providing a standard OpenIdConnectWithConfig scheme."""
  return OpenIdConnectWithConfig(
      type_="openIdConnect",
      openId_connect_url="https://example.com/.well-known/openid_configuration",
      authorization_endpoint="https://example.com/auth",
      token_endpoint="https://example.com/token",
      scopes=["openid", "profile"],
  )


def create_oauth2_auth_credential(
    auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
    token_endpoint_auth_method: Optional[str] = None,
):
  """Helper function to create OAuth2Auth credential with optional token_endpoint_auth_method."""
  oauth2_auth = OAuth2Auth(
      client_id="test_client_id",
      client_secret="test_client_secret",
      redirect_uri="https://example.com/callback",
      state="test_state",
  )
  if token_endpoint_auth_method is not None:
    oauth2_auth.token_endpoint_auth_method = token_endpoint_auth_method

  return AuthCredential(
      auth_type=auth_type,
      oauth2=oauth2_auth,
  )


class TestOAuth2CredentialUtil:
  """Test suite for OAuth2 credential utility functions."""

  def test_create_oauth2_session_openid_connect(self):
    """Test create_oauth2_session with OpenID Connect scheme."""
    scheme = OpenIdConnectWithConfig(
        type_="openIdConnect",
        openId_connect_url=(
            "https://example.com/.well-known/openid_configuration"
        ),
        authorization_endpoint="https://example.com/auth",
        token_endpoint="https://example.com/token",
        scopes=["openid", "profile"],
    )
    credential = create_oauth2_auth_credential(
        auth_type=AuthCredentialTypes.OAUTH2,
        token_endpoint_auth_method="client_secret_jwt",
    )

    client, token_endpoint = create_oauth2_session(scheme, credential)

    assert client is not None
    assert token_endpoint == "https://example.com/token"
    assert client.client_id == "test_client_id"
    assert client.client_secret == "test_client_secret"

  def test_create_oauth2_session_oauth2_scheme(self):
    """Test create_oauth2_session with OAuth2 scheme."""
    flows = OAuthFlows(
        authorizationCode=OAuthFlowAuthorizationCode(
            authorizationUrl="https://example.com/auth",
            tokenUrl="https://example.com/token",
            scopes={"read": "Read access", "write": "Write access"},
        )
    )
    scheme = OAuth2(type_="oauth2", flows=flows)
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OAUTH2,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
            redirect_uri="https://example.com/callback",
        ),
    )

    client, token_endpoint = create_oauth2_session(scheme, credential)

    assert client is not None
    assert token_endpoint == "https://example.com/token"

  def test_create_oauth2_session_invalid_scheme(self):
    """Test create_oauth2_session with invalid scheme."""
    scheme = Mock()  # Invalid scheme type
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OAUTH2,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
        ),
    )

    client, token_endpoint = create_oauth2_session(scheme, credential)

    assert client is None
    assert token_endpoint is None

  def test_create_oauth2_session_missing_credentials(self):
    """Test create_oauth2_session with missing credentials."""
    scheme = OpenIdConnectWithConfig(
        type_="openIdConnect",
        openId_connect_url=(
            "https://example.com/.well-known/openid_configuration"
        ),
        authorization_endpoint="https://example.com/auth",
        token_endpoint="https://example.com/token",
        scopes=["openid"],
    )
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            # Missing client_secret
        ),
    )

    client, token_endpoint = create_oauth2_session(scheme, credential)

    assert client is None
    assert token_endpoint is None

  def _google_openid_scheme(self) -> OpenIdConnectWithConfig:
    """OpenID Connect scheme that uses Google's OAuth2 token endpoint."""
    return OpenIdConnectWithConfig(
        type_="openIdConnect",
        openId_connect_url=(
            "https://accounts.google.com/.well-known/openid_configuration"
        ),
        authorization_endpoint="https://accounts.google.com/o/oauth2/v2/auth",
        token_endpoint="https://oauth2.googleapis.com/token",
        scopes=["openid"],
    )

  @patch.dict("os.environ", {}, clear=True)
  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
  def test_create_oauth2_session_google_endpoint_uses_mtls(
      self, mock_use_cert, mock_configure
  ):
    """Google token endpoint is switched to mTLS when a cert is mounted."""
    mock_use_cert.return_value = True
    mock_configure.return_value = True
    credential = create_oauth2_auth_credential(
        auth_type=AuthCredentialTypes.OAUTH2
    )

    client, token_endpoint = create_oauth2_session(
        self._google_openid_scheme(), credential
    )

    assert client is not None
    assert token_endpoint == "https://oauth2.mtls.googleapis.com/token"
    mock_configure.assert_called_once_with(client)

  @patch.dict("os.environ", {}, clear=True)
  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
  def test_create_oauth2_session_google_endpoint_no_cert_keeps_plain(
      self, mock_use_cert, mock_configure
  ):
    """Without a client cert the plain Google endpoint is kept."""
    mock_use_cert.return_value = False
    credential = create_oauth2_auth_credential(
        auth_type=AuthCredentialTypes.OAUTH2
    )

    _, token_endpoint = create_oauth2_session(
        self._google_openid_scheme(), credential
    )

    assert token_endpoint == "https://oauth2.googleapis.com/token"
    mock_configure.assert_not_called()

  @patch.dict("os.environ", {}, clear=True)
  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
  def test_create_oauth2_session_cert_unavailable_keeps_plain(
      self, mock_use_cert, mock_configure
  ):
    """If the adapter cannot be mounted, the endpoint is not switched."""
    mock_use_cert.return_value = True
    mock_configure.return_value = False
    credential = create_oauth2_auth_credential(
        auth_type=AuthCredentialTypes.OAUTH2
    )

    client, token_endpoint = create_oauth2_session(
        self._google_openid_scheme(), credential
    )

    assert token_endpoint == "https://oauth2.googleapis.com/token"
    mock_configure.assert_called_once_with(client)

  @patch.dict("os.environ", {}, clear=True)
  @patch("google.adk.utils._mtls_utils.configure_session_for_mtls")
  @patch("google.adk.utils._mtls_utils.use_client_cert_effective")
  def test_create_oauth2_session_non_google_endpoint_skips_mtls(
      self, mock_use_cert, mock_configure
  ):
    """Non-Google providers are never switched to an mTLS endpoint."""
    mock_use_cert.return_value = True
    credential = create_oauth2_auth_credential(
        auth_type=AuthCredentialTypes.OAUTH2
    )
    scheme = OpenIdConnectWithConfig(
        type_="openIdConnect",
        openId_connect_url=(
            "https://example.com/.well-known/openid_configuration"
        ),
        authorization_endpoint="https://example.com/auth",
        token_endpoint="https://example.com/token",
        scopes=["openid"],
    )

    _, token_endpoint = create_oauth2_session(scheme, credential)

    assert token_endpoint == "https://example.com/token"
    mock_configure.assert_not_called()

  @pytest.mark.parametrize(
      "token_endpoint_auth_method, expected_auth_method",
      [
          ("client_secret_post", "client_secret_post"),
          (None, "client_secret_basic"),
      ],
  )
  def test_create_oauth2_session_with_token_endpoint_auth_method(
      self,
      openid_connect_scheme,
      token_endpoint_auth_method,
      expected_auth_method,
  ):
    """Test create_oauth2_session with various token_endpoint_auth_method settings."""
    credential = create_oauth2_auth_credential(
        token_endpoint_auth_method=token_endpoint_auth_method
    )

    client, token_endpoint = create_oauth2_session(
        openid_connect_scheme, credential
    )

    assert client is not None
    assert token_endpoint == "https://example.com/token"
    assert client.client_id == "test_client_id"
    assert client.client_secret == "test_client_secret"
    assert client.token_endpoint_auth_method == expected_auth_method

  def test_create_oauth2_session_oauth2_scheme_with_token_endpoint_auth_method(
      self,
  ):
    """Test create_oauth2_session with OAuth2 scheme and token_endpoint_auth_method."""
    flows = OAuthFlows(
        authorizationCode=OAuthFlowAuthorizationCode(
            authorizationUrl="https://example.com/auth",
            tokenUrl="https://example.com/token",
            scopes={"read": "Read access", "write": "Write access"},
        )
    )
    scheme = OAuth2(type_="oauth2", flows=flows)
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OAUTH2,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
            redirect_uri="https://example.com/callback",
            token_endpoint_auth_method="client_secret_jwt",
        ),
    )

    client, token_endpoint = create_oauth2_session(scheme, credential)

    assert client is not None
    assert token_endpoint == "https://example.com/token"
    assert client.token_endpoint_auth_method == "client_secret_jwt"

  def _oauth2_scheme_with_scopes(self):
    """Build an OAuth2 scheme that declares scopes."""
    return OAuth2(
        type_="oauth2",
        flows=OAuthFlows(
            authorizationCode=OAuthFlowAuthorizationCode(
                authorizationUrl="https://example.com/auth",
                tokenUrl="https://example.com/token",
                scopes={"read": "Read access", "write": "Write access"},
            )
        ),
    )

  def _capturing_post(self, captured):
    """Stub for OAuth2Session.post that records the token-request body."""

    def _post(*args, **kwargs):
      captured["data"] = kwargs.get("data")
      response = Mock()
      response.status_code = 200
      response.json.return_value = {
          "access_token": "new_access_token",
          "token_type": "Bearer",
          "expires_in": 3600,
          "refresh_token": "new_refresh_token",
      }
      return response

    return _post

  def test_refresh_request_omits_scope(self):
    """Refresh requests must not carry scope (some providers reject it)."""
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OAUTH2,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
            redirect_uri="https://example.com/callback",
        ),
    )

    client, token_endpoint = create_oauth2_session(
        self._oauth2_scheme_with_scopes(), credential
    )
    assert client is not None

    captured = {}
    client.post = self._capturing_post(captured)
    client.refresh_token(token_endpoint, refresh_token="old_refresh_token")

    assert "scope" not in captured["data"]

  def test_token_exchange_omits_scope(self):
    """Authorization-code exchange must not carry scope (it is redundant)."""
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OAUTH2,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
            redirect_uri="https://example.com/callback",
        ),
    )

    client, token_endpoint = create_oauth2_session(
        self._oauth2_scheme_with_scopes(), credential
    )
    assert client is not None

    captured = {}
    client.post = self._capturing_post(captured)
    client.fetch_token(
        token_endpoint, grant_type="authorization_code", code="test_code"
    )

    assert "scope" not in captured["data"]

  def test_update_credential_with_tokens(self):
    """Test update_credential_with_tokens function."""
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
        oauth2=OAuth2Auth(
            client_id="test_client_id",
            client_secret="test_client_secret",
        ),
    )

    # Store the expected expiry time to avoid timing issues
    expected_expires_at = int(time.time()) + 3600
    tokens = OAuth2Token({
        "access_token": "new_access_token",
        "refresh_token": "new_refresh_token",
        "id_token": "new_id_token",
        "expires_at": expected_expires_at,
        "expires_in": 3600,
    })

    assert credential.oauth2 is not None

    update_credential_with_tokens(credential, tokens)

    assert credential.oauth2.access_token == "new_access_token"
    assert credential.oauth2.refresh_token == "new_refresh_token"
    assert credential.oauth2.id_token == "new_id_token"
    assert credential.oauth2.expires_at == expected_expires_at
    assert credential.oauth2.expires_in == 3600

  def test_update_credential_with_tokens_none(self) -> None:
    credential = AuthCredential(
        auth_type=AuthCredentialTypes.API_KEY,
    )
    tokens = OAuth2Token({"access_token": "new_access_token"})

    # Should not raise any exceptions when oauth2 is None
    update_credential_with_tokens(credential, tokens)
    assert credential.oauth2 is None