/*
 * 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.gravitino.s3.credential;

import java.io.IOException;
import java.net.URI;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Stream;
import org.apache.commons.lang3.StringUtils;
import org.apache.gravitino.credential.AwsIrsaCredential;
import org.apache.gravitino.credential.CredentialContext;
import org.apache.gravitino.credential.CredentialGenerator;
import org.apache.gravitino.credential.PathBasedCredentialContext;
import org.apache.gravitino.credential.config.S3CredentialConfig;
import org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSource;
import org.apache.gravitino.s3.credential.webidentity.WebIdentityTokenSources;
import software.amazon.awssdk.auth.credentials.AwsCredentials;
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
import software.amazon.awssdk.policybuilder.iam.IamConditionOperator;
import software.amazon.awssdk.policybuilder.iam.IamEffect;
import software.amazon.awssdk.policybuilder.iam.IamPolicy;
import software.amazon.awssdk.policybuilder.iam.IamResource;
import software.amazon.awssdk.policybuilder.iam.IamStatement;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.sts.StsClient;
import software.amazon.awssdk.services.sts.StsClientBuilder;
import software.amazon.awssdk.services.sts.auth.StsAssumeRoleWithWebIdentityCredentialsProvider;
import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityRequest;
import software.amazon.awssdk.services.sts.model.AssumeRoleWithWebIdentityResponse;
import software.amazon.awssdk.services.sts.model.Credentials;

/** Generate AWS IRSA credentials according to the read and write paths. */
public class AwsIrsaCredentialGenerator implements CredentialGenerator<AwsIrsaCredential> {

  private WebIdentityTokenSource tokenSource;
  private String roleArn;
  private int tokenExpireSecs;
  private String region;
  private String stsEndpoint;
  private boolean listLocationPrefix;
  private StsClient basicModeStsClient;
  private Map<String, StsAssumeRoleWithWebIdentityCredentialsProvider> basicModeProviders;

  @Override
  public void initialize(Map<String, String> properties) {
    this.tokenSource = WebIdentityTokenSources.create(properties);

    S3CredentialConfig s3CredentialConfig = new S3CredentialConfig(properties);
    this.roleArn = s3CredentialConfig.s3RoleArn();
    this.tokenExpireSecs = s3CredentialConfig.tokenExpireInSecs();
    this.region = s3CredentialConfig.region();
    this.stsEndpoint = s3CredentialConfig.stsEndpoint();
    this.listLocationPrefix = s3CredentialConfig.listLocationPrefix();
    this.basicModeStsClient = createStsClient();
    this.basicModeProviders = new ConcurrentHashMap<>();
  }

  @Override
  public AwsIrsaCredential generate(CredentialContext context) {
    if (!(context instanceof PathBasedCredentialContext)) {
      return toAwsIrsaCredential(createBasicModeCredentials(context.getUserName()));
    }

    PathBasedCredentialContext pathBasedCredentialContext = (PathBasedCredentialContext) context;

    Credentials s3Token =
        createCredentialsWithSessionPolicy(
            pathBasedCredentialContext.getReadPaths(),
            pathBasedCredentialContext.getWritePaths(),
            pathBasedCredentialContext.getUserName());
    return toAwsIrsaCredential(s3Token);
  }

  private Credentials createCredentialsWithSessionPolicy(
      Set<String> readLocations, Set<String> writeLocations, String userName) {
    validateInputParameters(readLocations, writeLocations, userName);

    IamPolicy sessionPolicy = createSessionPolicy(readLocations, writeLocations, region);
    return createCredentials(userName, Optional.of(sessionPolicy));
  }

  private Credentials createCredentials(String userName, Optional<IamPolicy> sessionPolicy) {
    String effectiveRoleArn = getValidatedRoleArn(roleArn);
    String tokenContent = tokenSource.getToken();
    return assumeRoleWithWebIdentity(effectiveRoleArn, userName, tokenContent, sessionPolicy);
  }

  private AwsCredentials createBasicModeCredentials(String userName) {
    return basicModeProviders
        .computeIfAbsent(userName, this::createBasicModeCredentialsProvider)
        .resolveCredentials();
  }

  private StsAssumeRoleWithWebIdentityCredentialsProvider createBasicModeCredentialsProvider(
      String userName) {
    return StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
        .stsClient(basicModeStsClient)
        .refreshRequest(
            () ->
                createAssumeRoleWithWebIdentityRequest(
                    getValidatedRoleArn(roleArn),
                    userName,
                    tokenSource.getToken(),
                    Optional.empty()))
        .build();
  }

  private AwsIrsaCredential toAwsIrsaCredential(Credentials credentials) {
    return new AwsIrsaCredential(
        credentials.accessKeyId(),
        credentials.secretAccessKey(),
        credentials.sessionToken(),
        credentials.expiration().toEpochMilli());
  }

