/*
    * Licensed to the Apache Software Foundation (ASF) under one or more
    * contributor license agreements.  See the NOTICE file distributed with
    * this work for additional information regarding copyright ownership.
    * The ASF licenses this file to You 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.apache.shardingsphere.data.pipeline.cdc.core.importer;
  
  import com.github.benmanes.caffeine.cache.Cache;
  import com.github.benmanes.caffeine.cache.Caffeine;
  import lombok.Getter;
  import lombok.RequiredArgsConstructor;
  import lombok.SneakyThrows;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.commons.lang3.RandomStringUtils;
  import org.apache.commons.lang3.tuple.Pair;
  import org.apache.shardingsphere.data.pipeline.cdc.core.ack.CDCAckId;
  import org.apache.shardingsphere.data.pipeline.cdc.core.ack.CDCAckPosition;
  import org.apache.shardingsphere.data.pipeline.core.channel.PipelineChannel;
  import org.apache.shardingsphere.data.pipeline.core.constant.PipelineSQLOperationType;
  import org.apache.shardingsphere.data.pipeline.core.execute.AbstractPipelineLifecycleRunnable;
  import org.apache.shardingsphere.data.pipeline.core.importer.Importer;
  import org.apache.shardingsphere.data.pipeline.core.importer.sink.PipelineSink;
  import org.apache.shardingsphere.data.pipeline.core.ingest.record.DataRecord;
  import org.apache.shardingsphere.data.pipeline.core.ingest.record.FinishedRecord;
  import org.apache.shardingsphere.data.pipeline.core.ingest.record.Record;
  import org.apache.shardingsphere.data.pipeline.core.job.progress.listener.PipelineJobProgressUpdatedParameter;
  import org.apache.shardingsphere.data.pipeline.core.ratelimit.JobRateLimitAlgorithm;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.LinkedList;
  import java.util.List;
  import java.util.PriorityQueue;
  import java.util.concurrent.TimeUnit;
  import java.util.stream.Collectors;
  
  /**
   * CDC importer.
   */
  @RequiredArgsConstructor
  @Slf4j
  public final class CDCImporter extends AbstractPipelineLifecycleRunnable implements Importer {
      
      @Getter
      private final String importerId = RandomStringUtils.randomAlphanumeric(8);
      
      private final List<CDCChannelProgressPair> channelProgressPairs;
      
      private final int batchSize;
      
      private final long timeoutMillis;
      
      private final PipelineSink sink;
      
      private final boolean needSorting;
      
      private final JobRateLimitAlgorithm rateLimitAlgorithm;
      
      private final PriorityQueue<CSNRecords> csnRecordsQueue = new PriorityQueue<>(new CSNRecordsComparator());
      
      private final Cache<String, List<Pair<CDCChannelProgressPair, CDCAckPosition>>> ackCache = Caffeine.newBuilder().maximumSize(10000).expireAfterAccess(5, TimeUnit.MINUTES).build();
      
      @Override
      protected void runBlocking() {
          CDCImporterManager.putImporter(this);
          for (CDCChannelProgressPair each : channelProgressPairs) {
              each.getJobProgressListener().onProgressUpdated(new PipelineJobProgressUpdatedParameter(0));
          }
          while (isRunning()) {
              if (needSorting) {
                  doWithSorting();
              } else {
                  doWithoutSorting();
              }
              if (channelProgressPairs.isEmpty()) {
                  break;
              }
          }
      }
      
      @SneakyThrows(InterruptedException.class)
      private void doWithSorting() {
          if (null != rateLimitAlgorithm) {
              rateLimitAlgorithm.intercept(PipelineSQLOperationType.INSERT, 1);
          }
          List<CSNRecords> csnRecordsList = getCsnRecordsList();
         if (csnRecordsList.isEmpty()) {
             TimeUnit.MILLISECONDS.sleep(timeoutMillis);
             return;
         }
         // TODO Combine small transactions into a large transaction, to improve transformation performance.
         String ackId = CDCAckId.build(importerId).marshal();
         if (1 == csnRecordsList.size()) {
             processCSNRecords(csnRecordsList.get(0), ackId);
         } else {
             processCSNRecordsList(csnRecordsList, ackId);
         }
     }
     
     private List<CSNRecords> getCsnRecordsList() {
         List<CSNRecords> result = new LinkedList<>();
         CSNRecords firstRecords = null;
         for (int i = 0, count = channelProgressPairs.size(); i < count; i++) {
             prepareTransactionRecords();
             CSNRecords csnRecords = csnRecordsQueue.peek();
             if (null == csnRecords) {
                 continue;
             }
             if (null == firstRecords) {
                 csnRecords = csnRecordsQueue.poll();
                 firstRecords = csnRecords;
                 result.add(csnRecords);
             } else if (csnRecords.getCsn() == firstRecords.getCsn()) {
                 csnRecords = csnRecordsQueue.poll();
                 result.add(csnRecords);
             }
         }
         return result;
     }
     
     // TODO openGauss CSN should be incremented for every transaction. Currently, CSN might be duplicated in transactions.
     // TODO Use channels watermark depth to improve performance.
     private void prepareTransactionRecords() {
         if (csnRecordsQueue.isEmpty()) {
             prepareWhenQueueIsEmpty();
         } else {
             prepareWhenQueueIsNotEmpty(csnRecordsQueue.peek().getCsn());
         }
     }
     
     private void prepareWhenQueueIsEmpty() {
         for (CDCChannelProgressPair each : channelProgressPairs) {
             PipelineChannel channel = each.getChannel();
             List<Record> records = channel.poll();
             if (records.isEmpty()) {
                 continue;
             }
             if (0 == getDataRecordsCount(records)) {
                 channel.ack(records);
                 continue;
             }
             csnRecordsQueue.add(new CSNRecords(findFirstDataRecord(records).getCsn(), each, records));
         }
     }
     
     private void prepareWhenQueueIsNotEmpty(final long oldestCSN) {
         for (CDCChannelProgressPair each : channelProgressPairs) {
             PipelineChannel channel = each.getChannel();
             List<Record> records = channel.peek();
             if (records.isEmpty()) {
                 continue;
             }
             if (0 == getDataRecordsCount(records)) {
                 records = channel.poll();
                 channel.ack(records);
                 continue;
             }
             long csn = findFirstDataRecord(records).getCsn();
             if (csn <= oldestCSN) {
                 records = channel.poll();
                 csnRecordsQueue.add(new CSNRecords(csn, each, records));
             }
         }
     }
     
     private int getDataRecordsCount(final List<Record> records) {
         return (int) records.stream().filter(DataRecord.class::isInstance).count();
     }
     
     private DataRecord findFirstDataRecord(final List<Record> records) {
         for (Record each : records) {
             if (each instanceof DataRecord) {
                 return (DataRecord) each;
             }
         }
         throw new IllegalStateException("No data record found");
     }
     
     private void processCSNRecords(final CSNRecords csnRecords, final String ackId) {
         List<Record> records = csnRecords.getRecords();
         ackCache.put(ackId, Collections.singletonList(Pair.of(csnRecords.getChannelProgressPair(), new CDCAckPosition(records.get(records.size() - 1), getDataRecordsCount(records)))));
         sink.write(ackId, filterDataRecords(records));
     }
     
     private void processCSNRecordsList(final List<CSNRecords> csnRecordsList, final String ackId) {
         List<Pair<CDCChannelProgressPair, CDCAckPosition>> ackValue = csnRecordsList.stream().map(each -> Pair.of(each.getChannelProgressPair(),
                 new CDCAckPosition(each.getRecords().get(each.getRecords().size() - 1), getDataRecordsCount(each.getRecords())))).collect(Collectors.toList());
         ackCache.put(ackId, ackValue);
         Collection<Record> records = new ArrayList<>(ackValue.stream().mapToInt(each -> each.getRight().getDataRecordCount()).sum());
         csnRecordsList.forEach(each -> records.addAll(filterDataRecords(each.getRecords())));
         sink.write(ackId, filterDataRecords(records));
     }
     
     private List<Record> filterDataRecords(final Collection<Record> records) {
         return records.stream().filter(DataRecord.class::isInstance).map(DataRecord.class::cast).collect(Collectors.toList());
     }
     
     private void doWithoutSorting() {
         for (CDCChannelProgressPair each : channelProgressPairs) {
             doWithoutSorting(each);
         }
     }
     
     private void doWithoutSorting(final CDCChannelProgressPair channelProgressPair) {
         PipelineChannel channel = channelProgressPair.getChannel();
         List<Record> records = channel.fetch(batchSize, timeoutMillis);
         if (records.isEmpty()) {
             return;
         }
         Record lastRecord = records.get(records.size() - 1);
         if (records.stream().noneMatch(DataRecord.class::isInstance)) {
             channel.ack(records);
             channelProgressPair.getJobProgressListener().onProgressUpdated(new PipelineJobProgressUpdatedParameter(0));
             if (lastRecord instanceof FinishedRecord) {
                 channelProgressPairs.remove(channelProgressPair);
             }
             return;
         }
         if (null != rateLimitAlgorithm) {
             rateLimitAlgorithm.intercept(PipelineSQLOperationType.INSERT, 1);
         }
         String ackId = CDCAckId.build(importerId).marshal();
         ackCache.put(ackId, Collections.singletonList(Pair.of(channelProgressPair, new CDCAckPosition(records.get(records.size() - 1), getDataRecordsCount(records)))));
         sink.write(ackId, records);
     }
     
     /**
      * Ack.
      *
      * @param ackId ack id
      */
     public void ack(final String ackId) {
         List<Pair<CDCChannelProgressPair, CDCAckPosition>> channelPositionPairList = ackCache.getIfPresent(ackId);
         if (null == channelPositionPairList) {
             log.warn("Could not find cached ack info, ack id: {}", ackId);
             return;
         }
         for (Pair<CDCChannelProgressPair, CDCAckPosition> each : channelPositionPairList) {
             CDCAckPosition ackPosition = each.getRight();
             Record lastRecord = ackPosition.getLastRecord();
             each.getLeft().getChannel().ack(Collections.singletonList(lastRecord));
             if (lastRecord instanceof FinishedRecord) {
                 channelProgressPairs.remove(each.getKey());
             }
             each.getLeft().getJobProgressListener().onProgressUpdated(new PipelineJobProgressUpdatedParameter(ackPosition.getDataRecordCount()));
         }
         ackCache.invalidate(ackId);
     }
     
     @Override
     protected void doStop() {
         CDCImporterManager.removeImporter(importerId);
     }
 }