/*
 * 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.carbondata.processing.loading.sort.impl;

import java.io.File;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.carbondata.common.CarbonIterator;
import org.apache.carbondata.common.logging.LogServiceFactory;
import org.apache.carbondata.core.constants.CarbonCommonConstants;
import org.apache.carbondata.core.datastore.row.CarbonRow;
import org.apache.carbondata.core.metadata.schema.ColumnRangeInfo;
import org.apache.carbondata.core.util.CarbonProperties;
import org.apache.carbondata.core.util.CarbonTimeStatisticsFactory;
import org.apache.carbondata.processing.loading.exception.CarbonDataLoadingException;
import org.apache.carbondata.processing.loading.row.CarbonRowBatch;
import org.apache.carbondata.processing.loading.sort.AbstractMergeSorter;
import org.apache.carbondata.processing.loading.sort.unsafe.UnsafeCarbonRowPage;
import org.apache.carbondata.processing.loading.sort.unsafe.UnsafeSortDataRows;
import org.apache.carbondata.processing.loading.sort.unsafe.merger.UnsafeIntermediateMerger;
import org.apache.carbondata.processing.loading.sort.unsafe.merger.UnsafeSingleThreadFinalSortFilesMerger;
import org.apache.carbondata.processing.sort.sortdata.SortParameters;
import org.apache.carbondata.processing.util.CarbonDataProcessorUtil;

import org.apache.commons.lang3.StringUtils;
import org.apache.log4j.Logger;

/**
 * It parallely reads data from array of iterates and do merge sort.
 * First it sorts the data and write to temp files. These temp files will be merge sorted to get
 * final merge sort result.
 * This step is specifically for the data loading with specifying column value range, such as
 * bucketing, sort_column_bounds, it sorts each range of data separately and write to temp files.
 */
public class UnsafeParallelReadMergeSorterWithColumnRangeImpl extends AbstractMergeSorter {

  private static final Logger LOGGER =
      LogServiceFactory.getLogService(
          UnsafeParallelReadMergeSorterWithColumnRangeImpl.class.getName());

  private SortParameters originSortParameters;
  private UnsafeIntermediateMerger[] intermediateFileMergers;
  private int inMemoryChunkSizeInMB;
  private AtomicLong rowCounter;
  private ColumnRangeInfo columnRangeInfo;
  /**
   * counters to collect information about rows processed by each range
   */
  private List<AtomicLong> insideRowCounterList;

  public UnsafeParallelReadMergeSorterWithColumnRangeImpl(AtomicLong rowCounter,
      ColumnRangeInfo columnRangeInfo) {
    this.rowCounter = rowCounter;
    this.columnRangeInfo = columnRangeInfo;
  }

  @Override
  public void initialize(SortParameters sortParameters) {
    this.originSortParameters = sortParameters;
    int totalInMemoryChunkSizeInMB = CarbonProperties.getInstance().getSortMemoryChunkSizeInMB();
    inMemoryChunkSizeInMB = totalInMemoryChunkSizeInMB / columnRangeInfo.getNumOfRanges();
    if (inMemoryChunkSizeInMB < 5) {
      inMemoryChunkSizeInMB = 5;
    }
    this.insideRowCounterList = new ArrayList<>(columnRangeInfo.getNumOfRanges());
    for (int i = 0; i < columnRangeInfo.getNumOfRanges(); i++) {
      insideRowCounterList.add(new AtomicLong(0));
    }
    // Delete if any older file exists in sort temp folder
    CarbonDataProcessorUtil.deleteSortLocationIfExists(sortParameters.getTempFileLocation());
    // create new sort temp directory
    CarbonDataProcessorUtil.createLocations(sortParameters.getTempFileLocation());
  }

