package io.shardingjdbc.orchestration.reg.etcd.internal;

import com.github.rholder.retry.RetryException;
import com.github.rholder.retry.RetryerBuilder;
import com.github.rholder.retry.StopStrategies;
import com.github.rholder.retry.WaitStrategies;
import com.google.common.base.Optional;
import com.google.common.collect.Lists;
import com.google.protobuf.ByteString;
import etcdserverpb.KVGrpc;
import etcdserverpb.LeaseGrpc;
import etcdserverpb.LeaseGrpc.LeaseFutureStub;
import etcdserverpb.Rpc.LeaseGrantRequest;
import etcdserverpb.Rpc.PutRequest;
import etcdserverpb.Rpc.PutResponse;
import etcdserverpb.Rpc.RangeRequest;
import etcdserverpb.Rpc.RangeResponse;
import etcdserverpb.Rpc.WatchCreateRequest;
import etcdserverpb.Rpc.WatchRequest;
import etcdserverpb.WatchGrpc;
import io.grpc.Channel;
import io.grpc.stub.StreamObserver;
import io.shardingjdbc.orchestration.reg.exception.RegException;
import io.shardingjdbc.orchestration.reg.exception.RegExceptionHandler;
import mvccpb.Kv.KeyValue;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Etcd client implementation client.
 *
 * @author junxiong
 */
public class EtcdClientImpl implements EtcdClient, AutoCloseable {
    
    private long timeoutMills = 500L;
    
    private long retryMills = 200L;
    
    private int retryTimes = 2;
    
    private LeaseFutureStub leaseStub;
    
    private KVGrpc.KVFutureStub kvStub;
    
    private WatchGrpc.WatchStub watchStub;
    
    private AtomicBoolean closed = new AtomicBoolean(false);
    
    /**
     * construct a etcd client from grpc channel.
     *
     * @param channel Channel
     */
    EtcdClientImpl(final Channel channel, final long timeoutMills, final long retryMills, final int retryTimes) {
        this.timeoutMills = timeoutMills;
        this.retryMills = retryMills;
        this.retryTimes = retryTimes;
        leaseStub = LeaseGrpc.newFutureStub(channel);
        kvStub = KVGrpc.newFutureStub(channel);
        watchStub = WatchGrpc.newStub(channel);
    }
    
    @Override
    public Optional<String> get(final String key) {
        final RangeRequest request = RangeRequest.newBuilder().setKey(ByteString.copyFromUtf8(key)).build();
        return retry(new Callable<String>() {
            
            @Override
            public String call() throws Exception {
                final RangeResponse rangeResponse = kvStub.range(request).get(timeoutMills, TimeUnit.MILLISECONDS);
                return rangeResponse.getKvsCount() > 0
                        ? rangeResponse.getKvs(0).getValue().toStringUtf8() : null;
            }
        });
    }
    
    @Override
    public Optional<List<String>> list(final String dir) {
        final RangeRequest request = RangeRequest.newBuilder()
                .setKey(ByteString.copyFromUtf8(dir))
                .setRangeEnd(getRangeEnd(dir))
                .build();
        return retry(new Callable<List<String>>() {
            @Override
            public List<String> call() throws Exception {
                RangeResponse rangeResponse = kvStub.range(request).get(timeoutMills, TimeUnit.MILLISECONDS);
                final List<String> keys = Lists.newArrayList();
                for (KeyValue keyValue : rangeResponse.getKvsList()) {
                    keys.add(keyValue.getKey().toStringUtf8());
                }
                return keys;
            }
        });
    }
    
    @Override
    public Optional<String> put(final String key, final String value) {
        final PutRequest request = PutRequest.newBuilder()
                .setPrevKv(true)
                .setKey(ByteString.copyFromUtf8(key))
                .setValue(ByteString.copyFromUtf8(value))
                .build();
        return retry(new Callable<String>() {
            @Override
            public String call() throws Exception {
                PutResponse putResponse = kvStub.put(request).get(timeoutMills, TimeUnit.MILLISECONDS);
                return putResponse.getPrevKv().getValue().toStringUtf8();
            }
        });
    }
    
