/*
    * 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 io.grpc.ManagedChannel;
  import io.grpc.health.v1.HealthCheckRequest;
  import io.grpc.health.v1.HealthCheckResponse;
  import io.grpc.health.v1.HealthGrpc;
  import io.grpc.stub.ClientCalls;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Supplier;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.tikv.common.ReadOnlyPDClient;
  import org.tikv.common.util.ChannelFactory;
  import org.tikv.common.util.ConcreteBackOffer;
  import org.tikv.kvproto.Metapb;
  import org.tikv.kvproto.Mpp;
  import org.tikv.kvproto.Mpp.IsAliveRequest;
  import org.tikv.kvproto.TikvGrpc;
  
  public class StoreHealthyChecker implements Runnable {
    private static final Logger logger = LoggerFactory.getLogger(StoreHealthyChecker.class);
    private static final long MAX_CHECK_STORE_TOMBSTONE_TICK = 60;
    private final BlockingQueue<TiStore> taskQueue;
    private final ChannelFactory channelFactory;
    private final ReadOnlyPDClient pdClient;
    private final RegionCache cache;
    private long checkTombstoneTick;
    private final long timeout;
  
    public StoreHealthyChecker(
        ChannelFactory channelFactory, ReadOnlyPDClient pdClient, RegionCache cache, long timeout) {
      this.taskQueue = new LinkedBlockingQueue<>();
      this.channelFactory = channelFactory;
      this.pdClient = pdClient;
      this.cache = cache;
      this.checkTombstoneTick = 0;
      this.timeout = timeout;
    }
  
    public boolean scheduleStoreHealthCheck(TiStore store) {
      // add queue false, mark it reachable so that it can be put again.
      return this.taskQueue.add(store);
    }
  
    private List<TiStore> getValidStores() {
      List<TiStore> unhealthStore = new LinkedList<>();
      while (!this.taskQueue.isEmpty()) {
        try {
          TiStore store = this.taskQueue.take();
          if (!store.isValid()) {
            continue;
          }
          unhealthStore.add(store);
        } catch (Exception e) {
          return unhealthStore;
        }
      }
      return unhealthStore;
    }
  
    private boolean checkStoreHealth(TiStore store) {
      String addressStr = store.getStore().getAddress();
      if (store.isTiFlash()) {
        return checkTiFlashHealth(addressStr);
      }
      return checkTiKVHealth(addressStr);
    }
  
    private boolean checkTiFlashHealth(String addressStr) {
      try {
        ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
        TikvGrpc.TikvBlockingStub stub =
            TikvGrpc.newBlockingStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
        Supplier<IsAliveRequest> factory = () -> Mpp.IsAliveRequest.newBuilder().build();
        Mpp.IsAliveResponse resp =
            ClientCalls.blockingUnaryCall(
                stub.getChannel(), TikvGrpc.getIsAliveMethod(), stub.getCallOptions(), factory.get());
        return resp != null && resp.getAvailable();
      } catch (Exception e) {
       logger.info(
           "fail to check TiFlash health, regard as unhealthy. TiFlash address: " + addressStr, e);
       return false;
     }
   }
 
   private boolean checkTiKVHealth(String addressStr) {
     try {
       ManagedChannel channel = channelFactory.getChannel(addressStr, pdClient.getHostMapping());
       HealthGrpc.HealthBlockingStub stub =
           HealthGrpc.newBlockingStub(channel).withDeadlineAfter(timeout, TimeUnit.MILLISECONDS);
       HealthCheckRequest req = HealthCheckRequest.newBuilder().build();
       HealthCheckResponse resp = stub.check(req);
       return resp.getStatus() == HealthCheckResponse.ServingStatus.SERVING;
     } catch (Exception e) {
       logger.info("fail to check TiKV health, regard as unhealthy. TiKV address: " + addressStr, e);
       return false;
     }
   }
 
   private boolean checkStoreTombstone(TiStore store) {
     try {
       Metapb.Store newStore =
           pdClient.getStore(
               ConcreteBackOffer.newRawKVBackOff(pdClient.getClusterId()), store.getId());
       if (newStore != null && newStore.getState() == Metapb.StoreState.Tombstone) {
         return true;
       }
     } catch (Exception e) {
       logger.info("fail to check tombstone stores", e);
       return false;
     }
     return false;
   }
 
   @Override
   public void run() {
     checkTombstoneTick += 1;
     boolean needCheckTombstoneStore = false;
     if (checkTombstoneTick >= MAX_CHECK_STORE_TOMBSTONE_TICK) {
       needCheckTombstoneStore = true;
       checkTombstoneTick = 0;
     }
     List<TiStore> allStores = getValidStores();
     List<TiStore> unreachableStore = new LinkedList<>();
     for (TiStore store : allStores) {
       if (needCheckTombstoneStore) {
         if (checkStoreTombstone(store)) {
           continue;
         }
       }
 
       if (checkStoreHealth(store)) {
         if (store.getProxyStore() != null) {
           TiStore newStore = store.withProxy(null);
           logger.warn(String.format("store [%s] recovers to be reachable", store.getAddress()));
           if (cache.putStore(newStore.getId(), newStore)) {
             this.taskQueue.add(newStore);
             continue;
           }
         } else {
           if (!store.isReachable()) {
             logger.warn(String.format("store [%s] recovers to be reachable", store.getAddress()));
             store.markReachable();
           }
         }
       } else if (store.isReachable()) {
         unreachableStore.add(store);
         continue;
       }
       this.taskQueue.add(store);
     }
     if (!unreachableStore.isEmpty()) {
       try {
         Thread.sleep(timeout);
       } catch (Exception e) {
         this.taskQueue.addAll(unreachableStore);
         return;
       }
       for (TiStore store : unreachableStore) {
         if (!checkStoreHealth(store)) {
           logger.warn(String.format("store [%s] is not reachable", store.getAddress()));
           store.markUnreachable();
         }
         this.taskQueue.add(store);
       }
     }
   }
 }