  private AwsIrsaCredential toAwsIrsaCredential(AwsCredentials credentials) {
    if (!(credentials instanceof AwsSessionCredentials)) {
      throw new IllegalStateException(
          "AWS IRSA credentials must be of type AwsSessionCredentials. "
              + "Check your EKS/IRSA configuration. Got: "
              + credentials.getClass().getName());
    }
    AwsSessionCredentials sessionCredentials = (AwsSessionCredentials) credentials;
    if (!sessionCredentials.expirationTime().isPresent()) {
      throw new IllegalStateException(
          "AWS IRSA session credentials must include an expiration time for vended credential"
              + " refresh");
    }
    return new AwsIrsaCredential(
        sessionCredentials.accessKeyId(),
        sessionCredentials.secretAccessKey(),
        sessionCredentials.sessionToken(),
        sessionCredentials.expirationTime().get().toEpochMilli());
  }

  private IamPolicy createSessionPolicy(
      Set<String> readLocations, Set<String> writeLocations, String region) {
    IamPolicy.Builder policyBuilder = IamPolicy.builder();
    String arnPrefix = getArnPrefix(region);

    addReadPermissions(policyBuilder, readLocations, writeLocations, arnPrefix);
    if (!writeLocations.isEmpty()) {
      addWritePermissions(policyBuilder, writeLocations, arnPrefix);
    }
    addBucketPermissions(policyBuilder, readLocations, writeLocations, arnPrefix);

    return policyBuilder.build();
  }

  private void addReadPermissions(
      IamPolicy.Builder policyBuilder,
      Set<String> readLocations,
      Set<String> writeLocations,
      String arnPrefix) {
    IamStatement.Builder allowGetObjectStatementBuilder =
        IamStatement.builder()
            .effect(IamEffect.ALLOW)
            .addAction("s3:GetObject")
            .addAction("s3:GetObjectVersion");

    Stream.concat(readLocations.stream(), writeLocations.stream())
        .distinct()
        .forEach(
            location -> {
              URI uri = URI.create(location);
              allowGetObjectStatementBuilder.addResource(
                  IamResource.create(getS3UriWithArn(arnPrefix, uri)));
            });

    policyBuilder.addStatement(allowGetObjectStatementBuilder.build());
  }

  private void addWritePermissions(
      IamPolicy.Builder policyBuilder, Set<String> writeLocations, String arnPrefix) {
    IamStatement.Builder allowPutObjectStatementBuilder =
        IamStatement.builder()
            .effect(IamEffect.ALLOW)
            .addAction("s3:PutObject")
            .addAction("s3:DeleteObject");

    writeLocations.forEach(
        location -> {
          URI uri = URI.create(location);
          allowPutObjectStatementBuilder.addResource(
              IamResource.create(getS3UriWithArn(arnPrefix, uri)));
        });

    policyBuilder.addStatement(allowPutObjectStatementBuilder.build());
  }

  private void addBucketPermissions(
      IamPolicy.Builder policyBuilder,
      Set<String> readLocations,
      Set<String> writeLocations,
      String arnPrefix) {
    Map<String, IamStatement.Builder> bucketListStatementBuilder = new HashMap<>();
    Map<String, IamStatement.Builder> bucketGetLocationStatementBuilder = new HashMap<>();

    Stream.concat(readLocations.stream(), writeLocations.stream())
        .distinct()
        .forEach(
            location -> {
              URI uri = URI.create(location);
              String bucketArn = arnPrefix + getBucketName(uri);
              String rawPath = trimLeadingSlash(uri.getPath());

              bucketListStatementBuilder
                  .computeIfAbsent(
                      bucketArn,
                      key ->
                          IamStatement.builder()
                              .effect(IamEffect.ALLOW)
                              .addAction("s3:ListBucket")
                              .addResource(key))
                  .addConditions(
                      IamConditionOperator.STRING_LIKE,
                      "s3:prefix",
                      listPrefixes(rawPath, listLocationPrefix));

              bucketGetLocationStatementBuilder.computeIfAbsent(
                  bucketArn,
                  key ->
                      IamStatement.builder()
                          .effect(IamEffect.ALLOW)
                          .addAction("s3:GetBucketLocation")
                          .addResource(key));
            });

    addStatementsToPolicy(policyBuilder, bucketListStatementBuilder);
    addStatementsToPolicy(policyBuilder, bucketGetLocationStatementBuilder);
  }

  private void addStatementsToPolicy(
      IamPolicy.Builder policyBuilder, Map<String, IamStatement.Builder> statementBuilders) {
    statementBuilders.values().forEach(builder -> policyBuilder.addStatement(builder.build()));
  }

  private String getS3UriWithArn(String arnPrefix, URI uri) {
    return arnPrefix + addWildcardToPath(removeSchemaFromS3Uri(uri));
  }

  private String getArnPrefix(String region) {
    if (StringUtils.isNotBlank(region)) {
      if (region.contains("cn-")) {
        return "arn:aws-cn:s3:::";
      } else if (region.contains("us-gov-")) {
        return "arn:aws-us-gov:s3:::";
      }
    }
    return "arn:aws:s3:::";
  }

  private static String addWildcardToPath(String path) {
    return path.endsWith("/") ? path + "*" : path + "/*";
  }