  @Override
  public Iterator<CarbonRowBatch>[] sort(Iterator<CarbonRowBatch>[] iterators)
      throws CarbonDataLoadingException {
    UnsafeSortDataRows[] sortDataRows = new UnsafeSortDataRows[columnRangeInfo.getNumOfRanges()];
    intermediateFileMergers = new UnsafeIntermediateMerger[columnRangeInfo.getNumOfRanges()];
    SortParameters[] sortParameterArray = new SortParameters[columnRangeInfo.getNumOfRanges()];
    try {
      for (int i = 0; i < columnRangeInfo.getNumOfRanges(); i++) {
        SortParameters parameters = originSortParameters.getCopy();
        parameters.setPartitionID(i + "");
        parameters.setRangeId(i);
        sortParameterArray[i] = parameters;
        setTempLocation(parameters);
        intermediateFileMergers[i] = new UnsafeIntermediateMerger(parameters);
        sortDataRows[i] =
            new UnsafeSortDataRows(parameters, intermediateFileMergers[i], inMemoryChunkSizeInMB);
        sortDataRows[i].initialize();
      }
    } catch (Exception e) {
      throw new CarbonDataLoadingException(e);
    }
    ExecutorService executorService = Executors.newFixedThreadPool(iterators.length);
    this.threadStatusObserver = new ThreadStatusObserver(executorService);
    final int batchSize = CarbonProperties.getInstance().getBatchSize();
    try {
      for (int i = 0; i < iterators.length; i++) {
        executorService.execute(new SortIteratorThread(iterators[i], sortDataRows, rowCounter,
            this.insideRowCounterList, this.threadStatusObserver));
      }
      executorService.shutdown();
      executorService.awaitTermination(2, TimeUnit.DAYS);
      processRowToNextStep(sortDataRows, originSortParameters);
    } catch (Exception e) {
      checkError();
      throw new CarbonDataLoadingException("Problem while shutdown the server ", e);
    }
    checkError();
    try {
      for (int i = 0; i < intermediateFileMergers.length; i++) {
        intermediateFileMergers[i].finish();
      }
    } catch (Exception e) {
      throw new CarbonDataLoadingException(e);
    }

    Iterator<CarbonRowBatch>[] batchIterator = new Iterator[columnRangeInfo.getNumOfRanges()];
    for (int i = 0; i < sortDataRows.length; i++) {
      batchIterator[i] =
          new MergedDataIterator(sortParameterArray[i], batchSize, intermediateFileMergers[i]);
    }

    return batchIterator;
  }

  private UnsafeSingleThreadFinalSortFilesMerger getFinalMerger(SortParameters sortParameters) {
    String[] storeLocation = CarbonDataProcessorUtil
        .getLocalDataFolderLocation(sortParameters.getCarbonTable(),
            String.valueOf(sortParameters.getTaskNo()), sortParameters.getSegmentId() + "", false,
            false);
    // Set the data file location
    String[] dataFolderLocation = CarbonDataProcessorUtil.arrayAppend(storeLocation,
        File.separator, CarbonCommonConstants.SORT_TEMP_FILE_LOCATION);
    return new UnsafeSingleThreadFinalSortFilesMerger(sortParameters, dataFolderLocation);
  }

  @Override
  public void close() {
    for (int i = 0; i < intermediateFileMergers.length; i++) {
      intermediateFileMergers[i].close();
    }
  }

  /**
   * Below method will be used to process data to next step
   */
  private boolean processRowToNextStep(UnsafeSortDataRows[] sortDataRows, SortParameters parameters)
      throws CarbonDataLoadingException {
    try {
      for (int i = 0; i < sortDataRows.length; i++) {
        // start sorting
        sortDataRows[i].startSorting();
      }
      // check any more rows are present
      LOGGER.info("Record Processed For table: " + parameters.getTableName());
      CarbonTimeStatisticsFactory.getLoadStatisticsInstance()
          .recordSortRowsStepTotalTime(parameters.getPartitionID(), System.currentTimeMillis());
      CarbonTimeStatisticsFactory.getLoadStatisticsInstance()
          .recordDictionaryValuesTotalTime(parameters.getPartitionID(), System.currentTimeMillis());
      return false;
    } catch (Exception e) {
      throw new CarbonDataLoadingException(e);
    }
  }

