package com.pinterest.secor.io.impl;

import java.io.IOException;

import com.pinterest.secor.common.SecorSchemaRegistryClient;
import org.apache.avro.generic.GenericRecord;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.parquet.avro.AvroParquetReader;
import org.apache.parquet.avro.AvroParquetWriter;
import org.apache.parquet.hadoop.ParquetReader;
import org.apache.parquet.hadoop.ParquetWriter;
import org.apache.parquet.hadoop.metadata.CompressionCodecName;

import com.pinterest.secor.common.LogFilePath;
import com.pinterest.secor.common.SecorConfig;
import com.pinterest.secor.io.FileReader;
import com.pinterest.secor.io.FileReaderWriterFactory;
import com.pinterest.secor.io.FileWriter;
import com.pinterest.secor.io.KeyValue;
import com.pinterest.secor.util.ParquetUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * Implementation for reading/writing avro messages to/from Parquet files.
 *
 * @author Michael Spector (spektom@gmail.com)
 */
public class AvroParquetFileReaderWriterFactory implements FileReaderWriterFactory {

    private static final Logger LOG = LoggerFactory.getLogger(AvroParquetFileReaderWriterFactory.class);
    protected final int blockSize;
    protected final int pageSize;
    protected final boolean enableDictionary;
    protected final boolean validating;
    protected final SecorSchemaRegistryClient schemaRegistryClient;

    public AvroParquetFileReaderWriterFactory(SecorConfig config) {
        blockSize = ParquetUtil.getParquetBlockSize(config);
        pageSize = ParquetUtil.getParquetPageSize(config);
        enableDictionary = ParquetUtil.getParquetEnableDictionary(config);
        validating = ParquetUtil.getParquetValidation(config);
        schemaRegistryClient = SecorSchemaRegistryClient.getInstance();
        schemaRegistryClient.init(config);
    }

    @Override
    public FileReader BuildFileReader(LogFilePath logFilePath, CompressionCodec codec) throws Exception {
        return new AvroParquetFileReader(logFilePath, codec);
    }

    @Override
    public FileWriter BuildFileWriter(LogFilePath logFilePath, CompressionCodec codec) throws Exception {
        return new AvroParquetFileWriter(logFilePath, codec);
    }

    protected class AvroParquetFileReader implements FileReader {

        private ParquetReader reader;
        private long offset;

        public AvroParquetFileReader(LogFilePath logFilePath, CompressionCodec codec) throws IOException {
            Path path = new Path(logFilePath.getLogFilePath());
            reader = AvroParquetReader.builder(path).build();
            offset = logFilePath.getOffset();
        }

        @Override
        public KeyValue next() throws IOException {
//            Builder messageBuilder = (Builder) reader.read();
//            if (messageBuilder != null) {
//                return new KeyValue(offset++, messageBuilder.build().toByteArray());
//            }
            return null;
        }

        @Override
        public void close() throws IOException {
            reader.close();
        }
    }

    protected class AvroParquetFileWriter implements FileWriter {

        private ParquetWriter writer;
        private String topic;

        public AvroParquetFileWriter(LogFilePath logFilePath, CompressionCodec codec) throws IOException {
            Path path = new Path(logFilePath.getLogFilePath());
            LOG.info("Creating Brand new Writer for path {}", path);
            topic = logFilePath.getTopic();
            // Not setting blockSize, pageSize, enableDictionary, and validating
            writer = AvroParquetWriter.builder(path)
                    .withSchema(schemaRegistryClient.getSchema(topic))
                    .withCompressionCodec(CompressionCodecName.GZIP)
                    .build();
        }

        @Override
        public long getLength() throws IOException {
            return writer.getDataSize();
        }

        @Override
        public void write(KeyValue keyValue) throws IOException {
            //GenericRecord record = avroUtil.decodeMessage(keyValue.getValue());
            GenericRecord record = schemaRegistryClient.decodeMessage(topic, keyValue.getValue());
            LOG.trace("Writing record {}", record);
            if (record != null){
                writer.write(record);
            }
        }

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


@@ -50,127 +50,390 @@
 .io.
-IOException;%0A%0Aimport com.pinterest.secor.common.SecorSchemaRegistryClient;%0Aimport org.apache.avro.generic.GenericRecord
+ByteArrayOutputStream;%0Aimport java.io.IOException;%0Aimport java.nio.ByteBuffer;%0A%0Aimport com.google.protobuf.Message;%0Aimport com.pinterest.secor.common.SecorSchemaRegistryClient;%0Aimport org.apache.avro.Schema;%0Aimport org.apache.avro.generic.GenericRecord;%0Aimport org.apache.avro.io.Encoder;%0Aimport org.apache.avro.io.EncoderFactory;%0Aimport org.apache.avro.specific.SpecificDatumWriter
 ;%0Aim
@@ -2627,21 +2627,95 @@
 etReader
- read
+%3CGenericRecord%3E reader;%0A        private SpecificDatumWriter%3CGenericRecord%3E writ
 er;%0A    
@@ -2903,32 +2903,150 @@
 LogFilePath());%0A
+            String topic = logFilePath.getTopic();%0A            Schema schema = schemaRegistryClient.getSchema(topic);%0A
             read
@@ -3064,24 +3064,39 @@
 rquetReader.
+%3CGenericRecord%3E
 builder(path
@@ -3098,32 +3098,86 @@
 (path).build();%0A
+            writer = new SpecificDatumWriter(schema);%0A
             offs
@@ -3287,18 +3287,16 @@
 ption %7B%0A
-//
         
@@ -3304,41 +3304,181 @@
    B
-uilder messageBuilder = (Builder)
+yteArrayOutputStream out = new ByteArrayOutputStream();%0A            Encoder encoder = EncoderFactory.get().directBinaryEncoder(out, null);%0A            GenericRecord record =
  rea
@@ -3489,18 +3489,16 @@
 read();%0A
-//
         
@@ -3509,36 +3509,244 @@
 if (
-messageBuilder != null) %7B%0A//
+record != null) %7B%0A                writer.write(record, encoder);%0A                encoder.flush();%0A                ByteBuffer serialized = ByteBuffer.allocate(out.toByteArray().length);%0A                serialized.put(out.toByteArray());%0A
     
@@ -3791,38 +3791,20 @@
 ++, 
-messageBuilder.build().toByteA
+serialized.a
 rray
@@ -3808,18 +3808,16 @@
 ray());%0A
-//
         

add reader to parquet factory