/*
 * 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.iotdb.db.queryengine.transformation.dag.intermediate;

import org.apache.iotdb.db.queryengine.plan.expression.Expression;
import org.apache.iotdb.db.queryengine.transformation.api.LayerReader;
import org.apache.iotdb.db.queryengine.transformation.api.LayerRowWindowReader;
import org.apache.iotdb.db.queryengine.transformation.api.YieldableState;
import org.apache.iotdb.db.queryengine.transformation.dag.adapter.ElasticSerializableTVListBackedSingleColumnWindow;
import org.apache.iotdb.db.queryengine.transformation.dag.memory.SafetyLine;
import org.apache.iotdb.db.queryengine.transformation.dag.memory.SafetyLine.SafetyPile;
import org.apache.iotdb.db.queryengine.transformation.dag.util.LayerCacheUtils;
import org.apache.iotdb.db.queryengine.transformation.dag.util.TransformUtils;
import org.apache.iotdb.db.queryengine.transformation.datastructure.iterator.TVListForwardIterator;
import org.apache.iotdb.db.queryengine.transformation.datastructure.tv.ElasticSerializableTVList;
import org.apache.iotdb.db.queryengine.transformation.datastructure.util.ValueRecorder;
import org.apache.iotdb.udf.api.access.RowWindow;
import org.apache.iotdb.udf.api.customizer.strategy.SessionTimeWindowAccessStrategy;
import org.apache.iotdb.udf.api.customizer.strategy.SlidingSizeWindowAccessStrategy;
import org.apache.iotdb.udf.api.customizer.strategy.SlidingTimeWindowAccessStrategy;
import org.apache.iotdb.udf.api.customizer.strategy.StateWindowAccessStrategy;

import org.apache.tsfile.block.column.Column;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.read.common.block.column.TimeColumn;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

public class SingleInputMultiReferenceLayer extends IntermediateLayer {

  private static final Logger LOGGER =
      LoggerFactory.getLogger(SingleInputMultiReferenceLayer.class);

  private final LayerReader parentLayerReader;
  private final TSDataType parentLayerReaderDataType;
  private final boolean isParentLayerReaderConstant;
  private final ElasticSerializableTVList tvList;
  private final SafetyLine safetyLine;

  public SingleInputMultiReferenceLayer(
      Expression expression,
      String queryId,
      float memoryBudgetInMB,
      LayerReader parentLayerReader) {
    super(expression, queryId, memoryBudgetInMB);
    this.parentLayerReader = parentLayerReader;

    parentLayerReaderDataType = this.parentLayerReader.getDataTypes()[0];
    isParentLayerReaderConstant = this.parentLayerReader.isConstantPointReader();
    tvList =
        ElasticSerializableTVList.construct(
            parentLayerReaderDataType, queryId, memoryBudgetInMB, CACHE_BLOCK_SIZE);
    safetyLine = new SafetyLine();
  }

  @Override
  public LayerReader constructReader() {
    return new LayerReader() {
      private final SafetyPile safetyPile = safetyLine.addSafetyPile();

      private TimeColumn cachedTimes = null;
      private Column cachedValues = null;
      private int cacheConsumed = 0;

      private final TVListForwardIterator iterator = tvList.constructIterator();

      @Override
      public boolean isConstantPointReader() {
        return isParentLayerReaderConstant;
      }

      @Override
      public YieldableState yield() throws Exception {
        // Column cached in reader is not yet consumed
        if (cachedTimes != null && cacheConsumed < cachedTimes.getPositionCount()) {
          return YieldableState.YIELDABLE;
        }

        // TVList still has some cached columns
        if (iterator.hasNext()) {
          iterator.next();
          cachedTimes = iterator.currentTimes();
          cachedValues = iterator.currentValues();

          return YieldableState.YIELDABLE;
        }

        // No data cached, yield from parent layer reader
        YieldableState state = LayerCacheUtils.yieldPoints(parentLayerReader, tvList);
        if (state == YieldableState.YIELDABLE) {
          iterator.next();
          cachedTimes = iterator.currentTimes();
          cachedValues = iterator.currentValues();
        }
        return state;
      }

      @Override
      public void consumedAll() {
        int steps = cachedTimes.getPositionCount() - cacheConsumed;
        safetyPile.moveForward(steps);
        tvList.setEvictionUpperBound(safetyLine.getSafetyLine());

        cacheConsumed = 0;
        cachedTimes = null;
        cachedValues = null;
      }

      @Override
      public Column[] current() {
        return cacheConsumed == 0
            ? new Column[] {cachedValues, cachedTimes}
            : new Column[] {
              cachedValues.subColumn(cacheConsumed), cachedTimes.subColumn(cacheConsumed)
            };
      }

      @Override
      public TSDataType[] getDataTypes() {
        return new TSDataType[] {parentLayerReaderDataType};
      }
    };
  }

  @Override
  protected LayerRowWindowReader constructRowSlidingSizeWindowReader(
      SlidingSizeWindowAccessStrategy strategy, float memoryBudgetInMB) {
    return new LayerRowWindowReader() {

      private final int windowSize = strategy.getWindowSize();
      private final int slidingStep = strategy.getSlidingStep();

      private final SafetyPile safetyPile = safetyLine.addSafetyPile();
      private final ElasticSerializableTVListBackedSingleColumnWindow window =
          new ElasticSerializableTVListBackedSingleColumnWindow(tvList);

      private boolean hasCached = false;
      private int beginIndex = -slidingStep;

      @Override
      public YieldableState yield() throws Exception {
        if (hasCached) {
          return YieldableState.YIELDABLE;
        }

        beginIndex += slidingStep;
        int endIndex = beginIndex + windowSize;
        if (beginIndex < 0 || endIndex < 0) {
          LOGGER.warn(
              "LayerRowWindowReader index overflow. beginIndex: {}, endIndex: {}, windowSize: {}.",
              beginIndex,
              endIndex,
              windowSize);
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }

        final int pointsToBeCollected = endIndex - tvList.size();
        if (pointsToBeCollected > 0) {
          final YieldableState yieldableState =
              LayerCacheUtils.yieldPoints(parentLayerReader, tvList, pointsToBeCollected);
          if (yieldableState == YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA) {
            beginIndex -= slidingStep;
            return YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA;
          }

          if (tvList.size() <= beginIndex) {
            return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
          }

          // TVList's size may be less than endIndex
          // When parent layer reader has no more data
          endIndex = Math.min(endIndex, tvList.size());
        }

        window.seek(beginIndex, endIndex, tvList.getTime(beginIndex), tvList.getTime(endIndex - 1));

        hasCached = true;
        return YieldableState.YIELDABLE;
      }

      @Override
      public void readyForNext() {
        hasCached = false;

        safetyPile.moveForwardTo(beginIndex + 1);
        tvList.setEvictionUpperBound(safetyLine.getSafetyLine());
      }

      @Override
      public TSDataType[] getDataTypes() {
        return new TSDataType[] {parentLayerReaderDataType};
      }

      @Override
      public RowWindow currentWindow() {
        return window;
      }
    };
  }

  @Override
  protected LayerRowWindowReader constructRowSlidingTimeWindowReader(
      SlidingTimeWindowAccessStrategy strategy, float memoryBudgetInMB) {

    final long timeInterval = strategy.getTimeInterval();
    final long slidingStep = strategy.getSlidingStep();
    final long displayWindowEnd = strategy.getDisplayWindowEnd();

    final SafetyPile safetyPile = safetyLine.addSafetyPile();
    final ElasticSerializableTVListBackedSingleColumnWindow window =
        new ElasticSerializableTVListBackedSingleColumnWindow(tvList);

    final long nextWindowTimeBeginGivenByStrategy = strategy.getDisplayWindowBegin();

    return new LayerRowWindowReader() {

      private boolean isFirstIteration = true;
      private boolean hasAtLeastOneRow = false;
      private boolean hasCached = false;
      private long nextWindowTimeBegin = nextWindowTimeBeginGivenByStrategy;
      private int nextIndexBegin = 0;
      private int nextIndexEnd = 0;
      private long currentEndTime = Long.MAX_VALUE;

      private final TVListForwardIterator beginIterator = tvList.constructIterator();
      private TimeColumn cachedBeginTimeColumn;
      private int cachedBeginConsumed;

      private TimeColumn cachedEndTimeColumn;
      private int cachedEndConsumed;

      @Override
      public YieldableState yield() throws Exception {
        if (isFirstIteration) {
          if (tvList.size() == 0) {
            final YieldableState state = parentLayerReader.yield();
            if (state != YieldableState.YIELDABLE) {
              return state;
            }

            Column[] columns = parentLayerReader.current();
            TimeColumn times = (TimeColumn) columns[1];
            Column values = columns[0];

            tvList.putColumn(times, values);
            parentLayerReader.consumedAll();

            cachedEndTimeColumn = times;
          }
          if (nextWindowTimeBegin == Long.MIN_VALUE) {
            // display window begin should be set to the same as the min timestamp of the query
            // result set
            nextWindowTimeBegin = cachedEndTimeColumn.getStartTime();
          }
          hasAtLeastOneRow = tvList.size() != 0;
          if (hasAtLeastOneRow) {
            currentEndTime = cachedEndTimeColumn.getEndTime();
          }
          isFirstIteration = false;
        }

        if (hasCached) {
          return YieldableState.YIELDABLE;
        }

        if (!hasAtLeastOneRow || displayWindowEnd <= nextWindowTimeBegin) {
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }

        long nextWindowTimeEnd = Math.min(nextWindowTimeBegin + timeInterval, displayWindowEnd);
        while (currentEndTime < nextWindowTimeEnd) {
          final YieldableState state = parentLayerReader.yield();
          if (state == YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA) {
            return YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA;
          }
          if (state == YieldableState.NOT_YIELDABLE_NO_MORE_DATA) {
            break;
          }
          // Generate data
          Column[] columns = parentLayerReader.current();
          TimeColumn times = (TimeColumn) columns[1];
          Column values = columns[0];
          // Put data into container
          tvList.putColumn(times, values);
          parentLayerReader.consumedAll();
          currentEndTime = times.getEndTime();
          // Increase nextIndexEnd
          nextIndexEnd += cachedEndTimeColumn.getPositionCount() - cachedEndConsumed;
          // Update cache
          cachedEndTimeColumn = times;
          cachedEndConsumed = 0;
        }

        // Set nextIndexEnd field
        while (cachedEndConsumed < cachedEndTimeColumn.getPositionCount()) {
          if (cachedEndTimeColumn.getLong(cachedEndConsumed) >= nextWindowTimeEnd) {
            break;
          }
          cachedEndConsumed++;
          nextIndexEnd++;
        }

        // Set nextIndexBegin field
        boolean findNextIndexBegin = false;
        while (!findNextIndexBegin) {
          while (cachedBeginTimeColumn != null
              && cachedBeginConsumed < cachedBeginTimeColumn.getPositionCount()) {
            if (cachedBeginTimeColumn.getLong(cachedBeginConsumed) >= nextWindowTimeBegin) {
              findNextIndexBegin = true;
              break;
            }
            cachedBeginConsumed++;
            nextIndexBegin++;
          }

          if (!findNextIndexBegin) {
            if (beginIterator.hasNext()) {
              beginIterator.next();

              cachedBeginConsumed = 0;
              cachedBeginTimeColumn = beginIterator.currentTimes();
            } else {
              // No more data
              // Set nextIndexBegin to list's size
              findNextIndexBegin = true;
            }
          }
        }

        if ((nextIndexEnd == nextIndexBegin)
            && nextWindowTimeEnd < cachedEndTimeColumn.getEndTime()) {
          window.setEmptyWindow(nextWindowTimeBegin, nextWindowTimeEnd);
          return YieldableState.YIELDABLE;
        }

        window.seek(
            nextIndexBegin,
            nextIndexEnd,
            nextWindowTimeBegin,
            nextWindowTimeBegin + timeInterval - 1);

        hasCached = !(nextIndexBegin == nextIndexEnd && nextIndexEnd == tvList.size());
        return hasCached ? YieldableState.YIELDABLE : YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
      }

      @Override
      public void readyForNext() {
        hasCached = false;
        nextWindowTimeBegin += slidingStep;

        safetyPile.moveForwardTo(nextIndexBegin + 1);
        tvList.setEvictionUpperBound(safetyLine.getSafetyLine());
      }

      @Override
      public TSDataType[] getDataTypes() {
        return new TSDataType[] {parentLayerReaderDataType};
      }

      @Override
      public RowWindow currentWindow() {
        return window;
      }
    };
  }

  @Override
  protected LayerRowWindowReader constructRowSessionTimeWindowReader(
      SessionTimeWindowAccessStrategy strategy, float memoryBudgetInMB) {
    final long displayWindowBegin = strategy.getDisplayWindowBegin();
    final long displayWindowEnd = strategy.getDisplayWindowEnd();
    final long sessionTimeGap = strategy.getSessionTimeGap();

    final SafetyPile safetyPile = safetyLine.addSafetyPile();
    final ElasticSerializableTVListBackedSingleColumnWindow window =
        new ElasticSerializableTVListBackedSingleColumnWindow(tvList);

    return new LayerRowWindowReader() {

      private boolean isFirstIteration = true;
      private boolean hasAtLeastOneRow = false;

      private long nextWindowTimeBegin = displayWindowBegin;
      private long nextWindowTimeEnd = 0;
      private int nextIndexBegin = 0;
      private int nextIndexEnd = 0;

      private TimeColumn cachedTimes;
      private int cachedConsumed;

      @Override
      public YieldableState yield() throws Exception {
        if (isFirstIteration) {
          YieldableState state = yieldInFirstIteration();
          if (state != YieldableState.YIELDABLE) {
            return state;
          }
        }

        if (!hasAtLeastOneRow || nextWindowTimeBegin >= displayWindowEnd) {
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }

        // Set nextIndexEnd
        long curTime = cachedTimes.getLong(cachedConsumed);
        if (cachedConsumed < cachedTimes.getPositionCount()) {
          nextIndexEnd++;
          cachedConsumed++;
        }
        boolean findWindow = false;
        // Find target window or no more data to exit
        while (!findWindow && cachedConsumed < cachedTimes.getPositionCount()) {
          while (cachedConsumed < cachedTimes.getPositionCount()) {
            long nextTime = cachedTimes.getLong(cachedConsumed);

            if (nextTime >= displayWindowEnd) {
              findWindow = true;
              break;
            }

            if (nextTime - curTime > sessionTimeGap) {
              findWindow = true;
              break;
            }
            nextIndexEnd++;
            cachedConsumed++;
            curTime = nextTime;
          }

          if (!findWindow) {
            if (cachedTimes.getEndTime() < displayWindowEnd) {
              YieldableState state = yieldAndCache();
              if (state == YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA) {
                return YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA;
              } else if (state == YieldableState.NOT_YIELDABLE_NO_MORE_DATA) {
                break;
              }
            }
          }
        }
        // Set nextWindowTimeEnd
        nextWindowTimeEnd = tvList.getTime(nextIndexEnd - 1);

        if (nextIndexBegin == nextIndexEnd) {
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }
        window.seek(nextIndexBegin, nextIndexEnd, nextWindowTimeBegin, nextWindowTimeEnd);
        return YieldableState.YIELDABLE;
      }

      private YieldableState yieldInFirstIteration() throws Exception {
        // Yield initial data in first iteration
        if (tvList.size() == 0) {
          YieldableState state = yieldAndCache();
          if (state != YieldableState.YIELDABLE) {
            return state;
          }
        }
        // Initialize essential information
        nextWindowTimeBegin = Math.max(displayWindowBegin, cachedTimes.getStartTime());
        hasAtLeastOneRow = tvList.size() != 0;
        isFirstIteration = false;

        // Set initial nextIndexBegin
        long currentEndTime = cachedTimes.getEndTime();
        // Find corresponding block
        while (currentEndTime < nextWindowTimeBegin) {
          // Consume all data
          cachedConsumed = cachedTimes.getPositionCount();
          nextIndexBegin += cachedConsumed;

          YieldableState state = yieldAndCache();
          if (state != YieldableState.YIELDABLE) {
            // Cannot find nextIndexBegin
            // Set nextIndexEnd to nextIndexBegin and exit
            nextIndexEnd = nextIndexBegin;
            return state;
          }
        }
        // Find nextIndexBegin in block
        while (cachedConsumed < cachedTimes.getPositionCount()) {
          if (cachedTimes.getLong(cachedConsumed) >= nextWindowTimeBegin) {
            break;
          }
          cachedConsumed++;
          nextIndexBegin++;
        }
        nextIndexEnd = nextIndexBegin;

        return YieldableState.YIELDABLE;
      }

      private YieldableState yieldAndCache() throws Exception {
        final YieldableState state = parentLayerReader.yield();
        if (state != YieldableState.YIELDABLE) {
          return state;
        }
        Column[] columns = parentLayerReader.current();
        TimeColumn times = (TimeColumn) columns[1];
        Column values = columns[0];

        tvList.putColumn(times, values);
        parentLayerReader.consumedAll();

        cachedTimes = times;
        cachedConsumed = 0;

        return YieldableState.YIELDABLE;
      }

      @Override
      public void readyForNext() throws IOException {
        if (nextIndexEnd < tvList.size()) {
          nextWindowTimeBegin = tvList.getTime(nextIndexEnd);
        }
        safetyPile.moveForwardTo(nextIndexBegin + 1);
        tvList.setEvictionUpperBound(safetyLine.getSafetyLine());
        nextIndexBegin = nextIndexEnd;
      }

      @Override
      public TSDataType[] getDataTypes() {
        return new TSDataType[] {parentLayerReaderDataType};
      }

      @Override
      public RowWindow currentWindow() {
        return window;
      }
    };
  }

  @Override
  protected LayerRowWindowReader constructRowStateWindowReader(
      StateWindowAccessStrategy strategy, float memoryBudgetInMB) {
    final long displayWindowBegin = strategy.getDisplayWindowBegin();
    final long displayWindowEnd = strategy.getDisplayWindowEnd();
    final double delta = strategy.getDelta();

    final SafetyPile safetyPile = safetyLine.addSafetyPile();
    final ElasticSerializableTVListBackedSingleColumnWindow window =
        new ElasticSerializableTVListBackedSingleColumnWindow(tvList);

    return new LayerRowWindowReader() {

      private boolean isFirstIteration = true;
      private boolean hasAtLeastOneRow = false;

      private long nextWindowTimeBegin = displayWindowBegin;
      private long nextWindowTimeEnd = 0;
      private int nextIndexBegin = 0;
      private int nextIndexEnd = 0;

      private TimeColumn cachedTimes;
      private Column cachedValues;
      private int cachedConsumed;

      private final ValueRecorder valueRecorder = new ValueRecorder();

      @Override
      public YieldableState yield() throws Exception {
        if (isFirstIteration) {
          YieldableState state = yieldInFirstIteration();
          if (state != YieldableState.YIELDABLE) {
            return state;
          }
        }

        if (!hasAtLeastOneRow || nextWindowTimeBegin >= displayWindowEnd) {
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }

        // Set nextIndexEnd
        if (cachedConsumed < cachedTimes.getPositionCount()) {
          nextIndexEnd++;
          cachedConsumed++;
        }
        boolean findWindow = false;
        // Find target window or no more data to exit
        while (!findWindow && cachedConsumed < cachedTimes.getPositionCount()) {
          while (cachedConsumed < cachedTimes.getPositionCount()) {
            long nextTime = cachedTimes.getLong(cachedConsumed);

            if (nextTime >= displayWindowEnd) {
              findWindow = true;
              break;
            }

            if (TransformUtils.splitWindowForStateWindow(
                parentLayerReaderDataType, valueRecorder, delta, cachedValues, cachedConsumed)) {
              findWindow = true;
              break;
            }
            nextIndexEnd++;
            cachedConsumed++;
          }

          if (!findWindow) {
            if (cachedTimes.getEndTime() < displayWindowEnd) {
              YieldableState state = yieldAndCache();
              if (state == YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA) {
                return YieldableState.NOT_YIELDABLE_WAITING_FOR_DATA;
              } else if (state == YieldableState.NOT_YIELDABLE_NO_MORE_DATA) {
                break;
              }
            }
          }
        }
        // Set nextWindowTimeEnd
        nextWindowTimeEnd = tvList.getTime(nextIndexEnd - 1);

        if (nextIndexBegin == nextIndexEnd) {
          return YieldableState.NOT_YIELDABLE_NO_MORE_DATA;
        }
        window.seek(nextIndexBegin, nextIndexEnd, nextWindowTimeBegin, nextWindowTimeEnd);
        return YieldableState.YIELDABLE;
      }

      private YieldableState yieldInFirstIteration() throws Exception {
        // Yield initial data in first iteration
        if (tvList.size() == 0) {
          YieldableState state = yieldAndCache();
          if (state != YieldableState.YIELDABLE) {
            return state;
          }
        }
        // Initialize essential information
        nextWindowTimeBegin = Math.max(displayWindowBegin, cachedTimes.getStartTime());
        hasAtLeastOneRow = tvList.size() != 0;
        isFirstIteration = false;

        // Set initial nextIndexBegin
        long currentEndTime = cachedTimes.getEndTime();
        // Find corresponding block
        while (currentEndTime < nextWindowTimeBegin) {
          // Consume all data
          cachedConsumed = cachedTimes.getPositionCount();
          nextIndexBegin += cachedConsumed;

          YieldableState state = yieldAndCache();
          if (state != YieldableState.YIELDABLE) {
            // Cannot find nextIndexBegin
            // Set nextIndexEnd to nextIndexBegin and exit
            nextIndexEnd = nextIndexBegin;
            return state;
          }
        }
        // Find nextIndexBegin in block
        while (cachedConsumed < cachedTimes.getPositionCount()) {
          if (cachedTimes.getLong(cachedConsumed) >= nextWindowTimeBegin) {
            break;
          }
          cachedConsumed++;
          nextIndexBegin++;
        }
        nextIndexEnd = nextIndexBegin;

        return YieldableState.YIELDABLE;
      }

      private YieldableState yieldAndCache() throws Exception {
        final YieldableState state = parentLayerReader.yield();
        if (state != YieldableState.YIELDABLE) {
          return state;
        }
        Column[] columns = parentLayerReader.current();
        TimeColumn times = (TimeColumn) columns[1];
        Column values = columns[0];

        tvList.putColumn(times, values);
        parentLayerReader.consumedAll();

        cachedTimes = times;
        cachedValues = values;
        cachedConsumed = 0;

        return YieldableState.YIELDABLE;
      }

      @Override
      public void readyForNext() throws IOException {
        if (nextIndexEnd < tvList.size()) {
          nextWindowTimeBegin = tvList.getTime(nextIndexEnd);
        }
        safetyPile.moveForwardTo(nextIndexBegin + 1);
        tvList.setEvictionUpperBound(safetyLine.getSafetyLine());
        nextIndexBegin = nextIndexEnd;
      }

      @Override
      public TSDataType[] getDataTypes() {
        return new TSDataType[] {parentLayerReaderDataType};
      }

      @Override
      public RowWindow currentWindow() {
        return window;
      }
    };
  }
}