  private void setTempLocation(SortParameters parameters) {
    String[] carbonDataDirectoryPath = CarbonDataProcessorUtil
        .getLocalDataFolderLocation(parameters.getCarbonTable(), parameters.getTaskNo(),
            parameters.getSegmentId(), false, false);
    String[] tmpLoc = CarbonDataProcessorUtil.arrayAppend(carbonDataDirectoryPath, File.separator,
        CarbonCommonConstants.SORT_TEMP_FILE_LOCATION);
    LOGGER.warn("set temp location: " + StringUtils.join(tmpLoc, ", "));
    parameters.setTempFileLocation(tmpLoc);
  }

  /**
   * This thread iterates the iterator and adds the rows to @{@link UnsafeSortDataRows}
   */
  private static class SortIteratorThread implements Runnable {

    private Iterator<CarbonRowBatch> iterator;

    private UnsafeSortDataRows[] sortDataRows;
    private AtomicLong rowCounter;
    private List<AtomicLong> insideRowCounterList;
    private ThreadStatusObserver threadStatusObserver;

    public SortIteratorThread(Iterator<CarbonRowBatch> iterator,
        UnsafeSortDataRows[] sortDataRows, AtomicLong rowCounter,
        List<AtomicLong> insideRowCounterList,
        ThreadStatusObserver threadStatusObserver) {
      this.iterator = iterator;
      this.sortDataRows = sortDataRows;
      this.rowCounter = rowCounter;
      this.insideRowCounterList = insideRowCounterList;
      this.threadStatusObserver = threadStatusObserver;
    }

    @Override
    public void run() {
      try {
        while (iterator.hasNext()) {
          CarbonRowBatch batch = iterator.next();
          while (batch.hasNext()) {
            CarbonRow row = batch.next();
            if (row != null) {
              UnsafeSortDataRows sortDataRow = sortDataRows[row.getRangeId()];
              synchronized (sortDataRow) {
                rowCounter.getAndIncrement();
                insideRowCounterList.get(row.getRangeId()).getAndIncrement();
                sortDataRow.addRow(row.getData());
              }
            }
          }
        }
        LOGGER.info("Rows processed by each range: " + insideRowCounterList);
      } catch (Exception e) {
        LOGGER.error(e.getMessage(), e);
        this.threadStatusObserver.notifyFailed(e);
      }
    }
  }

  private class MergedDataIterator extends CarbonIterator<CarbonRowBatch> {

    private SortParameters sortParameters;

    private int batchSize;

    private boolean firstRow;

    private UnsafeIntermediateMerger intermediateMerger;

    public MergedDataIterator(SortParameters sortParameters, int batchSize,
        UnsafeIntermediateMerger intermediateMerger) {
      this.sortParameters = sortParameters;
      this.batchSize = batchSize;
      this.intermediateMerger = intermediateMerger;
      this.firstRow = true;
    }

    private UnsafeSingleThreadFinalSortFilesMerger finalMerger;

    @Override
    public boolean hasNext() {
      if (firstRow) {
        firstRow = false;
        finalMerger = getFinalMerger(sortParameters);
        List<UnsafeCarbonRowPage> rowPages = intermediateMerger.getRowPages();
        finalMerger.startFinalMerge(rowPages.toArray(new UnsafeCarbonRowPage[rowPages.size()]),
            intermediateMerger.getMergedPages());
      }
      return finalMerger.hasNext();
    }

    @Override
    public CarbonRowBatch next() {
      int counter = 0;
      CarbonRowBatch rowBatch = new CarbonRowBatch(batchSize);
      while (finalMerger.hasNext() && counter < batchSize) {
        rowBatch.addRow(new CarbonRow(finalMerger.next()));
        counter++;
      }
      return rowBatch;
    }
  }
}