    @Override
    public Optional<String> put(final String key, final String value, final long ttl) {
        final Optional<Long> leaseId = lease(ttl);
        if (!leaseId.isPresent()) {
            throw new RegException("Unable to set up heat beat for key %s", key);
        }
        final PutRequest request = PutRequest.newBuilder()
                .setPrevKv(true)
                .setLease(leaseId.get())
                .setKey(ByteString.copyFromUtf8(key))
                .setValue(ByteString.copyFromUtf8(value))
                .build();
        return retry(new Callable<String>() {
            
            @Override
            public String call() throws Exception {
                PutResponse putResponse = kvStub.put(request).get(timeoutMills, TimeUnit.MILLISECONDS);
                return putResponse.getPrevKv().getValue().toStringUtf8();
            }
        });
    }
    
    private Optional<Long> lease(final long ttl) {
        final LeaseGrantRequest request = LeaseGrantRequest.newBuilder().setTTL(ttl).build();
        leaseStub.leaseGrant(request);
        return retry(new Callable<Long>() {
            
            @Override
            public Long call() throws Exception {
                return leaseStub.leaseGrant(request).get(timeoutMills, TimeUnit.MILLISECONDS).getID();
            }
        });
    }
    
    @Override
    public Optional<Watcher> watch(final String key) {
        WatchCreateRequest createWatchRequest = WatchCreateRequest.newBuilder().setKey(ByteString.copyFromUtf8(key)).setRangeEnd(getRangeEnd(key)).build();
        final WatchRequest watchRequest = WatchRequest.newBuilder().setCreateRequest(createWatchRequest).build();
        return retry(new Callable<Watcher>() {
            
            @Override
            public Watcher call() throws Exception {
                Watcher watcher = new Watcher();
                StreamObserver<WatchRequest> requestStream = watchStub.watch(new WatchStreamObserver(watcher));
                requestStream.onNext(watchRequest);
                return watcher;
            }
        });
    }
    
    @Override
    public void close() throws Exception {
        closed.compareAndSet(true, false);
    }
    
    private <T> Optional<T> retry(final Callable<T> command) {
        try {
            return Optional.fromNullable(RetryerBuilder.<T>newBuilder()
                    .retryIfExceptionOfType(TimeoutException.class)
                    .retryIfExceptionOfType(ExecutionException.class)
                    .retryIfExceptionOfType(InterruptedException.class)
                    .withWaitStrategy(WaitStrategies.fixedWait(retryMills, TimeUnit.MILLISECONDS))
                    .withStopStrategy(StopStrategies.stopAfterAttempt(retryTimes))
                    .build().call(command));
        } catch (final ExecutionException | RetryException ex) {
            RegExceptionHandler.handleException(ex);
            return Optional.absent();
        }
    }
    
    private ByteString getRangeEnd(final String key) {
        byte[] noPrefix = {0};
        byte[] endKey = key.getBytes().clone();
        for (int i = endKey.length - 1; i >= 0; i--) {
            if (endKey[i] < 0xff) {
                endKey[i] = (byte) (endKey[i] + 1);
                return ByteString.copyFrom(Arrays.copyOf(endKey, i + 1));
            }
        }
        return ByteString.copyFrom(noPrefix);
    }
}

@@ -1220,58 +1220,8 @@
 ion;
-%0Aimport java.util.concurrent.atomic.AtomicBoolean;
 %0A%0A/*
@@ -1340,23 +1340,8 @@
 ient
-, AutoCloseable
  %7B%0A 
@@ -1606,74 +1606,8 @@
    %0A
-    private AtomicBoolean closed = new AtomicBoolean(false);%0A    %0A
     
@@ -6250,119 +6250,8 @@
    %0A
-    @Override%0A    public void close() throws Exception %7B%0A        closed.compareAndSet(true, false);%0A    %7D%0A    %0A
     

remove useless close