/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
package org.apache.iceberg.rest.auth;

import static org.apache.hadoop.hdfs.web.oauth2.OAuth2Constants.BEARER;
import static org.apache.hadoop.hdfs.web.oauth2.OAuth2Constants.CLIENT_CREDENTIALS;
import static org.apache.hadoop.hdfs.web.oauth2.OAuth2Constants.GRANT_TYPE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.ArgumentMatchers.argThat;

import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.PlainJWT;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
import org.apache.iceberg.rest.RESTClient;
import org.apache.iceberg.rest.auth.OAuth2Util.AuthSession;
import org.apache.iceberg.rest.responses.OAuthTokenResponse;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;

public class TestOAuth2Util {
  @Test
  public void testOAuthScopeTokenValidation() {
    // valid scope tokens are from ascii 0x21 to 0x7E, excluding 0x22 (") and 0x5C (\)
    // test characters that are outside of the ! to ~ range and the excluded characters, " and \
    assertThat(OAuth2Util.isValidScopeToken("a\\b"))
        .as("Should reject scope token with \\")
        .isFalse();
    assertThat(OAuth2Util.isValidScopeToken("a b"))
        .as("Should reject scope token with space")
        .isFalse();
    assertThat(OAuth2Util.isValidScopeToken("a\"b"))
        .as("Should reject scope token with \"")
        .isFalse();
    assertThat(OAuth2Util.isValidScopeToken("\u007F"))
        .as("Should reject scope token with DEL")
        .isFalse();
    // test all characters that are inside of the ! to ~ range and are not excluded
    assertThat(OAuth2Util.isValidScopeToken("!#$%&'()*+,-./"))
        .as("Should accept scope token with !-/")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("0123456789"))
        .as("Should accept scope token with 0-9")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken(":;<=>?@"))
        .as("Should accept scope token with :-@")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("ABCDEFGHIJKLM"))
        .as("Should accept scope token with A-M")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("NOPQRSTUVWXYZ"))
        .as("Should accept scope token with N-Z")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("[]^_`"))
        .as("Should accept scope token with [-`, not \\")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("abcdefghijklm"))
        .as("Should accept scope token with a-m")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("nopqrstuvwxyz"))
        .as("Should accept scope token with n-z")
        .isTrue();
    assertThat(OAuth2Util.isValidScopeToken("{|}~"))
        .as("Should accept scope token with {-~")
        .isTrue();
  }

  @Test
  public void testExpiresAt() {
    assertThat(OAuth2Util.expiresAtMillis(null)).isNull();
    assertThat(OAuth2Util.expiresAtMillis("not a token")).isNull();
    assertThat(OAuth2Util.expiresAtMillis("a.b.c")).isNull();

    // expires at epoch second = 1
    String token =
        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyLCJleHAiOjF9.gQADTbdEv-rpDWKSkGLbmafyB5UUjTdm9B_1izpuZ6E";
    assertThat(OAuth2Util.expiresAtMillis(token)).isEqualTo(1000L);

    // expires at epoch second = 19999999999
    token =
        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyLCJleHAiOjE5OTk5OTk5OTk5fQ._3k92KJi2NTyTG6V1s2mzJ__GiQtL36DnzsZSkBdYPw";
    assertThat(OAuth2Util.expiresAtMillis(token)).isEqualTo(19999999999000L);

    // no expiration
    token =
        "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c";
    assertThat(OAuth2Util.expiresAtMillis(token)).isNull();
  }

  @Test
  public void testCredentialFlowForSessionRefresh() throws IOException {
    AuthConfig authConfig =
        ImmutableAuthConfig.builder()
            .keepRefreshed(true)
            .exchangeEnabled(false)
            .token("test_token")
            .credential("testClientId:testClientSecret")
            .oauth2ServerUri("/v1/token")
            .expiresAtMillis(System.currentTimeMillis())
            .build();

    OAuthTokenResponse response =
        OAuthTokenResponse.builder().withToken("refreshed_token").withTokenType(BEARER).build();

    try (RESTClient client = Mockito.mock(RESTClient.class);
        OAuth2Util.AuthSession session = new OAuth2Util.AuthSession(Map.of(), authConfig); ) {
      Mockito.when(client.postForm(any(), anyMap(), any(), anyMap(), any())).thenReturn(response);

      session.refresh(client);

      Mockito.verify(client)
          .postForm(
              any(),
              argThat(
                  formData ->
                      formData.containsKey(GRANT_TYPE)
                          && formData.get(GRANT_TYPE).equals(CLIENT_CREDENTIALS)),
              Mockito.eq(OAuthTokenResponse.class),
              anyMap(),
              any());
    }
  }

  @Test
  void fromTokenResponseUsesChildTokenExpiry() {
    AuthSession parent = parentSession(7200);
    OAuthTokenResponse response = childTokenResponse(tokenWithExp(300), 300);

    AuthSession child =
        AuthSession.fromTokenResponse(null, null, response, System.currentTimeMillis(), parent);
    assertThat(child.expiresAtMillis())
        .as("Child session should use the child token's exp, not the parent's")
        .isEqualTo(TimeUnit.SECONDS.toMillis(300));
  }

  @Test
  void fromTokenResponseOpaqueTokenDoesNotInheritParentExpiry() {
    AuthSession parent = parentSession(7200);
    OAuthTokenResponse response = childTokenResponse("opaque-access-token", 600);

    AuthSession child =
        AuthSession.fromTokenResponse(null, null, response, System.currentTimeMillis(), parent);

    assertThat(child.expiresAtMillis())
        .as("Child session with opaque token should not inherit parent's expiresAtMillis")
        .isNull();
  }

  @Test
  void fromAccessTokenUsesChildTokenExpiry() {
    AuthSession parent = parentSession(7200);
    String childToken = tokenWithExp(300);

    AuthSession child = AuthSession.fromAccessToken(null, null, childToken, null, parent);
    assertThat(child.expiresAtMillis())
        .as("Child session should use the child token's exp, not the parent's")
        .isEqualTo(TimeUnit.SECONDS.toMillis(300));
  }

  @Test
  void refreshUsesRefreshedTokenExpiry() throws IOException {
    String parentToken = tokenWithExp(7200);
    String refreshedToken = tokenWithExp(500);

    AuthConfig authConfig =
        AuthConfig.builder()
            .token(parentToken)
            .tokenType(OAuth2Properties.ACCESS_TOKEN_TYPE)
            .keepRefreshed(true)
            .credential("testClientId:testClientSecret")
            .oauth2ServerUri("/v1/token")
            .expiresAtMillis(OAuth2Util.expiresAtMillis(parentToken))
            .build();

    OAuthTokenResponse response = childTokenResponse(refreshedToken, 500);

    try (RESTClient client = Mockito.mock(RESTClient.class);
        AuthSession session = new AuthSession(Map.of(), authConfig)) {
      Mockito.when(client.postForm(any(), anyMap(), any(), anyMap(), any())).thenReturn(response);

      session.refresh(client);

      assertThat(session.expiresAtMillis())
          .as("After refresh, session should use the refreshed token's exp")
          .isEqualTo(TimeUnit.SECONDS.toMillis(500));
    }
  }

  @Test
  void refreshExpiredTokenShouldIncludeOptionalOAuthParams() throws IOException {
    assertRefreshIncludesOptionalOAuthParams(System.currentTimeMillis() - 10_000);
  }

  @Test
  void refreshCurrentTokenNonExchangeShouldIncludeOptionalOAuthParams() throws IOException {
    assertRefreshIncludesOptionalOAuthParams(System.currentTimeMillis() + 300_000);
  }

  private static void assertRefreshIncludesOptionalOAuthParams(long expiresAtMillis)
      throws IOException {
    String audience = "https://my-catalog.example.com";
    AuthConfig authConfig =
        ImmutableAuthConfig.builder()
            .keepRefreshed(true)
            .exchangeEnabled(false)
            .token("token")
            .credential("testClientId:testClientSecret")
            .oauth2ServerUri("/v1/token")
            .expiresAtMillis(expiresAtMillis)
            .optionalOAuthParams(ImmutableMap.of("audience", audience, "scope", "catalog"))
            .build();

    OAuthTokenResponse response =
        OAuthTokenResponse.builder().withToken("refreshed_token").withTokenType(BEARER).build();

    try (RESTClient client = Mockito.mock(RESTClient.class);
        OAuth2Util.AuthSession session = new OAuth2Util.AuthSession(Map.of(), authConfig)) {
      Mockito.when(client.postForm(any(), anyMap(), any(), anyMap(), any())).thenReturn(response);

      session.refresh(client);

      Mockito.verify(client)
          .postForm(
              any(),
              argThat(
                  formData ->
                      CLIENT_CREDENTIALS.equals(formData.get(GRANT_TYPE))
                          && audience.equals(formData.get("audience"))
                          && "catalog".equals(formData.get("scope"))),
              Mockito.eq(OAuthTokenResponse.class),
              anyMap(),
              any());
    }
  }

  private static AuthSession parentSession(long expSeconds) {
    String parentToken = tokenWithExp(expSeconds);
    AuthConfig parentConfig =
        AuthConfig.builder()
            .token(parentToken)
            .tokenType(OAuth2Properties.ACCESS_TOKEN_TYPE)
            .keepRefreshed(false)
            .build();
    AuthSession parent = new AuthSession(Map.of(), parentConfig);
    assertThat(parent.expiresAtMillis()).isEqualTo(TimeUnit.SECONDS.toMillis(expSeconds));
    return parent;
  }

  private static OAuthTokenResponse childTokenResponse(String token, int expiresInSeconds) {
    return OAuthTokenResponse.builder()
        .withToken(token)
        .withTokenType(BEARER)
        .withIssuedTokenType(OAuth2Properties.ACCESS_TOKEN_TYPE)
        .setExpirationInSeconds(expiresInSeconds)
        .build();
  }

  private static String tokenWithExp(long expSeconds) {
    JWTClaimsSet claimsSet =
        new JWTClaimsSet.Builder().subject("test").claim("exp", expSeconds).build();
    return new PlainJWT(claimsSet).serialize();
  }
}
