/*
 * Copyright 2016 Spotify AB.
 *
 * 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.spotify.scio.bigtable;

import com.google.bigtable.admin.v2.Cluster;
import com.google.bigtable.admin.v2.ListClustersRequest;
import com.google.bigtable.admin.v2.ListClustersResponse;
import com.google.cloud.bigtable.config.BigtableOptions;
import com.google.cloud.bigtable.grpc.BigtableInstanceClient;
import com.google.cloud.bigtable.grpc.BigtableInstanceGrpcClient;
import com.google.cloud.bigtable.grpc.io.ChannelPool;
import org.joda.time.Duration;

import java.io.IOException;

/**
 * Utilities to deal with Bigtable.
 */
public final class BigtableUtil {

  private BigtableUtil() { }

  /**
   * Updates all clusters within the specified Bigtable instance to a specified number of nodes.
   * Useful for increasing the number of nodes at the beginning of a job and decreasing it at
   * the end to lower costs yet still get high throughput during bulk ingests/dumps.
   *
   * @param bigtableOptions Bigtable Options
   * @param numberOfNodes New number of nodes in the cluster
   * @param sleepDuration How long to sleep after updating the number of nodes. Google recommends
   *                      at least 20 minutes before the new nodes are fully functional
   * @throws IOException If setting up channel pool fails
   * @throws InterruptedException If sleep fails
   */
  public static void updateNumberOfBigtableNodes(
      final BigtableOptions bigtableOptions,
      final int numberOfNodes,
      final Duration sleepDuration
  ) throws IOException, InterruptedException {
    final ChannelPool channelPool =
        ChannelPoolCreator.createPool(bigtableOptions.getInstanceAdminHost());

    try {
      final BigtableInstanceClient bigtableInstanceClient = new BigtableInstanceGrpcClient(channelPool);

      final String instanceName = bigtableOptions.getInstanceName().toString();

      // Fetch clusters in Bigtable instance
      final ListClustersRequest clustersRequest =
          ListClustersRequest.newBuilder().setParent(instanceName).build();
      final ListClustersResponse clustersResponse =
          bigtableInstanceClient.listCluster(clustersRequest);

      // For each cluster update the number of nodes
      for (Cluster cluster : clustersResponse.getClustersList()) {
        final Cluster updatedCluster = Cluster.newBuilder()
            .setName(cluster.getName())
            .setServeNodes(numberOfNodes)
            .build();
        bigtableInstanceClient.updateCluster(updatedCluster);
      }

      // Wait for the new nodes to be provisioned
      Thread.sleep(sleepDuration.getMillis());
    } finally {
      channelPool.shutdownNow();
    }
  }
}

@@ -1063,145 +1063,662 @@
 on;%0A
-%0Aimport java.io.IOException;%0A%0A/**%0A * Utilities to deal with Bigtable.%0A */%0Apublic final class BigtableUtil %7B%0A%0A  private BigtableUtil() %7B %7D
+import org.joda.time.format.PeriodFormatter;%0Aimport org.joda.time.format.PeriodFormatterBuilder;%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A%0Aimport java.io.IOException;%0A%0A/**%0A * Utilities to deal with Bigtable.%0A */%0Apublic final class BigtableUtil %7B%0A%0A  private BigtableUtil() %7B %7D%0A%0A  private static final Logger LOG = LoggerFactory.getLogger(BigtableUtil.class);%0A%0A  private static final PeriodFormatter formatter = new PeriodFormatterBuilder()%0A     .appendDays()%0A     .appendSuffix(%22d%22)%0A     .appendHours()%0A     .appendSuffix(%22h%22)%0A     .appendMinutes()%0A     .appendSuffix(%22m%22)%0A     .appendSeconds()%0A     .appendSuffix(%22s%22)%0A     .toFormatter();
 %0A%0A  
@@ -3495,24 +3495,125 @@
   .build();%0A
+        LOG.info(%22Updating number of nodes to %7B%7D for cluster %7B%7D%22, numberOfNodes, cluster.getName());%0A
         bigt
@@ -3721,16 +3721,154 @@
 isioned%0A
+      if (sleepDuration.getMillis() %3E 0) %7B%0A        LOG.info(%22Sleeping for %7B%7D after update%22, formatter.print(sleepDuration.toPeriod()));%0A  
       Th
@@ -3902,24 +3902,32 @@
 tMillis());%0A
+      %7D%0A
     %7D finall
@@ -3973,10 +3973,11 @@
   %7D%0A  %7D%0A
+%0A
 %7D%0A

add logging for updateNumberOfBigtableNodes