/*
    * 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.operation.iterator;
  
  import static java.util.Objects.requireNonNull;
  
  import com.google.protobuf.ByteString;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.tikv.common.TiConfiguration;
  import org.tikv.common.exception.GrpcException;
  import org.tikv.common.exception.KeyException;
  import org.tikv.common.key.Key;
  import org.tikv.common.region.RegionStoreClient;
  import org.tikv.common.region.RegionStoreClient.RegionStoreClientBuilder;
  import org.tikv.common.region.TiRegion;
  import org.tikv.common.region.TiStore;
  import org.tikv.common.util.BackOffer;
  import org.tikv.common.util.ConcreteBackOffer;
  import org.tikv.common.util.Pair;
  import org.tikv.kvproto.Kvrpcpb;
  
  public class ConcreteScanIterator extends ScanIterator {
    private final long version;
    private final Logger logger = LoggerFactory.getLogger(ConcreteScanIterator.class);
  
    public ConcreteScanIterator(
        TiConfiguration conf,
        RegionStoreClientBuilder builder,
        ByteString startKey,
        long version,
        int limit) {
      // Passing endKey as ByteString.EMPTY means that endKey is +INF by default,
      this(conf, builder, startKey, ByteString.EMPTY, version, limit);
    }
  
    public ConcreteScanIterator(
        TiConfiguration conf,
        RegionStoreClientBuilder builder,
        ByteString startKey,
        ByteString endKey,
        long version) {
      // Passing endKey as ByteString.EMPTY means that endKey is +INF by default,
      this(conf, builder, startKey, endKey, version, Integer.MAX_VALUE);
    }
  
    private ConcreteScanIterator(
        TiConfiguration conf,
        RegionStoreClientBuilder builder,
        ByteString startKey,
        ByteString endKey,
        long version,
        int limit) {
      super(conf, builder, startKey, endKey, limit, false);
      this.version = version;
    }
  
    @Override
    TiRegion loadCurrentRegionToCache() throws GrpcException {
      TiRegion region;
      try (RegionStoreClient client = builder.build(startKey)) {
        client.setTimeout(conf.getScanTimeout());
        BackOffer backOffer = ConcreteBackOffer.newScannerNextMaxBackOff();
        currentCache = client.scan(backOffer, startKey, version);
        // If we get region before scan, we will use region from cache which
        // may have wrong end key. This may miss some regions that split from old region.
        // Client will get the newest region during scan. So we need to
        // update region after scan.
        region = client.getRegion();
        return region;
      }
    }
  
    private ByteString resolveCurrentLock(Kvrpcpb.KvPair current) {
      logger.warn(String.format("resolve current key error %s", current.getError().toString()));
      Pair<TiRegion, TiStore> pair =
          builder.getRegionManager().getRegionStorePairByKey(current.getKey());
      TiRegion region = pair.first;
      TiStore store = pair.second;
      BackOffer backOffer =
          ConcreteBackOffer.newGetBackOff(builder.getRegionManager().getPDClient().getClusterId());
      try (RegionStoreClient client = builder.build(region, store)) {
        return client.get(backOffer, current.getKey(), version);
      } catch (Exception e) {
       throw new KeyException(current.getError());
     }
   }
 
   @Override
   public boolean hasNext() {
     Kvrpcpb.KvPair current;
     // continue when cache is empty but not null
     do {
       current = getCurrent();
       if (isCacheDrained() && cacheLoadFails()) {
         endOfScan = true;
         return false;
       }
     } while (currentCache != null && current == null);
     // for last batch to be processed, we have to check if
     return !processingLastBatch
         || current == null
         || (hasEndKey && Key.toRawKey(current.getKey()).compareTo(endKey) < 0);
   }
 
   @Override
   public Kvrpcpb.KvPair next() {
     --limit;
     Kvrpcpb.KvPair current = currentCache.get(index++);
 
     requireNonNull(current, "current kv pair cannot be null");
     if (current.hasError()) {
       ByteString val = resolveCurrentLock(current);
       current = Kvrpcpb.KvPair.newBuilder().setKey(current.getKey()).setValue(val).build();
     }
 
     return current;
   }
 
   /**
    * Cache is drained when - no data extracted - scan limit was not defined - have read the last
    * index of cache - index not initialized
    *
    * @return whether cache is drained
    */
   private boolean isCacheDrained() {
     return currentCache == null || limit <= 0 || index >= currentCache.size() || index == -1;
   }
 
   private Kvrpcpb.KvPair getCurrent() {
     if (isCacheDrained()) {
       return null;
     }
     return currentCache.get(index);
   }
 }