/*
 * 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.streaming;

import java.io.DataOutputStream;
import java.io.File;
import java.io.IOException;
import java.math.BigDecimal;
import java.util.BitSet;
import java.util.List;

import org.apache.carbondata.common.logging.LogServiceFactory;
import org.apache.carbondata.core.constants.CarbonCommonConstants;
import org.apache.carbondata.core.datastore.filesystem.CarbonFile;
import org.apache.carbondata.core.datastore.impl.FileFactory;
import org.apache.carbondata.core.datastore.row.CarbonRow;
import org.apache.carbondata.core.metadata.blocklet.index.BlockletMinMaxIndex;
import org.apache.carbondata.core.metadata.datatype.DataType;
import org.apache.carbondata.core.metadata.datatype.DataTypes;
import org.apache.carbondata.core.metadata.schema.table.CarbonTable;
import org.apache.carbondata.core.metadata.schema.table.column.ColumnSchema;
import org.apache.carbondata.core.reader.CarbonHeaderReader;
import org.apache.carbondata.core.util.ByteUtil;
import org.apache.carbondata.core.util.CarbonMetadataUtil;
import org.apache.carbondata.core.util.CarbonProperties;
import org.apache.carbondata.core.util.CarbonUtil;
import org.apache.carbondata.core.util.DataTypeUtil;
import org.apache.carbondata.core.util.path.CarbonTablePath;
import org.apache.carbondata.format.FileHeader;
import org.apache.carbondata.processing.loading.BadRecordsLogger;
import org.apache.carbondata.processing.loading.BadRecordsLoggerProvider;
import org.apache.carbondata.processing.loading.CarbonDataLoadConfiguration;
import org.apache.carbondata.processing.loading.DataField;
import org.apache.carbondata.processing.loading.DataLoadProcessBuilder;
import org.apache.carbondata.processing.loading.converter.RowConverter;
import org.apache.carbondata.processing.loading.converter.impl.RowConverterImpl;
import org.apache.carbondata.processing.loading.model.CarbonLoadModel;
import org.apache.carbondata.processing.loading.parser.RowParser;
import org.apache.carbondata.processing.loading.parser.impl.RowParserImpl;
import org.apache.carbondata.processing.store.writer.AbstractFactDataWriter;
import org.apache.carbondata.processing.util.CarbonDataProcessorUtil;
import org.apache.carbondata.streaming.segment.StreamSegment;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.TaskID;
import org.apache.log4j.Logger;

/**
 * Stream record writer
 */
public class CarbonStreamRecordWriter extends RecordWriter<Void, Object> {

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

  // basic info
  private Configuration hadoopConf;
  private CarbonLoadModel carbonLoadModel;
  private CarbonDataLoadConfiguration configuration;
  private CarbonTable carbonTable;
  private int maxRowNums;
  private int maxCacheSize;

  // parser and converter
  private RowParser rowParser;
  private BadRecordsLogger badRecordLogger;
  private RowConverter converter;
  private final CarbonRow currentRow = new CarbonRow(null);

  // encoder
  private DataField[] dataFields;
  private BitSet nullBitSet;
  private boolean[] isNoDictionaryDimensionColumn;
  private int dimensionWithComplexCount;
  private int measureCount;
  private boolean[] dimensionsIsVarcharTypeMap;
  private DataType[] measureDataTypes;
  private StreamBlockletWriter output = null;
  private String compressorName;

  // data write
  private String segmentDir;
  private String fileName;
  private DataOutputStream outputStream;
  private boolean isFirstRow = true;
  private boolean hasException = false;

  // batch level stats collector
  private BlockletMinMaxIndex batchMinMaxIndex;
  private boolean isClosed = false;

  CarbonStreamRecordWriter(TaskAttemptContext job) throws IOException {
    initialize(job);
  }

  public CarbonStreamRecordWriter(TaskAttemptContext job, CarbonLoadModel carbonLoadModel)
      throws IOException {
    this.carbonLoadModel = carbonLoadModel;
    initialize(job);
  }