  /**
   * Builds the {@code s3:prefix} condition values for a {@code ListBucket} statement. By default
   * the prefixes are restricted to the location and its descendants by requiring a trailing slash,
   * so a vended credential cannot enumerate keys in adjacent locations (e.g. {@code
   * path/to/table_new}) that merely share the location's string prefix. The empty path (bucket
   * root) is preserved so listing the whole bucket still works.
   *
   * <p>When {@code includeLocationPrefix} is {@code true}, the bare location prefix is also
   * allowed. This is required for the Hadoop FileSystem API, whose {@code getFileStatus} issues a
   * HEAD on the bare location key: that HEAD returns 404 (instead of 403) only if {@code
   * ListBucket} permits the bare prefix. The trade-off is that the credential can then enumerate
   * sibling key names sharing the prefix.
   *
   * @param rawPath the object key prefix of the location, without a leading slash
   * @param includeLocationPrefix whether to also allow the bare location prefix
   * @return the allowed {@code s3:prefix} values
   */
  static List<String> listPrefixes(String rawPath, boolean includeLocationPrefix) {
    String dirPrefix = rawPath.isEmpty() || rawPath.endsWith("/") ? rawPath : rawPath + "/";
    if (includeLocationPrefix && !rawPath.equals(dirPrefix)) {
      return Arrays.asList(rawPath, dirPrefix, addWildcardToPath(rawPath));
    }
    return Arrays.asList(dirPrefix, addWildcardToPath(rawPath));
  }

  private static String removeSchemaFromS3Uri(URI uri) {
    String bucket = uri.getHost();
    String path = trimLeadingSlash(uri.getPath());
    return String.join(
        "/", Stream.of(bucket, path).filter(Objects::nonNull).toArray(String[]::new));
  }

  private static String trimLeadingSlash(String path) {
    return path.startsWith("/") ? path.substring(1) : path;
  }

  private static String getBucketName(URI uri) {
    return uri.getHost();
  }

  private void validateInputParameters(
      Set<String> readLocations, Set<String> writeLocations, String userName) {
    if (StringUtils.isBlank(userName)) {
      throw new IllegalArgumentException("userName cannot be null or empty");
    }
    if ((readLocations == null || readLocations.isEmpty())
        && (writeLocations == null || writeLocations.isEmpty())) {
      throw new IllegalArgumentException("At least one read or write location must be specified");
    }
  }

  private String getValidatedRoleArn(String configRoleArn) {
    String effectiveRoleArn =
        StringUtils.isNotBlank(configRoleArn) ? configRoleArn : System.getenv("AWS_ROLE_ARN");
    if (StringUtils.isBlank(effectiveRoleArn)) {
      throw new IllegalStateException(
          "No role ARN available. Either configure s3-role-arn or ensure AWS_ROLE_ARN environment variable is set.");
    }
    if (!effectiveRoleArn.startsWith("arn:aws")) {
      throw new IllegalArgumentException("Invalid role ARN format: " + effectiveRoleArn);
    }
    return effectiveRoleArn;
  }

  private Credentials assumeRoleWithWebIdentity(
      String roleArn, String userName, String webIdentityToken, Optional<IamPolicy> sessionPolicy) {
    try (StsClient stsClient = createStsClient()) {
      AssumeRoleWithWebIdentityResponse response =
          stsClient.assumeRoleWithWebIdentity(
              createAssumeRoleWithWebIdentityRequest(
                  roleArn, userName, webIdentityToken, sessionPolicy));
      return response.credentials();
    }
  }

  private StsClient createStsClient() {
    StsClientBuilder stsBuilder = StsClient.builder();
    if (StringUtils.isNotBlank(region)) {
      stsBuilder.region(Region.of(region));
    }
    if (StringUtils.isNotBlank(stsEndpoint)) {
      stsBuilder.endpointOverride(URI.create(stsEndpoint));
    }
    return stsBuilder.build();
  }

  private AssumeRoleWithWebIdentityRequest createAssumeRoleWithWebIdentityRequest(
      String roleArn, String userName, String webIdentityToken, Optional<IamPolicy> sessionPolicy) {
    AssumeRoleWithWebIdentityRequest.Builder requestBuilder =
        AssumeRoleWithWebIdentityRequest.builder()
            .roleArn(roleArn)
            .roleSessionName("gravitino_irsa_session_" + userName)
            .durationSeconds(tokenExpireSecs)
            .webIdentityToken(webIdentityToken);
    sessionPolicy.ifPresent(policy -> requestBuilder.policy(policy.toJson()));
    return requestBuilder.build();
  }

  @Override
  public void close() throws IOException {
    if (basicModeProviders != null) {
      basicModeProviders.values().forEach(StsAssumeRoleWithWebIdentityCredentialsProvider::close);
      basicModeProviders.clear();
    }
    if (basicModeStsClient != null) {
      basicModeStsClient.close();
    }
    if (tokenSource != null) {
      tokenSource.close();
    }
  }
}
