/*
    * Copyright 2021 TiKV Project Authors.
    *
    * 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 org.tikv.raw;
  
  import com.google.protobuf.ByteString;
  import io.prometheus.client.Counter;
  import io.prometheus.client.Histogram;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.tikv.common.TiSession;
  import org.tikv.common.exception.CircuitBreakerOpenException;
  import org.tikv.common.util.HistogramUtils;
  import org.tikv.common.util.Pair;
  import org.tikv.common.util.ScanOption;
  import org.tikv.kvproto.Kvrpcpb;
  import org.tikv.service.failsafe.CircuitBreaker;
  
  public class SmartRawKVClient implements RawKVClientBase {
    private static final Logger logger = LoggerFactory.getLogger(SmartRawKVClient.class);
  
    private static final Histogram REQUEST_LATENCY =
        HistogramUtils.buildDuration()
            .name("client_java_smart_raw_requests_latency")
            .help("client smart raw request latency.")
            .labelNames("type", "cluster")
            .register();
  
    private static final Counter REQUEST_SUCCESS =
        Counter.build()
            .name("client_java_smart_raw_requests_success")
            .help("client smart raw request success.")
            .labelNames("type", "cluster")
            .register();
  
    private static final Counter REQUEST_FAILURE =
        Counter.build()
            .name("client_java_smart_raw_requests_failure")
            .help("client smart raw request failure.")
            .labelNames("type", "cluster")
            .register();
  
    private static final Counter CIRCUIT_BREAKER_OPENED =
        Counter.build()
            .name("client_java_smart_raw_circuit_breaker_opened")
            .help("client smart raw circuit breaker opened.")
            .labelNames("type", "cluster")
            .register();
  
    private final RawKVClientBase client;
    private final CircuitBreaker circuitBreaker;
  
    public SmartRawKVClient(RawKVClientBase client, CircuitBreaker breaker) {
      this.client = client;
      this.circuitBreaker = breaker;
    }
  
    @Override
    public void put(ByteString key, ByteString value) {
      callWithCircuitBreaker("put", () -> client.put(key, value));
    }
  
    @Override
    public void put(ByteString key, ByteString value, long ttl) {
      callWithCircuitBreaker("put", () -> client.put(key, value, ttl));
    }
  
    @Override
    public Optional<ByteString> putIfAbsent(ByteString key, ByteString value) {
      return callWithCircuitBreaker("putIfAbsent", () -> client.putIfAbsent(key, value));
    }
  
    @Override
    public Optional<ByteString> putIfAbsent(ByteString key, ByteString value, long ttl) {
      return callWithCircuitBreaker("putIfAbsent", () -> client.putIfAbsent(key, value, ttl));
    }
  
    @Override
    public void compareAndSet(ByteString key, Optional<ByteString> prevValue, ByteString value) {
      callWithCircuitBreaker("compareAndSet", () -> client.compareAndSet(key, prevValue, value));
    }
  
   @Override
   public void compareAndSet(
       ByteString key, Optional<ByteString> prevValue, ByteString value, long ttl) {
     callWithCircuitBreaker("compareAndSet", () -> client.compareAndSet(key, prevValue, value, ttl));
   }
 
   @Override
   public void batchPut(Map<ByteString, ByteString> kvPairs) {
     callWithCircuitBreaker("batchPut", () -> client.batchPut(kvPairs));
   }
 
   @Override
   public void batchPut(Map<ByteString, ByteString> kvPairs, long ttl) {
     callWithCircuitBreaker("batchPut", () -> client.batchPut(kvPairs, ttl));
   }
 
   @Override
   public Optional<ByteString> get(ByteString key) {
     return callWithCircuitBreaker("get", () -> client.get(key));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> batchGet(List<ByteString> keys) {
     return callWithCircuitBreaker("batchGet", () -> client.batchGet(keys));
   }
 
   @Override
   public void batchDelete(List<ByteString> keys) {
     callWithCircuitBreaker("batchDelete", () -> client.batchDelete(keys));
   }
 
   @Override
   public Optional<Long> getKeyTTL(ByteString key) {
     return callWithCircuitBreaker("getKeyTTL", () -> client.getKeyTTL(key));
   }
 
   @Override
   public List<List<ByteString>> batchScanKeys(
       List<Pair<ByteString, ByteString>> ranges, int eachLimit) {
     return callWithCircuitBreaker("batchScanKeys", () -> client.batchScanKeys(ranges, eachLimit));
   }
 
   @Override
   public List<List<Kvrpcpb.KvPair>> batchScan(List<ScanOption> ranges) {
     return callWithCircuitBreaker("batchScan", () -> client.batchScan(ranges));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, int limit) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, limit));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(
       ByteString startKey, ByteString endKey, int limit, boolean keyOnly) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, limit, keyOnly));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, limit));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit, boolean keyOnly) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, limit, keyOnly));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey, boolean keyOnly) {
     return callWithCircuitBreaker("scan", () -> client.scan(startKey, endKey, keyOnly));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, int limit, boolean keyOnly) {
     return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey, limit, keyOnly));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey) {
     return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey));
   }
 
   @Override
   public List<Kvrpcpb.KvPair> scanPrefix(ByteString prefixKey, boolean keyOnly) {
     return callWithCircuitBreaker("scanPrefix", () -> client.scanPrefix(prefixKey, keyOnly));
   }
 
   @Override
   public void delete(ByteString key) {
     callWithCircuitBreaker("delete", () -> client.delete(key));
   }
 
   @Override
   public void deleteRange(ByteString startKey, ByteString endKey) {
     callWithCircuitBreaker("deleteRange", () -> client.deleteRange(startKey, endKey));
   }
 
   @Override
   public void deletePrefix(ByteString key) {
     callWithCircuitBreaker("deletePrefix", () -> client.deletePrefix(key));
   }
 
   @Override
   public TiSession getSession() {
     return client.getSession();
   }
 
   <T> T callWithCircuitBreaker(String funcName, Function1<T> func) {
     String[] labels =
         new String[] {funcName, client.getSession().getPDClient().getClusterId().toString()};
 
     Histogram.Timer requestTimer = REQUEST_LATENCY.labels(labels).startTimer();
     try {
       T result = callWithCircuitBreaker0(funcName, func);
       REQUEST_SUCCESS.labels(labels).inc();
       return result;
     } catch (Exception e) {
       REQUEST_FAILURE.labels(labels).inc();
       throw e;
     } finally {
       requestTimer.observeDuration();
     }
   }
 
   private <T> T callWithCircuitBreaker0(String funcName, Function1<T> func) {
     if (circuitBreaker.allowRequest()) {
       try {
         T result = func.apply();
         circuitBreaker.getMetrics().recordSuccess();
         return result;
       } catch (Exception e) {
         circuitBreaker.getMetrics().recordFailure();
         throw e;
       }
     } else if (circuitBreaker.attemptExecution()) {
       logger.debug("attemptExecution");
       try {
         T result = func.apply();
         circuitBreaker.getMetrics().recordSuccess();
         circuitBreaker.recordAttemptSuccess();
         logger.debug("markSuccess");
         return result;
       } catch (Exception e) {
         circuitBreaker.getMetrics().recordFailure();
         circuitBreaker.recordAttemptFailure();
         logger.debug("markNonSuccess");
         throw e;
       }
     } else {
       logger.debug("Circuit Breaker Opened");
       CIRCUIT_BREAKER_OPENED
           .labels(funcName, client.getSession().getPDClient().getClusterId().toString())
           .inc();
       throw new CircuitBreakerOpenException();
     }
   }
 
   private void callWithCircuitBreaker(String funcName, Function0 func) {
     callWithCircuitBreaker(
         funcName,
         (Function1<Void>)
             () -> {
               func.apply();
               return null;
             });
   }
 
   @Override
   public void close() throws Exception {
     client.close();
   }
 
   public interface Function1<T> {
     T apply();
   }
 
   public interface Function0 {
     void apply();
   }
 }