/*
 * 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.aws;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.when;

import java.net.URI;
import java.util.Map;
import org.apache.iceberg.catalog.Namespace;
import org.apache.iceberg.rest.HTTPHeaders;
import org.apache.iceberg.rest.HTTPHeaders.HTTPHeader;
import org.apache.iceberg.rest.HTTPRequest;
import org.apache.iceberg.rest.HTTPRequest.HTTPMethod;
import org.apache.iceberg.rest.ImmutableHTTPRequest;
import org.apache.iceberg.rest.auth.AuthSession;
import org.apache.iceberg.rest.requests.CreateNamespaceRequest;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.signer.Aws4Signer;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.utils.SdkAutoCloseable;

class TestRESTSigV4AuthSession {

  private final Aws4Signer signer = Aws4Signer.create();

  private final AwsProperties awsProperties =
      new AwsProperties(
          Map.of(
              // CI environment doesn't have credentials, but a value must be set for signing
              AwsProperties.REST_SIGNER_REGION,
              "us-west-2",
              AwsProperties.REST_ACCESS_KEY_ID,
              "id",
              AwsProperties.REST_SECRET_ACCESS_KEY,
              "secret"));

  @Test
  void nullArguments() {
    AuthSession delegate = Mockito.mock(AuthSession.class);
    assertThatThrownBy(() -> new RESTSigV4AuthSession(null, delegate, awsProperties))
        .isInstanceOf(NullPointerException.class)
        .hasMessage("Invalid signer: null");
    assertThatThrownBy(() -> new RESTSigV4AuthSession(signer, null, awsProperties))
        .isInstanceOf(NullPointerException.class)
        .hasMessage("Invalid delegate: null");
    assertThatThrownBy(() -> new RESTSigV4AuthSession(signer, delegate, null))
        .isInstanceOf(NullPointerException.class)
        .hasMessage("Invalid AWS properties: null");
  }

  @Test
  void authenticateWithoutBody() {
    HTTPRequest request = Mockito.mock(HTTPRequest.class);
    AuthSession delegate = Mockito.mock(AuthSession.class);
    when(delegate.authenticate(any()))
        .thenReturn(
            ImmutableHTTPRequest.builder()
                .method(HTTPMethod.GET)
                .baseUri(URI.create("http://localhost:8080"))
                .path("path")
                .headers(
                    HTTPHeaders.of(
                        HTTPHeader.of("Content-Type", "application/json"),
                        HTTPHeader.of("Content-Encoding", "gzip")))
                .build());
    try (RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, awsProperties)) {
      HTTPRequest actual = session.authenticate(request);
      assertThat(actual.headers().entries())
          .hasSize(6)
          // original headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Type");
                assertThat(header.value()).isEqualTo("application/json");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Encoding");
                assertThat(header.value()).isEqualTo("gzip");
              })
          // generated by the signer
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Host");
                assertThat(header.value()).isEqualTo("localhost:8080");
              })
          // Sigv4 headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Authorization");
                assertThat(header.value()).startsWith("AWS4-HMAC-SHA256 Credential=");
                assertThat(header.value())
                    .contains(
                        "SignedHeaders=content-encoding;content-type;host;x-amz-content-sha256;x-amz-date");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("x-amz-content-sha256");
                assertThat(header.value())
                    .isEqualTo("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
              })
          .satisfiesOnlyOnce(header -> assertThat(header.name()).isEqualTo("X-Amz-Date"));
    }
  }

  @Test
  void authenticateWithBody() {
    HTTPRequest request = Mockito.mock(HTTPRequest.class);
    AuthSession delegate = Mockito.mock(AuthSession.class);
    when(delegate.authenticate(any()))
        .thenReturn(
            ImmutableHTTPRequest.builder()
                .method(HTTPMethod.POST)
                .baseUri(URI.create("http://localhost:8080"))
                .path("path")
                .headers(
                    HTTPHeaders.of(
                        HTTPHeader.of("Content-Type", "application/x-www-form-urlencoded"),
                        HTTPHeader.of("Content-Encoding", "gzip")))
                .body(CreateNamespaceRequest.builder().withNamespace(Namespace.of("ns")).build())
                .build());
    try (RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, awsProperties)) {
      HTTPRequest actual = session.authenticate(request);
      assertThat(actual.headers().entries())
          .hasSize(6)
          // original headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Type");
                assertThat(header.value()).isEqualTo("application/x-www-form-urlencoded");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Encoding");
                assertThat(header.value()).isEqualTo("gzip");
              })
          // generated by the signer
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Host");
                assertThat(header.value()).isEqualTo("localhost:8080");
              })
          // Sigv4 headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Authorization");
                assertThat(header.value()).startsWith("AWS4-HMAC-SHA256 Credential=");
                assertThat(header.value())
                    .contains(
                        "SignedHeaders=content-encoding;content-type;host;x-amz-content-sha256;x-amz-date");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("x-amz-content-sha256");
                assertThat(header.value())
                    .isEqualTo("yc5oAKPWjHY4sW8XQq0l/3aNrrXJKBycVFNnDEGMfww=");
              })
          .satisfiesOnlyOnce(header -> assertThat(header.name()).isEqualTo("X-Amz-Date"));
    }
  }

  @Test
  void authenticateConflictingAuthorizationHeader() {
    HTTPRequest request = Mockito.mock(HTTPRequest.class);
    AuthSession delegate = Mockito.mock(AuthSession.class);
    when(delegate.authenticate(any()))
        .thenReturn(
            ImmutableHTTPRequest.builder()
                .method(HTTPMethod.GET)
                .baseUri(URI.create("http://localhost:8080"))
                .path("path")
                .headers(
                    HTTPHeaders.of(
                        HTTPHeader.of("Content-Type", "application/json"),
                        HTTPHeader.of("Authorization", "Bearer token")))
                .build());
    try (RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, awsProperties)) {
      HTTPRequest actual = session.authenticate(request);
      assertThat(actual.headers().entries())
          .hasSize(6)
          // original header
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Type");
                assertThat(header.value()).isEqualTo("application/json");
              })
          // relocated before signing
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Original-Authorization");
                assertThat(header.value()).isEqualTo("Bearer token");
              })
          // generated by the signer
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Host");
                assertThat(header.value()).isEqualTo("localhost:8080");
              })
          // Sigv4 headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Authorization");
                assertThat(header.value()).startsWith("AWS4-HMAC-SHA256 Credential=");
                assertThat(header.value())
                    .contains(
                        "SignedHeaders=content-type;host;original-authorization;x-amz-content-sha256;x-amz-date");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("x-amz-content-sha256");
                assertThat(header.value())
                    .isEqualTo("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
              })
          .satisfiesOnlyOnce(header -> assertThat(header.name()).isEqualTo("X-Amz-Date"));
    }
  }

  @Test
  void authenticateConflictingSigv4Headers() {
    HTTPRequest request = Mockito.mock(HTTPRequest.class);
    AuthSession delegate = Mockito.mock(AuthSession.class);
    when(delegate.authenticate(any()))
        .thenReturn(
            ImmutableHTTPRequest.builder()
                .method(HTTPMethod.GET)
                .baseUri(URI.create("http://localhost:8080"))
                .path("path")
                .headers(
                    HTTPHeaders.of(
                        HTTPHeader.of("Content-Type", "application/json"),
                        HTTPHeader.of("x-amz-content-sha256", "fake"),
                        // corner case: conflicting header with same value as generated by the
                        // signer: will not be relocated
                        HTTPHeader.of(
                            "x-amz-content-sha256",
                            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"),
                        HTTPHeader.of("X-Amz-Date", "fake")))
                .build());
    try (RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, awsProperties)) {
      HTTPRequest actual = session.authenticate(request);
      assertThat(actual.headers().entries())
          .hasSize(7)
          // original header
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Content-Type");
                assertThat(header.value()).isEqualTo("application/json");
              })
          // relocated after signing
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Original-x-amz-content-sha256");
                assertThat(header.value()).isEqualTo("fake");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Original-X-Amz-Date");
                assertThat(header.value()).isEqualTo("fake");
              })
          // generated by the signer
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Host");
                assertThat(header.value()).isEqualTo("localhost:8080");
              })
          // Sigv4 headers
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("Authorization");
                assertThat(header.value()).startsWith("AWS4-HMAC-SHA256 Credential=");
                assertThat(header.value())
                    .contains("SignedHeaders=content-type;host;x-amz-content-sha256;x-amz-date");
              })
          .satisfiesOnlyOnce(
              header -> {
                assertThat(header.name()).isEqualTo("x-amz-content-sha256");
                assertThat(header.value())
                    .isEqualTo("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
              })
          .satisfiesOnlyOnce(header -> assertThat(header.name()).isEqualTo("X-Amz-Date"));
    }
  }

  @Test
  void close() {
    AuthSession delegate = Mockito.mock(AuthSession.class);
    RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, awsProperties);
    session.close();
    Mockito.verify(delegate).close();
  }

  @Test
  void closeWithCloseableCredentialsProvider() {
    AuthSession delegate = Mockito.mock(AuthSession.class);
    CloseableAwsCredentialsProvider credentialsProvider =
        Mockito.mock(CloseableAwsCredentialsProvider.class);
    closeWithCloseableCredentialsProvider(delegate, credentialsProvider);
  }

  @Test
  void closeSuppressesFailure() {
    AuthSession delegate = Mockito.mock(AuthSession.class);
    Mockito.doThrow(new RuntimeException("delegate close failed")).when(delegate).close();
    CloseableAwsCredentialsProvider credentialsProvider =
        Mockito.mock(CloseableAwsCredentialsProvider.class);
    Mockito.doThrow(new RuntimeException("credentials provider close failed"))
        .when(credentialsProvider)
        .close();
    closeWithCloseableCredentialsProvider(delegate, credentialsProvider);
  }

  private void closeWithCloseableCredentialsProvider(
      AuthSession delegate, CloseableAwsCredentialsProvider credentialsProvider) {
    AwsProperties properties = Mockito.mock(AwsProperties.class);
    when(properties.restSigningRegion()).thenReturn(Region.US_WEST_2);
    when(properties.restSigningName()).thenReturn("execute-api");
    when(properties.restCredentialsProvider()).thenReturn(credentialsProvider);

    RESTSigV4AuthSession session = new RESTSigV4AuthSession(signer, delegate, properties);
    session.close();

    Mockito.verify(delegate).close();
    Mockito.verify(credentialsProvider).close();
  }

  interface CloseableAwsCredentialsProvider extends AwsCredentialsProvider, SdkAutoCloseable {
    @Override
    void close();
  }
}
