/*
 * 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 com.pinterest.secor.io.impl;

import static org.junit.Assert.assertArrayEquals;

import java.util.HashMap;
import java.util.Map;

import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.parquet.hadoop.ParquetWriter;

import org.apache.thrift.TDeserializer;
import org.apache.thrift.TSerializer;
import org.apache.thrift.protocol.TCompactProtocol;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.powermock.core.classloader.annotations.PowerMockIgnore;
import org.powermock.modules.junit4.PowerMockRunner;

import com.google.common.io.Files;
import com.pinterest.secor.common.LogFilePath;
import com.pinterest.secor.common.SecorConfig;
import com.pinterest.secor.io.FileReader;
import com.pinterest.secor.io.FileWriter;
import com.pinterest.secor.io.KeyValue;
import com.pinterest.secor.thrift.UnitTestMessage;
import com.pinterest.secor.util.ParquetUtil;
import com.pinterest.secor.util.ReflectionUtil;

import junit.framework.TestCase;

@RunWith(PowerMockRunner.class)
@PowerMockIgnore({
    "com.ctc.wstx.stax.*",
    "com.ctc.wstx.io.*",
    "com.sun.*",
    "com.sun.org.apache.xalan.",
    "com.sun.org.apache.xerces.",
    "com.sun.xml.internal.stream.*",
    "javax.activation.*",
    "javax.management.",
    "javax.xml.",
    "javax.xml.stream.*",
    "javax.security.auth.login.*",
    "javax.security.auth.spi.*",
    "org.apache.hadoop.security.*",
    "org.codehaus.stax2.*",
    "org.w3c.",
    "org.xml.",
    "org.w3c.dom."})
public class ThriftParquetFileReaderWriterFactoryTest extends TestCase {

    private SecorConfig config;

    @Override
    public void setUp() throws Exception {
        config = Mockito.mock(SecorConfig.class);
    }

    @Test
    public void testThriftParquetReadWriteRoundTrip() throws Exception {
        Map<String, String> classPerTopic = new HashMap<String, String>();
        classPerTopic.put("test-pb-topic", UnitTestMessage.class.getName());
        Mockito.when(config.getThriftMessageClassPerTopic()).thenReturn(classPerTopic);
        Mockito.when(config.getFileReaderWriterFactory())
                .thenReturn(ThriftParquetFileReaderWriterFactory.class.getName());
        Mockito.when(config.getThriftProtocolClass())
        .thenReturn(TCompactProtocol.class.getName());
        Mockito.when(ParquetUtil.getParquetBlockSize(config))
                .thenReturn(ParquetWriter.DEFAULT_BLOCK_SIZE);
        Mockito.when(ParquetUtil.getParquetPageSize(config))
                .thenReturn(ParquetWriter.DEFAULT_PAGE_SIZE);
        Mockito.when(ParquetUtil.getParquetEnableDictionary(config))
                .thenReturn(ParquetWriter.DEFAULT_IS_DICTIONARY_ENABLED);
        Mockito.when(ParquetUtil.getParquetValidation(config))
                .thenReturn(ParquetWriter.DEFAULT_IS_VALIDATING_ENABLED);


        LogFilePath tempLogFilePath = new LogFilePath(Files.createTempDir().toString(), "test-pb-topic",
                new String[] { "part-1" }, 0, 1, 23232, ".log");

        FileWriter fileWriter = ReflectionUtil.createFileWriter(config.getFileReaderWriterFactory(), tempLogFilePath,
                null, config);

        UnitTestMessage msg1 = new UnitTestMessage().setRequiredField("abc").setTimestamp(1467176315L);
        UnitTestMessage msg2 = new UnitTestMessage().setRequiredField("XYZ").setTimestamp(1467176344L);

        TSerializer serializer = new TSerializer(new TCompactProtocol.Factory());
        KeyValue kv1 = new KeyValue(23232, serializer.serialize(msg1));
        KeyValue kv2 = new KeyValue(23233, serializer.serialize(msg2));
        fileWriter.write(kv1);
        fileWriter.write(kv2);
        fileWriter.close();

        FileReader fileReader = ReflectionUtil.createFileReader(config.getFileReaderWriterFactory(), tempLogFilePath,
                null, config);
        TDeserializer deserializer = new TDeserializer(new TCompactProtocol.Factory());
        
        KeyValue kvout = fileReader.next();
        assertEquals(kv1.getOffset(), kvout.getOffset());
        assertArrayEquals(kv1.getValue(), kvout.getValue());
        UnitTestMessage actual = new UnitTestMessage();
        deserializer.deserialize(actual, kvout.getValue());
        assertEquals(msg1.getRequiredField(), actual.getRequiredField());

        kvout = fileReader.next();
        assertEquals(kv2.getOffset(), kvout.getOffset());
        assertArrayEquals(kv2.getValue(), kvout.getValue());
        actual = new UnitTestMessage();
        deserializer.deserialize(actual, kvout.getValue());
        assertEquals(msg2.getRequiredField(), actual.getRequiredField());
    }
}