/*
 * Copyright 2025 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.
 */

package com.google.cloud.auth.samples.customcredentials.aws;

// [START auth_custom_credential_supplier_aws]
import com.google.auth.oauth2.AwsCredentials;
import com.google.auth.oauth2.AwsSecurityCredentials;
import com.google.auth.oauth2.AwsSecurityCredentialsSupplier;
import com.google.auth.oauth2.ExternalAccountSupplierContext;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.cloud.storage.Bucket;
import com.google.cloud.storage.Storage;
import com.google.cloud.storage.StorageOptions;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import java.io.IOException;
import java.io.Reader;
import java.lang.reflect.Type;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Map;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;

// [END auth_custom_credential_supplier_aws]

/**
 * This sample demonstrates how to use a custom AWS security credentials supplier to authenticate to
 * Google Cloud Storage using AWS Workload Identity Federation.
 */
public class CustomCredentialSupplierAwsWorkload {

  public static void main(String[] args) throws IOException {

    // Reads the custom-credentials-aws-secrets.json if running locally.
    loadConfigFromFile();

    // The audience for the workload identity federation.
    // Format: //iam.googleapis.com/projects/<project-number>/locations/global/
    //         workloadIdentityPools/<pool-id>/providers/<provider-id>
    String gcpWorkloadAudience = getConfiguration("GCP_WORKLOAD_AUDIENCE");

    // The bucket to fetch data from.
    String gcsBucketName = getConfiguration("GCS_BUCKET_NAME");

    // (Optional) The service account impersonation URL.
    String saImpersonationUrl = getConfiguration("GCP_SERVICE_ACCOUNT_IMPERSONATION_URL");

    if (gcpWorkloadAudience == null || gcsBucketName == null) {
      System.err.println(
          "Required configuration missing. Please provide it in a "
              + "custom-credentials-aws-secrets.json file or as environment variables: "
              + "GCP_WORKLOAD_AUDIENCE, GCS_BUCKET_NAME");
      return;
    }

    try {
      System.out.println("Retrieving metadata for bucket: " + gcsBucketName + "...");
      Bucket bucket =
          authenticateWithAwsCredentials(gcpWorkloadAudience, saImpersonationUrl, gcsBucketName);

      System.out.println(" --- SUCCESS! ---");
      System.out.println("Bucket details:");
      System.out.printf("  Name: %s%n", bucket.getName());
      System.out.printf("  Location: %s%n", bucket.getLocation());
      System.out.printf("  Storage Class: %s%n", bucket.getStorageClass());
      System.out.printf("  Metageneration: %s%n", bucket.getMetageneration());
    } catch (Exception e) {
      System.err.println("Authentication or Request failed: " + e.getMessage());
    }
  }

  /**
   * Helper method to retrieve configuration. It checks Environment variables first, then System
   * properties (populated by loadConfigFromFile).
   */
  static String getConfiguration(String key) {
    String value = System.getenv(key);
    if (value == null) {
      value = System.getProperty(key);
    }
    return value;
  }

  /**
   * If a local secrets file is present, load it into the System Properties. This is a
   * "just-in-time" configuration for local development. These variables are only set for the
   * current process.
   */
  static void loadConfigFromFile() {
    // By default, this expects the file to be in the project root.
    String secretsFilePath = "custom-credentials-aws-secrets.json";
    if (!Files.exists(Paths.get(secretsFilePath))) {
      return;
    }

    try (Reader reader = Files.newBufferedReader(Paths.get(secretsFilePath))) {
      // Use Gson to parse the JSON file into a Map
      Gson gson = new Gson();
      Type type = new TypeToken<Map<String, String>>() {}.getType();
      Map<String, String> secrets = gson.fromJson(reader, type);

      if (secrets == null) {
        return;
      }

      // AWS SDK for Java looks for System Properties with specific names (camelCase)
      // if environment variables are missing.
      if (secrets.containsKey("aws_access_key_id")) {
        System.setProperty("aws.accessKeyId", secrets.get("aws_access_key_id"));
      }
      if (secrets.containsKey("aws_secret_access_key")) {
        System.setProperty("aws.secretAccessKey", secrets.get("aws_secret_access_key"));
      }
      if (secrets.containsKey("aws_region")) {
        System.setProperty("aws.region", secrets.get("aws_region"));
      }

      // Set custom GCP variables as System Properties so getConfiguration() can find them.
      if (secrets.containsKey("gcp_workload_audience")) {
        System.setProperty("GCP_WORKLOAD_AUDIENCE", secrets.get("gcp_workload_audience"));
      }
      if (secrets.containsKey("gcs_bucket_name")) {
        System.setProperty("GCS_BUCKET_NAME", secrets.get("gcs_bucket_name"));
      }
      if (secrets.containsKey("gcp_service_account_impersonation_url")) {
        System.setProperty(
            "GCP_SERVICE_ACCOUNT_IMPERSONATION_URL",
            secrets.get("gcp_service_account_impersonation_url"));
      }

    } catch (IOException e) {
      System.err.println("Error reading secrets file: " + e.getMessage());
    }
  }

