/*
 * Copyright (c) 2015 Spotify AB.
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.heroic.metric.bigtable;

import com.google.appengine.repackaged.com.google.common.collect.ImmutableList;
import com.google.cloud.bigtable.config.BigtableOptions;
import com.google.cloud.bigtable.config.CredentialOptions;
import com.google.cloud.bigtable.grpc.BigtableSession;
import com.spotify.heroic.metric.bigtable.api.BigtableDataClient;
import com.spotify.heroic.metric.bigtable.api.BigtableDataClientImpl;
import com.spotify.heroic.metric.bigtable.api.BigtableMutator;
import com.spotify.heroic.metric.bigtable.api.BigtableMutatorImpl;
import com.spotify.heroic.metric.bigtable.api.BigtableTableAdminClient;
import com.spotify.heroic.metric.bigtable.api.BigtableTableTableAdminClientImpl;
import eu.toolchain.async.AsyncFramework;
import eu.toolchain.async.AsyncFuture;
import lombok.RequiredArgsConstructor;
import lombok.ToString;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;

@ToString(of = {"project", "instance", "credentials"})
@RequiredArgsConstructor
public class BigtableConnectionBuilder implements Callable<BigtableConnection> {
    private static final String USER_AGENT = "heroic";

    private final String project;
    private final String instance;

    private final CredentialsBuilder credentials;

    private final AsyncFramework async;
    private final ExecutorService executorService;

    private final boolean disableBulkMutations;
    private final int flushIntervalSeconds;

    @Override
    public BigtableConnection call() throws Exception {
        final CredentialOptions credentials = this.credentials.build();

        final BigtableOptions options = new BigtableOptions.Builder()
            .setProjectId(project)
            .setInstanceId(instance)
            .setUserAgent(USER_AGENT)
            .setDataChannelCount(64)
            .setCredentialOptions(credentials)
            .build();

        final BigtableSession session = new BigtableSession(options, executorService);

        final BigtableTableAdminClient adminClient =
            new BigtableTableTableAdminClientImpl(async, session.getTableAdminClient(), project,
              instance);

        final BigtableMutator mutator =
            new BigtableMutatorImpl(async, session, disableBulkMutations, flushIntervalSeconds);

        final BigtableDataClient client =
            new BigtableDataClientImpl(async, session, mutator, project, instance);

        return new GrpcBigtableConnection(
            async, project, instance, session, mutator, adminClient, client);
    }

    @RequiredArgsConstructor
    @ToString(of = {"project", "instance"})
    public static class GrpcBigtableConnection implements BigtableConnection {
        private final AsyncFramework async;
        private final String project;
        private final String instance;

        final BigtableSession session;
        final BigtableMutator mutator;
        final BigtableTableAdminClient tableAdminClient;
        final BigtableDataClient dataClient;

        @Override
        public BigtableTableAdminClient tableAdminClient() {
            return tableAdminClient;
        }

        @Override
        public BigtableDataClient dataClient() {
            return dataClient;
        }

        @Override
        public AsyncFuture<Void> close() {
            final AsyncFuture<Void> closeSession = async.call(() -> {
                session.close();
                return null;
            });

            return async.collectAndDiscard(ImmutableList.of(mutator.close(), closeSession));
        }
    }
}

@@ -1117,28 +1117,130 @@
 ble.
-grpc.BigtableSession
+config.RetryOptions;%0Aimport com.google.cloud.bigtable.grpc.BigtableSession;%0Aimport com.spotify.heroic.bigtable.grpc.Status
 ;%0Aim
@@ -2543,32 +2543,229 @@
 tials.build();%0A%0A
+        final RetryOptions retryOptions = new RetryOptions.Builder()%0A            .addStatusToRetryOn(Status.Code.UNKNOWN)%0A            .setAllowRetriesWithoutTimestamp(true)%0A            .build();%0A%0A%0A
         final Bi
@@ -3012,16 +3012,59 @@
 ntials)%0A
+            .setRetryOptions(retryOptions)%0A
         

enable retries without timestamp and also on UNKNOWN