  private void initialize(TaskAttemptContext job) throws IOException {
    // set basic information
    hadoopConf = job.getConfiguration();
    if (carbonLoadModel == null) {
      carbonLoadModel = CarbonStreamOutputFormat.getCarbonLoadModel(hadoopConf);
      if (carbonLoadModel == null) {
        throw new IOException(
            "CarbonStreamRecordWriter require configuration: mapreduce.output.carbon.load.model");
      }
    }
    String segmentId = CarbonStreamOutputFormat.getSegmentId(hadoopConf);
    carbonLoadModel.setSegmentId(segmentId);
    carbonTable = carbonLoadModel.getCarbonDataLoadSchema().getCarbonTable();
    long taskNo = TaskID.forName(hadoopConf.get("mapred.tip.id")).getId();
    carbonLoadModel.setTaskNo("" + taskNo);
    configuration = DataLoadProcessBuilder.createConfiguration(carbonLoadModel);
    maxRowNums = hadoopConf.getInt(CarbonStreamOutputFormat.CARBON_STREAM_BLOCKLET_ROW_NUMS,
        CarbonStreamOutputFormat.CARBON_STREAM_BLOCKLET_ROW_NUMS_DEFAULT) - 1;
    maxCacheSize = hadoopConf.getInt(CarbonStreamOutputFormat.CARBON_STREAM_CACHE_SIZE,
        CarbonStreamOutputFormat.CARBON_STREAM_CACHE_SIZE_DEFAULT);

    segmentDir = CarbonTablePath.getSegmentPath(
        carbonTable.getAbsoluteTableIdentifier().getTablePath(), segmentId);
    fileName = CarbonTablePath.getCarbonDataFileName(
        0, taskNo + "", 0, 0, "0",
        segmentId, CarbonProperties.getInstance().getDefaultCompressor());

    // initialize metadata
    isNoDictionaryDimensionColumn =
        CarbonDataProcessorUtil.getNoDictionaryMapping(configuration.getDataFields());
    dimensionWithComplexCount = configuration.getDimensionCount();
    measureCount = configuration.getMeasureCount();
    dataFields = configuration.getDataFields();
    dimensionsIsVarcharTypeMap = new boolean[dimensionWithComplexCount];
    for (int i = 0; i < dimensionWithComplexCount; i++) {
      dimensionsIsVarcharTypeMap[i] = dataFields[i].getColumn().getDataType() == DataTypes.VARCHAR;
    }
    measureDataTypes = new DataType[measureCount];
    for (int i = 0; i < measureCount; i++) {
      measureDataTypes[i] =
          dataFields[dimensionWithComplexCount + i].getColumn().getDataType();
    }
  }

  private void initializeAtFirstRow() throws IOException {
    // initialize parser and converter
    rowParser = new RowParserImpl(dataFields, configuration);
    badRecordLogger = BadRecordsLoggerProvider.createBadRecordLogger(configuration);
    converter =
        new RowConverterImpl(configuration.getDataFields(), configuration, badRecordLogger, true);
    converter.initialize();

    // initialize data writer and compressor
    String filePath = segmentDir + File.separator + fileName;
    CarbonFile carbonFile = FileFactory.getCarbonFile(filePath);
    if (carbonFile.exists()) {
      // if the file is existed, use the append api
      outputStream = FileFactory.getDataOutputStreamUsingAppend(filePath);
      // get the compressor from the file header. In legacy store,
      // the compressor name is not set and it use snappy compressor
      FileHeader header = new CarbonHeaderReader(filePath).readHeader();
      if (header.isSetCompressor_name()) {
        compressorName = header.getCompressor_name();
      } else {
        compressorName = CarbonProperties.getInstance().getDefaultCompressor();
      }
    } else {
      // IF the file is not existed, use the create api
      outputStream = FileFactory.getDataOutputStream(filePath);
      compressorName = carbonTable.getTableInfo().getFactTable().getTableProperties().get(
          CarbonCommonConstants.COMPRESSOR);
      if (null == compressorName) {
        compressorName = CarbonProperties.getInstance().getDefaultCompressor();
      }
      writeFileHeader();
    }

    // initialize encoder
    nullBitSet = new BitSet(dataFields.length);
    int rowBufferSize = hadoopConf.getInt(CarbonStreamOutputFormat.CARBON_ENCODER_ROW_BUFFER_SIZE,
        CarbonStreamOutputFormat.CARBON_ENCODER_ROW_BUFFER_SIZE_DEFAULT);
    output = new StreamBlockletWriter(maxCacheSize, maxRowNums, rowBufferSize,
        isNoDictionaryDimensionColumn.length, measureCount,
        measureDataTypes, compressorName);

    isFirstRow = false;
  }

