package com.google.cloud.hadoop.gcsio;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Strings.isNullOrEmpty;

import com.google.api.ClientProto;
import com.google.api.client.auth.oauth2.Credential;
import com.google.api.client.googleapis.compute.ComputeCredential;
import com.google.auth.Credentials;
import com.google.auth.oauth2.ComputeEngineCredentials;
import com.google.cloud.hadoop.util.CredentialAdapter;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.protobuf.util.Durations;
import com.google.storage.v2.StorageGrpc;
import com.google.storage.v2.StorageGrpc.StorageBlockingStub;
import com.google.storage.v2.StorageGrpc.StorageStub;
import com.google.storage.v2.StorageProto;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.Status;
import io.grpc.alts.GoogleDefaultChannelBuilder;
import io.grpc.auth.MoreCallCredentials;
import io.grpc.stub.AbstractStub;
import java.util.Map;
import java.util.concurrent.ExecutorService;

/** Provides gRPC stubs for accessing the Storage gRPC API. */
class StorageStubProvider {

  // The maximum number of times to automatically retry gRPC requests.
  private static final double GRPC_MAX_RETRY_ATTEMPTS = 10;

  private static final ImmutableSet<Status.Code> STUB_BROKEN_ERROR_CODES =
      ImmutableSet.of(Status.Code.DEADLINE_EXCEEDED, Status.Code.UNAVAILABLE);

  // The GCS gRPC server address.
  private static final String DEFAULT_GCS_GRPC_SERVER_ADDRESS =
      StorageProto.getDescriptor()
          .findServiceByName("Storage")
          .getOptions()
          .getExtension(ClientProto.defaultHost);

  private final GoogleCloudStorageReadOptions readOptions;
  private final String userAgent;
  private final ExecutorService backgroundTasksThreadPool;
  private final GrpcDecorator grpcDecorator;
  private ManagedChannel channel;

  @VisibleForTesting
  GrpcDecorator getGrpcDecorator() {
    return grpcDecorator;
  }

  StorageStubProvider(
      GoogleCloudStorageOptions options,
      ExecutorService backgroundTasksThreadPool,
      GrpcDecorator grpcDecorator) {
    this.readOptions = options.getReadChannelOptions();
    this.userAgent = options.getAppName();
    this.backgroundTasksThreadPool = backgroundTasksThreadPool;
    this.grpcDecorator = checkNotNull(grpcDecorator, "grpcDecorator cannot be null");
  }

  private ManagedChannel buildManagedChannel() {
    String target =
        isNullOrEmpty(readOptions.getGrpcServerAddress())
            ? DEFAULT_GCS_GRPC_SERVER_ADDRESS
            : readOptions.getGrpcServerAddress();
    ManagedChannel channel =
        grpcDecorator.createChannelBuilder(target).enableRetry().userAgent(userAgent).build();
    return channel;
  }

  public static boolean isStubBroken(Status.Code statusCode) {
    return STUB_BROKEN_ERROR_CODES.contains(statusCode);
  }

  public StorageBlockingStub newBlockingStub() {
    StorageBlockingStub stub = StorageGrpc.newBlockingStub(getManagedChannel());
    return (StorageBlockingStub) grpcDecorator.applyCallOption(stub);
  }

  public StorageStub newAsyncStub() {
    StorageStub stub =
        StorageGrpc.newStub(getManagedChannel()).withExecutor(backgroundTasksThreadPool);
    return (StorageStub) grpcDecorator.applyCallOption(stub);
  }

  private synchronized ManagedChannel getManagedChannel() {
    if (channel == null) {
      channel = buildManagedChannel();
    }
    return channel;
  }

  public void shutdown() {
    if (channel != null) channel.shutdown();
  }

  interface GrpcDecorator {
    ManagedChannelBuilder<?> createChannelBuilder(String target);

    AbstractStub<?> applyCallOption(AbstractStub<?> stub);
  }

  static class CloudPathGrpcDecorator implements GrpcDecorator {
    private final Credentials credentials;

    CloudPathGrpcDecorator(Credentials credentials) {
      this.credentials = credentials;
    }

    public ManagedChannelBuilder<?> createChannelBuilder(String target) {
      return ManagedChannelBuilder.forTarget(target);
    }

    public AbstractStub<?> applyCallOption(AbstractStub<?> stub) {
      return stub.withCallCredentials(MoreCallCredentials.from(credentials));
    }
  }

  static class DirectPathGrpcDecorator implements GrpcDecorator {
    private final GoogleCloudStorageReadOptions readOptions;

    DirectPathGrpcDecorator(GoogleCloudStorageReadOptions readOptions) {
      this.readOptions = readOptions;
    }

    public ManagedChannelBuilder<?> createChannelBuilder(String target) {
      return GoogleDefaultChannelBuilder.forTarget(target)
          .defaultServiceConfig(getGrpcServiceConfig());
    }

