/*
 * 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.drill.exec.store.kudu;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

import org.apache.drill.common.exceptions.ExecutionSetupException;
import org.apache.drill.common.exceptions.UserException;
import org.apache.drill.common.expression.SchemaPath;
import org.apache.drill.common.types.TypeProtos.MajorType;
import org.apache.drill.common.types.TypeProtos.MinorType;
import org.apache.drill.common.types.Types;
import org.apache.drill.exec.exception.SchemaChangeException;
import org.apache.drill.exec.expr.TypeHelper;
import org.apache.drill.exec.ops.OperatorContext;
import org.apache.drill.exec.physical.impl.OutputMutator;
import org.apache.drill.exec.record.MaterializedField;
import org.apache.drill.exec.store.AbstractRecordReader;
import org.apache.drill.exec.store.kudu.KuduSubScan.KuduSubScanSpec;
import org.apache.drill.exec.vector.BigIntVector;
import org.apache.drill.exec.vector.BitVector;
import org.apache.drill.exec.vector.Float4Vector;
import org.apache.drill.exec.vector.Float8Vector;
import org.apache.drill.exec.vector.IntVector;
import org.apache.drill.exec.vector.NullableBigIntVector;
import org.apache.drill.exec.vector.NullableBitVector;
import org.apache.drill.exec.vector.NullableFloat4Vector;
import org.apache.drill.exec.vector.NullableFloat8Vector;
import org.apache.drill.exec.vector.NullableIntVector;
import org.apache.drill.exec.vector.NullableTimeStampVector;
import org.apache.drill.exec.vector.NullableVarBinaryVector;
import org.apache.drill.exec.vector.NullableVarCharVector;
import org.apache.drill.exec.vector.TimeStampVector;
import org.apache.drill.exec.vector.ValueVector;
import org.apache.drill.exec.vector.VarBinaryVector;
import org.apache.drill.exec.vector.VarCharVector;
import org.apache.kudu.ColumnSchema;
import org.apache.kudu.Schema;
import org.apache.kudu.Type;
import org.apache.kudu.client.KuduClient;
import org.apache.kudu.client.KuduScanner;
import org.apache.kudu.client.KuduScanner.KuduScannerBuilder;
import org.apache.kudu.client.KuduTable;
import org.apache.kudu.client.RowResult;
import org.apache.kudu.client.RowResultIterator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;

public class KuduRecordReader extends AbstractRecordReader {
  private static final Logger logger = LoggerFactory.getLogger(KuduRecordReader.class);

  private static final int TARGET_RECORD_COUNT = 4000;

  private final KuduClient client;
  private final KuduSubScanSpec scanSpec;
  private KuduScanner scanner;
  private RowResultIterator iterator;

  private OutputMutator output;
  private OperatorContext context;

  private String lastColumnName;
  private Type lastColumnType;

  private static class ProjectedColumnInfo {
    int index;
    ValueVector vv;
    ColumnSchema kuduColumn;
  }

  private ImmutableList<ProjectedColumnInfo> projectedCols;

  public KuduRecordReader(KuduClient client, KuduSubScan.KuduSubScanSpec subScanSpec, List<SchemaPath> projectedColumns) {
    setColumns(projectedColumns);
    this.client = client;
    scanSpec = subScanSpec;
    logger.debug("Scan spec: {}", subScanSpec);
  }

  @Override
  public void setup(OperatorContext context, OutputMutator output) throws ExecutionSetupException {
    this.output = output;
    this.context = context;
    try {
      KuduTable table = client.openTable(scanSpec.getTableName());

      KuduScannerBuilder builder = client.newScannerBuilder(table);
      if (!isStarQuery()) {
        List<String> colNames = Lists.newArrayList();
        for (SchemaPath p : this.getColumns()) {
          colNames.add(p.getRootSegmentPath());
        }
        builder.setProjectedColumnNames(colNames);
      }

      context.getStats().startWait();
      try {
        scanner = builder
            .lowerBoundRaw(scanSpec.getStartKey())
            .exclusiveUpperBoundRaw(scanSpec.getEndKey())
            .build();
      } finally {
        context.getStats().stopWait();
      }
    } catch (Exception e) {
      throw new ExecutionSetupException(e);
    }
  }

  static final Map<Type, MinorType> TYPES;

  static {
    TYPES = ImmutableMap.<Type, MinorType> builder()
        .put(Type.BINARY, MinorType.VARBINARY)
        .put(Type.BOOL, MinorType.BIT)
        .put(Type.DOUBLE, MinorType.FLOAT8)
        .put(Type.FLOAT, MinorType.FLOAT4)
        .put(Type.INT8, MinorType.INT)
        .put(Type.INT16, MinorType.INT)
        .put(Type.INT32, MinorType.INT)
        .put(Type.INT64, MinorType.BIGINT)
        .put(Type.STRING, MinorType.VARCHAR)
        .put(Type.UNIXTIME_MICROS, MinorType.TIMESTAMP)
        .build();
  }

  @Override
  public int next() {
    int rowCount = 0;
    try {
      while (iterator == null || !iterator.hasNext()) {
        if (!scanner.hasMoreRows()) {
          iterator = null;
          return 0;
        }
        context.getStats().startWait();
        try {
          iterator = scanner.nextRows();
        } finally {
          context.getStats().stopWait();
        }
      }
      for (; rowCount < TARGET_RECORD_COUNT && iterator.hasNext(); rowCount++) {
        addRowResult(iterator.next(), rowCount);
      }
    } catch (Exception ex) {
      throw new RuntimeException(ex);
    }
    for (ProjectedColumnInfo pci : projectedCols) {
      pci.vv.getMutator().setValueCount(rowCount);
    }
    return rowCount;
  }

  private void initCols(Schema schema) throws SchemaChangeException {
    ImmutableList.Builder<ProjectedColumnInfo> pciBuilder = ImmutableList.builder();

    for (int i = 0; i < schema.getColumnCount(); i++) {
      ColumnSchema col = schema.getColumnByIndex(i);

      final String name = col.getName();
      final Type kuduType = col.getType();
      lastColumnName = name;
      lastColumnType = kuduType;
      MinorType minorType = TYPES.get(kuduType);
      if (minorType == null) {
        logger.warn("Ignoring column that is unsupported.", UserException
            .unsupportedError()
            .message(
                "A column you queried has a data type that is not currently supported by the Kudu storage plugin. "
                    + "The column's name was %s and its Kudu data type was %s. ",
                name, kuduType.toString())
            .addContext("column Name", name)
            .addContext("plugin", "kudu")
            .build(logger));

        continue;
      }
      MajorType majorType;
      if (col.isNullable()) {
        majorType = Types.optional(minorType);
      } else {
        majorType = Types.required(minorType);
      }
      MaterializedField field = MaterializedField.create(name, majorType);
      final Class<? extends ValueVector> clazz = TypeHelper.getValueVectorClass(
          minorType, majorType.getMode());
      ValueVector vector = output.addField(field, clazz);
      vector.allocateNew();

      ProjectedColumnInfo pci = new ProjectedColumnInfo();
      pci.vv = vector;
      pci.kuduColumn = col;
      pci.index = i;
      pciBuilder.add(pci);
    }

    projectedCols = pciBuilder.build();
  }

  private void addRowResult(RowResult result, int rowIndex) throws SchemaChangeException {
    if (projectedCols == null) {
      initCols(result.getColumnProjection());
    }

    for (ProjectedColumnInfo pci : projectedCols) {
      if (result.isNull(pci.index)) {
        continue;
      }
      switch (pci.kuduColumn.getType()) {
      case BINARY: {
        ByteBuffer value = result.getBinary(pci.index);
        if (pci.kuduColumn.isNullable()) {
          ((NullableVarBinaryVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, value, 0, value.remaining());
        } else {
          ((VarBinaryVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, value, 0, value.remaining());
        }
        break;
      }
      case STRING: {
        ByteBuffer value = ByteBuffer.wrap(result.getString(pci.index).getBytes(StandardCharsets.UTF_8));
        if (pci.kuduColumn.isNullable()) {
          ((NullableVarCharVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, value, 0, value.remaining());
        } else {
          ((VarCharVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, value, 0, value.remaining());
        }
        break;
      }
      case BOOL:
        if (pci.kuduColumn.isNullable()) {
          ((NullableBitVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getBoolean(pci.index) ? 1 : 0);
        } else {
          ((BitVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getBoolean(pci.index) ? 1 : 0);
        }
        break;
      case DOUBLE:
        if (pci.kuduColumn.isNullable()) {
          ((NullableFloat8Vector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getDouble(pci.index));
        } else {
          ((Float8Vector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getDouble(pci.index));
        }
        break;
      case FLOAT:
        if (pci.kuduColumn.isNullable()) {
          ((NullableFloat4Vector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getFloat(pci.index));
        } else {
          ((Float4Vector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getFloat(pci.index));
        }
        break;
      case INT16:
        if (pci.kuduColumn.isNullable()) {
          ((NullableIntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getShort(pci.index));
        } else {
          ((IntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getShort(pci.index));
        }
        break;
      case INT32:
        if (pci.kuduColumn.isNullable()) {
          ((NullableIntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getInt(pci.index));
        } else {
          ((IntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getInt(pci.index));
        }
        break;
      case INT8:
        if (pci.kuduColumn.isNullable()) {
          ((NullableIntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getByte(pci.index));
        } else {
          ((IntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getByte(pci.index));
        }
        break;
      case INT64:
        if (pci.kuduColumn.isNullable()) {
          ((NullableBigIntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getLong(pci.index));
        } else {
          ((BigIntVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getLong(pci.index));
        }
        break;
        case UNIXTIME_MICROS:
        if (pci.kuduColumn.isNullable()) {
          ((NullableTimeStampVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getLong(pci.index) / 1000);
        } else {
          ((TimeStampVector.Mutator) pci.vv.getMutator())
              .setSafe(rowIndex, result.getLong(pci.index) / 1000);
        }
        break;
      default:
        throw new SchemaChangeException("unknown type"); // TODO make better
      }
    }
  }

  @Override
  public void close() {
  }

  @Override
  public String toString() {
    return "KuduRecordReader[Column=" + lastColumnName
        + ", Type=" + lastColumnType
        + "]";
  }
}