/*
 * 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.common.region;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.util.concurrent.ListenableFuture;
import com.google.protobuf.ByteString;
import io.grpc.ManagedChannel;
import io.grpc.Metadata;
import io.grpc.stub.MetadataUtils;
import io.prometheus.client.Histogram;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.tikv.common.AbstractGRPCClient;
import org.tikv.common.TiConfiguration;
import org.tikv.common.apiversion.RequestKeyCodec;
import org.tikv.common.exception.GrpcException;
import org.tikv.common.log.SlowLog;
import org.tikv.common.log.SlowLogSpan;
import org.tikv.common.util.BackOffer;
import org.tikv.common.util.ChannelFactory;
import org.tikv.common.util.HistogramUtils;
import org.tikv.kvproto.Kvrpcpb;
import org.tikv.kvproto.Metapb;
import org.tikv.kvproto.TikvGrpc;
import org.tikv.kvproto.Tracepb;

public abstract class AbstractRegionStoreClient
    extends AbstractGRPCClient<TikvGrpc.TikvBlockingStub, TikvGrpc.TikvFutureStub>
    implements RegionErrorReceiver {

  private static final Logger logger = LoggerFactory.getLogger(AbstractRegionStoreClient.class);

  public static final Histogram SEEK_LEADER_STORE_DURATION =
      HistogramUtils.buildDuration()
          .name("client_java_seek_leader_store_duration")
          .help("seek leader store duration.")
          .labelNames("cluster")
          .register();

  public static final Histogram SEEK_PROXY_STORE_DURATION =
      HistogramUtils.buildDuration()
          .name("client_java_seek_proxy_store_duration")
          .help("seek proxy store duration.")
          .labelNames("cluster")
          .register();

  protected final RegionManager regionManager;
  protected final RequestKeyCodec codec;
  protected TiRegion region;
  protected TiStore store;

  protected AbstractRegionStoreClient(
      TiConfiguration conf,
      TiRegion region,
      TiStore store,
      ChannelFactory channelFactory,
      TikvGrpc.TikvBlockingStub blockingStub,
      TikvGrpc.TikvFutureStub asyncStub,
      RegionManager regionManager) {
    super(conf, channelFactory, blockingStub, asyncStub);
    checkNotNull(region, "Region is empty");
    checkNotNull(region.getLeader(), "Leader Peer is null");
    checkArgument(region.getLeader() != null, "Leader Peer is null");
    this.region = region;
    this.regionManager = regionManager;
    this.store = store;
    this.codec = regionManager.getPDClient().getCodec();
    if (this.store.getProxyStore() != null) {
      this.timeout = conf.getForwardTimeout();
    }
  }

  @Override
  public TiRegion getRegion() {
    return region;
  }

  @Override
  protected TikvGrpc.TikvBlockingStub getBlockingStub() {
    return blockingStub.withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
  }

  @Override
  protected TikvGrpc.TikvFutureStub getAsyncStub() {
    return asyncStub.withDeadlineAfter(getTimeout(), TimeUnit.MILLISECONDS);
  }

  @Override
  public void close() throws GrpcException {}

  /**
   * onNotLeader deals with NotLeaderError and returns whether re-splitting key range is needed
   *
   * @param newRegion the new region presented by NotLeader Error
   * @return false when re-split is needed.
   */
  @Override
  public boolean onNotLeader(TiRegion newRegion, BackOffer backOffer) {
    if (logger.isDebugEnabled()) {
      logger.debug(region + ", new leader = " + newRegion.getLeader().getStoreId());
    }
    // When switch leader fails or the region changed its region epoch,
    // it would be necessary to re-split task's key range for new region.
    if (!region.getRegionEpoch().equals(newRegion.getRegionEpoch())) {
      return false;
    }

    // If we try one peer but find the leader has not changed, we do not need to try other peers.
    if (region.getLeader().getStoreId() == newRegion.getLeader().getStoreId()) {
      store = null;
    }
    region = newRegion;
    store = regionManager.getStoreById(region.getLeader().getStoreId(), backOffer);
    updateClientStub();
    return true;
  }

  @Override
  public boolean onStoreUnreachable(BackOffer backOffer) {
    if (!store.isValid()) {
      logger.warn(String.format("store [%d] has been invalid", store.getId()));
      store = regionManager.getStoreById(store.getId(), backOffer);
      updateClientStub();
      return true;
    }

    // seek an available leader store to send request
    backOffer.checkTimeout();
    Boolean result = seekLeaderStore(backOffer);
    if (result != null) {
      return result;
    }
    if (conf.getEnableGrpcForward()) {
      // seek an available proxy store to forward request
      backOffer.checkTimeout();
      return seekProxyStore(backOffer);
    }
    return false;
  }

  private Kvrpcpb.Context addTraceId(Kvrpcpb.Context context, SlowLog slowLog) {
    if (slowLog.getThresholdMS() < 0) {
      // disable tikv tracing
      return context;
    }
    long traceId = slowLog.getTraceId();
    return Kvrpcpb.Context.newBuilder(context)
        .setTraceContext(
            Tracepb.TraceContext.newBuilder()
                .setDurationThresholdMs(
                    (int) (slowLog.getThresholdMS() * conf.getRawKVServerSlowLogFactor()))
                .addRemoteParentSpans(Tracepb.RemoteParentSpan.newBuilder().setTraceId(traceId)))
        .build();
  }

  protected Kvrpcpb.Context makeContext(TiStoreType storeType, SlowLog slowLog) {
    Kvrpcpb.Context context = region.getReplicaContext(java.util.Collections.emptySet(), storeType);
    return addTraceId(context, slowLog);
  }

  protected Kvrpcpb.Context makeContext(
      Set<Long> resolvedLocks, TiStoreType storeType, SlowLog slowLog) {
    Kvrpcpb.Context context = region.getReplicaContext(resolvedLocks, storeType);
    return addTraceId(context, slowLog);
  }

  protected Kvrpcpb.Context makeContext() {
    return region.getLeaderContext();
  }

  protected Kvrpcpb.Context makeContext(Metapb.Peer peer) {
    return region.getReplicaContext(peer);
  }

  private void updateClientStub() {
    String addressStr = store.getStore().getAddress();
    long deadline = timeout;
    if (store.getProxyStore() != null) {
      addressStr = store.getProxyStore().getAddress();
      deadline = conf.getForwardTimeout();
    }
    ManagedChannel channel =
        channelFactory.getChannel(addressStr, regionManager.getPDClient().getHostMapping());
    blockingStub =
        TikvGrpc.newBlockingStub(channel).withDeadlineAfter(deadline, TimeUnit.MILLISECONDS);
    asyncStub = TikvGrpc.newFutureStub(channel).withDeadlineAfter(deadline, TimeUnit.MILLISECONDS);
    if (store.getProxyStore() != null) {
      Metadata header = new Metadata();
      header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
      blockingStub =
          blockingStub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(header));
      asyncStub = asyncStub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(header));
    }
  }

  private Boolean seekLeaderStore(BackOffer backOffer) {
    Histogram.Timer switchLeaderDurationTimer =
        SEEK_LEADER_STORE_DURATION
            .labels(regionManager.getPDClient().getClusterId().toString())
            .startTimer();
    SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("seekLeaderStore");
    try {
      List<Metapb.Peer> peers = region.getFollowerList();
      if (peers.isEmpty()) {
        // no followers available, retry
        logger.warn(String.format("no followers of region[%d] available, retry", region.getId()));
        regionManager.onRequestFail(region);
        return false;
      }

      logger.info(String.format("try switch leader: region[%d]", region.getId()));

      Metapb.Peer peer = switchLeaderStore(backOffer);
      if (peer != null) {
        // we found a leader
        TiStore currentLeaderStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
        if (currentLeaderStore.isReachable()) {
          logger.info(
              String.format(
                  "update leader using switchLeader logic from store[%d] to store[%d]",
                  region.getLeader().getStoreId(), peer.getStoreId()));
          // update region cache
          TiRegion result = regionManager.updateLeader(region, peer.getStoreId());
          if (result != null) {
            region = result;
            // switch to leader store
            store = currentLeaderStore;
            updateClientStub();
            return true;
          }
          return false;
        }
      } else {
        // no leader found, some response does not return normally, there may be network partition.
        logger.warn(
            String.format(
                "leader for region[%d] is not found, it is possible that network partition occurred",
                region.getId()));
      }
    } finally {
      switchLeaderDurationTimer.observeDuration();
      slowLogSpan.end();
    }
    return null;
  }

  private boolean seekProxyStore(BackOffer backOffer) {
    SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("seekProxyStore");
    Histogram.Timer grpcForwardDurationTimer =
        SEEK_PROXY_STORE_DURATION
            .labels(regionManager.getPDClient().getClusterId().toString())
            .startTimer();
    try {
      logger.info(String.format("try grpc forward: region[%d]", region.getId()));
      // when current leader cannot be reached
      TiStore storeWithProxy = switchProxyStore(backOffer);
      if (storeWithProxy == null) {
        // no store available, retry
        logger.warn(String.format("No store available, retry: region[%d]", region.getId()));
        return false;
      }
      // use proxy store to forward requests
      regionManager.updateStore(store, storeWithProxy);
      store = storeWithProxy;
      updateClientStub();
      return true;
    } finally {
      grpcForwardDurationTimer.observeDuration();
      slowLogSpan.end();
    }
  }

  // first: leader peer, second: true if any responses returned with grpc error
  private Metapb.Peer switchLeaderStore(BackOffer backOffer) {
    List<SwitchLeaderTask> responses = new LinkedList<>();
    for (Metapb.Peer peer : region.getFollowerList()) {
      ByteString key = region.getStartKey();
      try {
        TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
        ManagedChannel channel =
            channelFactory.getChannel(
                peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
        TikvGrpc.TikvFutureStub stub =
            TikvGrpc.newFutureStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
        Kvrpcpb.RawGetRequest rawGetRequest =
            Kvrpcpb.RawGetRequest.newBuilder()
                .setContext(makeContext(peer))
                .setKey(codec.encodeKey(key))
                .build();
        ListenableFuture<Kvrpcpb.RawGetResponse> task = stub.rawGet(rawGetRequest);
        responses.add(new SwitchLeaderTask(task, peer));
      } catch (Exception e) {
        logger.warn(
            "switch region[{}] leader store to {} failed: {}",
            region.getId(),
            peer.getStoreId(),
            e);
      }
    }
    while (true) {
      try {
        Thread.sleep(2);
      } catch (InterruptedException e) {
        throw new GrpcException(e);
      }
      List<SwitchLeaderTask> unfinished = new LinkedList<>();
      for (SwitchLeaderTask task : responses) {
        if (!task.task.isDone()) {
          unfinished.add(task);
          continue;
        }
        try {
          Kvrpcpb.RawGetResponse resp = task.task.get();
          if (resp != null) {
            if (!resp.hasRegionError()) {
              // the peer is leader
              logger.info(
                  String.format("rawGet response indicates peer[%d] is leader", task.peer.getId()));
              return task.peer;
            }
          }
        } catch (Exception ignored) {
        }
      }
      if (unfinished.isEmpty()) {
        return null;
      }
      responses = unfinished;
    }
  }

  private TiStore switchProxyStore(BackOffer backOffer) {
    long forwardTimeout = conf.getForwardTimeout();
    List<ForwardCheckTask> responses = new LinkedList<>();
    for (Metapb.Peer peer : region.getFollowerList()) {
      ByteString key = region.getStartKey();
      try {
        TiStore peerStore = regionManager.getStoreById(peer.getStoreId(), backOffer);
        ManagedChannel channel =
            channelFactory.getChannel(
                peerStore.getAddress(), regionManager.getPDClient().getHostMapping());
        TikvGrpc.TikvFutureStub stub =
            TikvGrpc.newFutureStub(channel)
                .withDeadlineAfter(forwardTimeout, TimeUnit.MILLISECONDS);
        Metadata header = new Metadata();
        header.put(TiConfiguration.FORWARD_META_DATA_KEY, store.getStore().getAddress());
        Kvrpcpb.RawGetRequest rawGetRequest =
            Kvrpcpb.RawGetRequest.newBuilder()
                .setContext(makeContext())
                .setKey(codec.encodeKey(key))
                .build();
        ListenableFuture<Kvrpcpb.RawGetResponse> task =
            stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(header))
                .rawGet(rawGetRequest);
        responses.add(new ForwardCheckTask(task, peerStore.getStore()));
      } catch (Exception e) {
        logger.warn(
            "switch region[{}] leader store to {} failed: {}",
            region.getId(),
            peer.getStoreId(),
            e);
      }
    }
    while (true) {
      try {
        Thread.sleep(2);
      } catch (InterruptedException e) {
        throw new GrpcException(e);
      }
      List<ForwardCheckTask> unfinished = new LinkedList<>();
      for (ForwardCheckTask task : responses) {
        if (!task.task.isDone()) {
          unfinished.add(task);
          continue;
        }
        try {
          // any answer will do
          Kvrpcpb.RawGetResponse resp = task.task.get();
          logger.info(
              String.format(
                  "rawGetResponse indicates forward from [%s] to [%s]",
                  task.store.getAddress(), store.getAddress()));
          return store.withProxy(task.store);
        } catch (Exception ignored) {
        }
      }
      if (unfinished.isEmpty()) {
        return null;
      }
      responses = unfinished;
    }
  }

  private static class SwitchLeaderTask {

    private final ListenableFuture<Kvrpcpb.RawGetResponse> task;
    private final Metapb.Peer peer;

    private SwitchLeaderTask(ListenableFuture<Kvrpcpb.RawGetResponse> task, Metapb.Peer peer) {
      this.task = task;
      this.peer = peer;
    }
  }

  private static class ForwardCheckTask {

    private final ListenableFuture<Kvrpcpb.RawGetResponse> task;
    private final Metapb.Store store;

    private ForwardCheckTask(ListenableFuture<Kvrpcpb.RawGetResponse> task, Metapb.Store store) {
      this.task = task;
      this.store = store;
    }
  }
}