/*
 * Jooby https://jooby.io
 * Apache License Version 2.0 https://jooby.io/LICENSE.txt
 * Copyright 2014 Edgar Espina
 */
package io.jooby.awssdkv1;

import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;

import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.AWSCredentialsProviderChain;
import com.amazonaws.auth.EC2ContainerCredentialsProviderWrapper;
import com.amazonaws.auth.EnvironmentVariableCredentialsProvider;
import com.amazonaws.auth.SystemPropertiesCredentialsProvider;
import com.amazonaws.auth.WebIdentityTokenCredentialsProvider;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.services.s3.transfer.TransferManager;
import com.typesafe.config.Config;
import edu.umd.cs.findbugs.annotations.NonNull;
import io.jooby.Extension;
import io.jooby.Jooby;
import io.jooby.ServiceRegistry;
import io.jooby.internal.awssdkv1.ServiceShutdown;

/**
 * Aws module for aws-java-sdk 1.x. This module:
 *
 * <p>- Integrates AWS credentials within application properties. - Register AWS services as
 * application services (so they can be used by require calls or DI). - Add graceful shutdown to AWS
 * services.
 *
 * <p>Usage:
 *
 * <pre>{@code
 * {
 *     install(
 *         new AwsModule()
 *             .setup(credentials -> {
 *               return TransferManagerBuilder.standard()
 *                   .withS3Client(
 *                       AmazonS3ClientBuilder.standard()
 *                           .withRegion(Regions.US_EAST_1)
 *                           .withCredentials(credentials)
 *                           .build()
 *                   ).build();
 *             })
 *     );
 * }
 * }</pre>
 *
 * <p>Previous example register AmazonS3Client and TransferManager services
 *
 * <p>NOTE: You need to add the required service dependency to your project.
 *
 * @author edgar
 */
public class AwsModule implements Extension {

  private final List<Function<AWSCredentialsProvider, Object>> factoryList = new ArrayList<>();
  private final AWSCredentialsProvider credentialsProvider;

  public AwsModule(@NonNull AWSCredentialsProvider credentialsProvider) {
    this.credentialsProvider = credentialsProvider;
  }

  public AwsModule() {
    this.credentialsProvider = null;
  }

  /**
   * Setup a new AWS service. Supported outputs are:
   *
   * <p>- Single amazon service - Stream of amazon services - Collection of amazon services
   *
   * <p>Each of the services returned by this function are added to the application service registry
   * and shutdown at application shutdown time.
   *
   * @param provider Service provider/factory.
   * @return AWS service.
   */
  public @NonNull AwsModule setup(@NonNull Function<AWSCredentialsProvider, Object> provider) {
    factoryList.add(provider);
    return this;
  }

  @Override
  public void install(@NonNull Jooby application) throws Exception {
    var credentialsProvider =
        Optional.ofNullable(this.credentialsProvider)
            .orElseGet(() -> newCredentialsProvider(application.getConfig()));
    List<Object> serviceList = new ArrayList<>(factoryList.size());
    for (Function<AWSCredentialsProvider, Object> factory : factoryList) {
      Object value = factory.apply(credentialsProvider);
      if (value instanceof Stream) {
        ((Stream) value).forEach(serviceList::add);
      } else if (value instanceof Collection) {
        ((Collection) value).forEach(serviceList::add);
      } else {
        extractServices(value, serviceList::add);
      }
    }
    ServiceRegistry services = application.getServices();
    // for each service
    for (Object service : serviceList) {
      Stream.of(
              service.getClass(),
              Stream.of(service.getClass().getInterfaces()).findFirst().orElse(null))
          .filter(Objects::nonNull)
          .forEach(
              serviceType -> {
                services.putIfAbsent((Class) serviceType, service);
              });
    }
    serviceList.stream()
        .distinct()
        .forEach(service -> application.onStop(new ServiceShutdown(application.getLog(), service)));
    serviceList.clear();
    factoryList.clear();
  }

  /**
   * Creates a credentials provider, exactly like {@link
   * com.amazonaws.auth.DefaultAWSCredentialsProviderChain} appending the application properties
   * provider.
   *
   * @param config Application properties.
   * @return Credentials provider.
   */
  public static @NonNull AWSCredentialsProvider newCredentialsProvider(@NonNull Config config) {
    return new AWSCredentialsProviderChain(
        new EnvironmentVariableCredentialsProvider(),
        new SystemPropertiesCredentialsProvider(),
        WebIdentityTokenCredentialsProvider.create(),
        new ProfileCredentialsProvider(),
        new EC2ContainerCredentialsProviderWrapper(),
        // Application configuration
        new ConfigCredentialsProvider(config));
  }

  private void extractServices(Object value, Consumer<Object> consumer) {
    transferManager(value, consumer);
    consumer.accept(value);
  }

  private void transferManager(Object value, Consumer<Object> consumer) {
    try {
      if (value instanceof TransferManager) {
        consumer.accept(((TransferManager) value).getAmazonS3Client());
      }
    } catch (NoClassDefFoundError x) {
      // s3 dependency is optional
    }
  }
}