  /**
   * Authenticates using a custom AWS credential supplier and retrieves bucket metadata.
   *
   * @param gcpWorkloadAudience The WIF provider audience.
   * @param saImpersonationUrl Optional service account impersonation URL.
   * @param gcsBucketName The GCS bucket name.
   * @return The Bucket object containing metadata.
   * @throws IOException If authentication fails.
   */
  // [START auth_custom_credential_supplier_aws]
  public static Bucket authenticateWithAwsCredentials(
      String gcpWorkloadAudience, String saImpersonationUrl, String gcsBucketName)
      throws IOException {

    CustomAwsSupplier customSupplier = new CustomAwsSupplier();

    AwsCredentials.Builder credentialsBuilder =
        AwsCredentials.newBuilder()
            .setAudience(gcpWorkloadAudience)
            // This token type indicates that the subject token is an AWS Signature Version 4 signed
            // request. This is required for AWS Workload Identity Federation.
            .setSubjectTokenType("urn:ietf:params:aws:token-type:aws4_request")
            .setAwsSecurityCredentialsSupplier(customSupplier);

    if (saImpersonationUrl != null) {
      credentialsBuilder.setServiceAccountImpersonationUrl(saImpersonationUrl);
    }

    GoogleCredentials credentials = credentialsBuilder.build();

    Storage storage = StorageOptions.newBuilder().setCredentials(credentials).build().getService();

    return storage.get(gcsBucketName);
  }

  /**
   * Custom AWS Security Credentials Supplier.
   *
   * <p>This implementation resolves AWS credentials and regions using the default provider chains
   * from the AWS SDK (v2). This supports environment variables, ~/.aws/credentials, and EC2/EKS
   * metadata.
   */
  private static class CustomAwsSupplier implements AwsSecurityCredentialsSupplier {
    private final AwsCredentialsProvider awsCredentialsProvider;
    private String region;

    public CustomAwsSupplier() {
      // The AWS SDK handles caching internally.
      this.awsCredentialsProvider = DefaultCredentialsProvider.create();
    }

    @Override
    public String getRegion(ExternalAccountSupplierContext context) {
      if (this.region == null) {
        Region awsRegion = new DefaultAwsRegionProviderChain().getRegion();
        if (awsRegion == null) {
          throw new IllegalStateException(
              "Unable to resolve AWS region. Ensure AWS_REGION is set or configured.");
        }
        this.region = awsRegion.id();
      }
      return this.region;
    }

    @Override
    public AwsSecurityCredentials getCredentials(ExternalAccountSupplierContext context) {
      software.amazon.awssdk.auth.credentials.AwsCredentials credentials =
          this.awsCredentialsProvider.resolveCredentials();

      if (credentials == null) {
        throw new IllegalStateException("Unable to resolve AWS credentials.");
      }

      String sessionToken = null;
      if (credentials instanceof AwsSessionCredentials) {
        sessionToken = ((AwsSessionCredentials) credentials).sessionToken();
      }

      return new AwsSecurityCredentials(
          credentials.accessKeyId(), credentials.secretAccessKey(), sessionToken);
    }
  }
  // [END auth_custom_credential_supplier_aws]
}