  @Override
  public void write(Void key, Object value) throws IOException {
    if (isFirstRow) {
      initializeAtFirstRow();
    }
    // null bit set
    nullBitSet.clear();
    Object[] rowData = (Object[]) value;
    currentRow.setRawData(rowData);
    // parse and convert row
    currentRow.setData(rowParser.parseRow(rowData));
    CarbonRow updatedCarbonRow = converter.convert(currentRow);
    if (updatedCarbonRow == null) {
      output.skipRow();
      currentRow.clearData();
    } else {
      for (int i = 0; i < dataFields.length; i++) {
        if (null == currentRow.getObject(i)) {
          nullBitSet.set(i);
        }
      }
      output.nextRow();
      byte[] b = nullBitSet.toByteArray();
      output.writeShort(b.length);
      if (b.length > 0) {
        output.writeBytes(b);
      }
      int dimCount = 0;
      Object columnValue;
      // primitive type dimension
      for (; dimCount < isNoDictionaryDimensionColumn.length; dimCount++) {
        columnValue = currentRow.getObject(dimCount);
        if (null != columnValue) {
          if (isNoDictionaryDimensionColumn[dimCount]) {
            byte[] col = (byte[]) columnValue;
            if (dimensionsIsVarcharTypeMap[dimCount]) {
              output.writeInt(col.length);
            } else {
              output.writeShort(col.length);
            }
            output.writeBytes(col);
            output.dimStatsCollectors[dimCount].update(col);
          } else {
            output.writeInt((int) columnValue);
            output.dimStatsCollectors[dimCount].update(ByteUtil.toBytes((int) columnValue));
          }
        } else {
          output.dimStatsCollectors[dimCount].updateNull(0);
        }
      }
      // complex type dimension
      for (; dimCount < dimensionWithComplexCount; dimCount++) {
        columnValue = currentRow.getObject(dimCount);
        if (null != columnValue) {
          byte[] col = (byte[]) columnValue;
          output.writeShort(col.length);
          output.writeBytes(col);
        }
      }
      // measure
      DataType dataType;
      for (int msrCount = 0; msrCount < measureCount; msrCount++) {
        columnValue = currentRow.getObject(dimCount + msrCount);
        if (null != columnValue) {
          dataType = measureDataTypes[msrCount];
          if (dataType == DataTypes.BOOLEAN) {
            output.writeBoolean((boolean) columnValue);
            output.msrStatsCollectors[msrCount].update((byte) ((boolean) columnValue ? 1 : 0));
          } else if (dataType == DataTypes.SHORT) {
            output.writeShort((short) columnValue);
            output.msrStatsCollectors[msrCount].update((short) columnValue);
          } else if (dataType == DataTypes.INT) {
            output.writeInt((int) columnValue);
            output.msrStatsCollectors[msrCount].update((int) columnValue);
          } else if (dataType == DataTypes.LONG) {
            output.writeLong((long) columnValue);
            output.msrStatsCollectors[msrCount].update((long) columnValue);
          } else if (dataType == DataTypes.DOUBLE) {
            output.writeDouble((double) columnValue);
            output.msrStatsCollectors[msrCount].update((double) columnValue);
          } else if (DataTypes.isDecimal(dataType)) {
            BigDecimal val = (BigDecimal) columnValue;
            byte[] bigDecimalInBytes = DataTypeUtil.bigDecimalToByte(val);
            output.writeShort(bigDecimalInBytes.length);
            output.writeBytes(bigDecimalInBytes);
            output.msrStatsCollectors[msrCount].update((BigDecimal) columnValue);
          } else {
            String msg =
                "unsupported data type:" + dataFields[dimCount + msrCount].getColumn().getDataType()
                .getName();
            LOGGER.error(msg);
            throw new IOException(msg);
          }
        } else {
          output.msrStatsCollectors[msrCount].updateNull(0);
        }
      }
    }

    if (output.isFull()) {
      appendBlockletToDataFile();
    }
  }

  private void writeFileHeader() throws IOException {
    List<ColumnSchema> wrapperColumnSchemaList = CarbonUtil
        .getColumnSchemaList(carbonTable.getVisibleDimensions(), carbonTable.getVisibleMeasures());
    List<org.apache.carbondata.format.ColumnSchema> columnSchemaList = AbstractFactDataWriter
        .getColumnSchemaListAndCardinality(wrapperColumnSchemaList);
    FileHeader fileHeader =
        CarbonMetadataUtil.getFileHeader(true, columnSchemaList, System.currentTimeMillis());
    fileHeader.setIs_footer_present(false);
    fileHeader.setIs_splitable(true);
    fileHeader.setSync_marker(CarbonStreamOutputFormat.CARBON_SYNC_MARKER);
    fileHeader.setCompressor_name(compressorName);
    outputStream.write(CarbonUtil.getByteArray(fileHeader));
  }

  /**
   * write a blocklet to file
   */
  private void appendBlockletToDataFile() throws IOException {
    if (output.getRowIndex() == -1) {
      return;
    }
    output.appendBlocklet(outputStream);
    outputStream.flush();
    if (!isClosed) {
      batchMinMaxIndex = StreamSegment.mergeBlockletMinMax(
          batchMinMaxIndex, output.generateBlockletMinMax(), measureDataTypes);
    }
    // reset data
    output.reset();
  }

  public BlockletMinMaxIndex getBatchMinMaxIndex() {
    if (output == null) {
      return StreamSegment.mergeBlockletMinMax(
          batchMinMaxIndex, null, measureDataTypes);
    }
    return StreamSegment.mergeBlockletMinMax(
        batchMinMaxIndex, output.generateBlockletMinMax(), measureDataTypes);
  }

  public DataType[] getMeasureDataTypes() {
    return measureDataTypes;
  }

  @Override
  public void close(TaskAttemptContext context) throws IOException {
    try {
      isClosed = true;
      // append remain buffer data
      if (!hasException && !isFirstRow) {
        appendBlockletToDataFile();
        converter.finish();
      }
    } finally {
      // close resource
      CarbonUtil.closeStreams(outputStream);
      if (output != null) {
        output.close();
      }
      if (badRecordLogger != null) {
        badRecordLogger.closeStreams();
      }
    }
  }

  public String getSegmentDir() {
    return segmentDir;
  }

  public String getFileName() {
    return fileName;
  }

  public void setHasException(boolean hasException) {
    this.hasException = hasException;
  }
}