/*
 * Copyright 2018 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
 *
 *     https://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.bigtable.admin.v2.stub;

import com.google.api.core.InternalApi;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.rpc.ClientContext;
import com.google.api.gax.rpc.UnaryCallable;
import com.google.bigtable.admin.v2.TableName;
import java.io.IOException;
import org.threeten.bp.Duration;

/**
 * Extension of the autogenerated {@link GrpcBigtableTableAdminStub}. It acts as a decorator to add
 * enhanced abilities to the autogenerated stub.
 *
 * <p>This class is considered an internal implementation detail and not meant to be used by
 * applications.
 */
@InternalApi
public class EnhancedBigtableTableAdminStub extends GrpcBigtableTableAdminStub {
  private final BigtableTableAdminStubSettings settings;
  private final ClientContext clientContext;

  private final AwaitReplicationCallable awaitReplicationCallable;

  public static EnhancedBigtableTableAdminStub createEnhanced(
      BigtableTableAdminStubSettings settings)
      throws IOException {
    return new EnhancedBigtableTableAdminStub(settings, ClientContext.create(settings));
  }

  private EnhancedBigtableTableAdminStub(BigtableTableAdminStubSettings settings,
      ClientContext clientContext) throws IOException {
    super(settings, clientContext);

    this.settings = settings;
    this.clientContext = clientContext;
    this.awaitReplicationCallable = createAwaitReplicationCallable();
  }

  private AwaitReplicationCallable createAwaitReplicationCallable() {
    // TODO(igorbernstein2): expose polling settings
    RetrySettings pollingSettings = RetrySettings.newBuilder()
        // use overall timeout from checkConsistencyCallable
        // NOTE: The overall timeout might exceed this value due to underlying retries
        .setTotalTimeout(settings.checkConsistencySettings().getRetrySettings().getTotalTimeout())
        // Use constant polling with jitter
        .setInitialRetryDelay(Duration.ofSeconds(10))
        .setRetryDelayMultiplier(1.0)
        .setMaxRetryDelay(Duration.ofSeconds(10))
        .setJittered(true)
        // These rpc timeouts are ignored, instead the rpc timeouts defined for
        // generateConsistencyToken and checkConsistency callables will be used.
        .setInitialRpcTimeout(Duration.ZERO)
        .setMaxRpcTimeout(Duration.ZERO)
        .setRpcTimeoutMultiplier(1.0)
        .build();

    return AwaitReplicationCallable.create(
        generateConsistencyTokenCallable(),
        checkConsistencyCallable(),
        clientContext,
        pollingSettings
    );
  }

  public UnaryCallable<TableName, Void> awaitReplicationCallable() {
    return awaitReplicationCallable;
  }
}