    public AbstractStub<?> applyCallOption(AbstractStub<?> stub) {
      return stub;
    }

    private Map<String, Object> getGrpcServiceConfig() {
      Map<String, Object> name = ImmutableMap.of("service", "google.storage.v1.Storage");

      Map<String, Object> retryPolicy =
          ImmutableMap.<String, Object>builder()
              .put("maxAttempts", GRPC_MAX_RETRY_ATTEMPTS)
              .put(
                  "initialBackoff",
                  Durations.toString(
                      Durations.fromMillis(readOptions.getBackoffInitialIntervalMillis())))
              .put(
                  "maxBackoff",
                  Durations.toString(
                      Durations.fromMillis(readOptions.getBackoffMaxIntervalMillis())))
              .put("backoffMultiplier", readOptions.getBackoffMultiplier())
              .put("retryableStatusCodes", ImmutableList.of("UNAVAILABLE", "RESOURCE_EXHAUSTED"))
              .build();

      Map<String, Object> methodConfig =
          ImmutableMap.of("name", ImmutableList.of(name), "retryPolicy", retryPolicy);

      Map<String, Object> childLbStrategy = ImmutableMap.of("round_robin", ImmutableMap.of());

      Map<String, Object> childPolicy =
          ImmutableMap.of("childPolicy", ImmutableList.of(childLbStrategy));

      Map<String, Object> grpcLbPolicy = ImmutableMap.of("grpclb", childPolicy);

      return ImmutableMap.of(
          "methodConfig",
          ImmutableList.of(methodConfig),
          "loadBalancingConfig",
          ImmutableList.of(grpcLbPolicy));
    }
  }

  public static StorageStubProvider newInstance(
      GoogleCloudStorageOptions options,
      ExecutorService backgroundTasksThreadPool,
      Credential credential) {
    boolean useDirectpath =
        options.isDirectPathPreferred()
            && credential != null
            && java.util.Objects.equals(
                credential.getTokenServerEncodedUrl(), ComputeCredential.TOKEN_SERVER_ENCODED_URL);
    return new StorageStubProvider(
        options,
        backgroundTasksThreadPool,
        useDirectpath
            ? new DirectPathGrpcDecorator(options.getReadChannelOptions())
            : new CloudPathGrpcDecorator(new CredentialAdapter(credential)));
  }

  public static StorageStubProvider newInstance(
      GoogleCloudStorageOptions options,
      ExecutorService backgroundTasksThreadPool,
      Credentials credentials) {
    boolean useDirectpath =
        options.isDirectPathPreferred() && credentials instanceof ComputeEngineCredentials;
    return new StorageStubProvider(
        options,
        backgroundTasksThreadPool,
        useDirectpath
            ? new DirectPathGrpcDecorator(options.getReadChannelOptions())
            : new CloudPathGrpcDecorator(credentials));
  }
}

@@ -663,51 +663,8 @@
 et;%0A
-import com.google.protobuf.util.Durations;%0A
 impo
@@ -4905,939 +4905,8 @@
 ) %7B%0A
-      Map%3CString, Object%3E name = ImmutableMap.of(%22service%22, %22google.storage.v1.Storage%22);%0A%0A      Map%3CString, Object%3E retryPolicy =%0A          ImmutableMap.%3CString, Object%3Ebuilder()%0A              .put(%22maxAttempts%22, GRPC_MAX_RETRY_ATTEMPTS)%0A              .put(%0A                  %22initialBackoff%22,%0A                  Durations.toString(%0A                      Durations.fromMillis(readOptions.getBackoffInitialIntervalMillis())))%0A              .put(%0A                  %22maxBackoff%22,%0A                  Durations.toString(%0A                      Durations.fromMillis(readOptions.getBackoffMaxIntervalMillis())))%0A              .put(%22backoffMultiplier%22, readOptions.getBackoffMultiplier())%0A              .put(%22retryableStatusCodes%22, ImmutableList.of(%22UNAVAILABLE%22, %22RESOURCE_EXHAUSTED%22))%0A              .build();%0A%0A      Map%3CString, Object%3E methodConfig =%0A          ImmutableMap.of(%22name%22, ImmutableList.of(name), %22retryPolicy%22, retryPolicy);%0A%0A
     
@@ -4992,25 +4992,24 @@
 eMap.of());%0A
-%0A
       Map%3CSt
@@ -5113,17 +5113,16 @@
 tegy));%0A
-%0A
       Ma
@@ -5194,17 +5194,16 @@
 olicy);%0A
-%0A
       re
@@ -5227,87 +5227,8 @@
 .of(
-%0A          %22methodConfig%22,%0A          ImmutableList.of(methodConfig),%0A          
 %22loa
@@ -5245,26 +5245,16 @@
 Config%22,
-%0A         
  Immutab

Remove grpclb retry config (#654)