package com.chenyang.nse.bussiness.tools.dataprocess.hive.hdfs.format;

import com.chenyang.nse.bussiness.entity.db.ColumnInfo;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.ql.exec.vector.BytesColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.ColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DecimalColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.DoubleColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.StructColumnVector;
import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch;
import org.apache.hadoop.hive.serde2.io.HiveDecimalWritable;
import org.apache.orc.OrcFile;
import org.apache.orc.TypeDescription;
import org.apache.orc.Writer;

public class OrcFileWriter extends HdfsFileWriter {
   private final OrcFileReader reader;
   private Path output;
   private Writer writer;
   private TypeDescription schema;
   private VectorizedRowBatch batch;
   private int index = -1;

   public OrcFileWriter(Path output, OrcFileReader reader, List<ColumnInfo> cols) {
      this.output = output;
      this.reader = reader;
      this.cols = cols;
   }

   public void initialize() throws IOException {
      Configuration conf = new Configuration();
      OrcFile.WriterOptions options = OrcFile.writerOptions(conf);
      this.schema = this.reader.getSchema();
      options.setSchema(this.schema);
      this.writer = OrcFile.createWriter(this.output, options);
      this.batch = this.schema.createRowBatch();
   }

   public void append(DataBatch.DataRow row) throws IOException {
      if (row != null && row.cells != null && !row.cells.isEmpty()) {
         ++this.batch.size;
         ++this.index;
         LongColumnVector operation = (LongColumnVector)this.batch.cols[0];
         operation.vector[this.index] = 0L;
         StructColumnVector rowVector = (StructColumnVector)this.batch.cols[5];
         ColumnVector[] fields = rowVector.fields;

         for(int i = 0; i < row.cells.size(); ++i) {
            DataBatch.DataCell cell = (DataBatch.DataCell)row.cells.get(i);
            DataBatch.DataValue value = cell.value;
            ColumnVector field = fields[i];
            ColumnInfo col = (ColumnInfo)this.cols.get(i);
            value.type = col.getDatatype();
            switch (col.getDatatype().toLowerCase()) {
               case "string":
               case "char":
               case "varchar":
                  BytesColumnVector data = (BytesColumnVector)field;
                  String sval = (String)value.value;
                  byte[] buffer = sval == null ? null : sval.getBytes(StandardCharsets.UTF_8);
                  data.setVal(this.index, buffer);
                  break;
               case "float":
               case "double":
                  DoubleColumnVector dataNew = (DoubleColumnVector)fields[i];
                  Double dval = (Double)value.value;
                  dataNew.vector[this.index] = dval;
                  break;
               case "decimal":
                  DecimalColumnVector dataNew1 = (DecimalColumnVector)fields[i];
                  HiveDecimalWritable dvalNew2 = (HiveDecimalWritable)value.value;
                  dataNew1.vector[this.index] = dvalNew2;
                  break;
               case "int":
               case "integer":
               case "long":
                  LongColumnVector dataNew3 = (LongColumnVector)fields[i];
                  long ival = (Long)value.value;
                  dataNew3.vector[this.index] = ival;
                  break;
               case "date":
               case "datetime":
               case "timestamp":
                  LongColumnVector dataNew4 = (LongColumnVector)fields[i];
                  long ivalNew5 = (Long)value.value;
                  dataNew4.vector[this.index] = ivalNew5;
            }
         }

         this.writer.addRowBatch(this.batch);
      }
   }

   public void close() throws IOException {
      this.writer.close();
   }

   public byte[] getBytes() {
      return new byte[0];
   }

   public static void main(String[] args) throws IOException {
      List<ColumnInfo> cols = new ArrayList();
      ColumnInfo col = new ColumnInfo();
      col.setDatatype("string");
      col.setColumnname("id");
      cols.add(col);
      col = new ColumnInfo();
      col.setDatatype("string");
      col.setColumnname("name");
      cols.add(col);
      col = new ColumnInfo();
      col.setDatatype("long");
      col.setColumnname("age");
      cols.add(col);
      col = new ColumnInfo();
      col.setDatatype("date");
      col.setColumnname("birth");
      cols.add(col);
      new Configuration();
      Path path = new Path("file:///d:/b.orc");
      OrcFileReader reader = new OrcFileReader(path, cols);
      reader.load();
      DataBatch next = null;

      while((next = reader.next()) != null) {
         for(DataBatch.DataRow row : next.rows) {
            List<String> datas = (List)row.cells.stream().map((x) -> x.value.value + "").collect(Collectors.toList());
            System.out.println(String.join(",", datas));
         }
      }

      Path outputPath = new Path("file:///d:/c.orc");
      OrcFileWriter writer = new OrcFileWriter(outputPath, reader, cols);
      writer.initialize();
      DataBatch.DataRow dr = new DataBatch.DataRow();
      DataBatch.DataCell cell = new DataBatch.DataCell();
      DataBatch.DataValue value = new DataBatch.DataValue();
      value.value = "5";
      cell.value = value;
      dr.cells.add(cell);
      cell = new DataBatch.DataCell();
      value = new DataBatch.DataValue();
      value.value = "顶顶顶的";
      cell.value = value;
      dr.cells.add(cell);
      cell = new DataBatch.DataCell();
      value = new DataBatch.DataValue();
      value.value = 32L;
      cell.value = value;
      dr.cells.add(cell);
      cell = new DataBatch.DataCell();
      value = new DataBatch.DataValue();
      value.value = System.currentTimeMillis();
      cell.value = value;
      dr.cells.add(cell);
      writer.append(dr);
      writer.close();
   }
}
