package org.iknowledge.hbase.client;

import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Admin;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Durability;
import org.apache.hadoop.hbase.client.Get;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.ResultScanner;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.io.compress.Compression;
import org.apache.hadoop.hbase.io.encoding.DataBlockEncoding;
import org.apache.hadoop.hbase.regionserver.BloomType;
import org.apache.hadoop.hbase.util.Bytes;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.stream.IntStream;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static java.util.stream.Collectors.toList;
import static org.assertj.core.api.Assertions.assertThat;

/**
 * {@link Batch} test class
 */
@Test( suiteName = "HBase batch test" )
public class BatchIt {

    private static final String TEST_CF = "ts";
    private static final byte[] TEST_CF_BYTES = Bytes.toBytes("ts");
    private Table testTable;
    private TableName tableName;
    private Admin admin;

    /**
     * test setup
     * @throws Exception
     */
    @BeforeClass
    public void setUp() throws Exception {
        final String name = "testTable" + new Random().nextInt(5000);
        tableName = TableName.valueOf(name);
        final HConfig config = HConfig.newBuilder()
                                      .retryCount(5)
                                      .retryBackoff(3000)
                                      .scanBatchSize(50)
                                      .scanCacheSize(50)
                                      .zkQuorum("localhost:2181")
                                      .connectionThreads(4)
                                      .metaLookupThreads(2)
                                      .metaOperationTimeout(5000)
                                      .metricsEnabled(true)
                                      .operationTimeout(5000)
                                      .perRegionMaxTasks(20)
                                      .perServerMaxTasks(40)
                                      .rpcTimeout(7000)
                                      .scannerTimeout(20000)
                                      .threadPoolMaxTasks(100)
                                      .zkSessionTimeout(15000)
                                      .znode("/hbase")
                                      .build();
        final Connection connection = ConnectionFactory.createConnection(config.asConfiguration());

        this.testTable = connection.getTable(this.tableName);
        admin = connection.getAdmin();
        final HColumnDescriptor cfTestDesc
                = new HColumnDescriptor(TEST_CF)
                .setBloomFilterType(BloomType.ROW)
                .setCompactionCompressionType(Compression.Algorithm.SNAPPY)
                .setCompressionType(Compression.Algorithm.SNAPPY)
                .setDataBlockEncoding(DataBlockEncoding.PREFIX)
                .setVersions(1, 1);
        final HTableDescriptor descriptor
                = new HTableDescriptor(tableName)
                .setCompactionEnabled(true)
                .setDurability(Durability.SYNC_WAL)
                .addFamily(cfTestDesc);
        admin.createTable(descriptor);
    }

    /**
     * Tear down test
     * @throws Exception
     */
    @AfterClass
    public void tearDown() throws Exception {
        System.out.println("Close HBase test table");
        admin.disableTable(tableName);
        admin.deleteTable(tableName);
        testTable.close();
    }

    @AfterMethod
    public void tearDownAfterTest() throws Exception {
        admin.disableTable(tableName);
        admin.truncateTable(tableName, true);
    }

    @Test( description = "Test invalid batch parameter: batch size",
            expectedExceptions = IllegalArgumentException.class )
    public void testInvalidBatchSize() throws Exception {
        Batch.newBuilder().withBatchSize(-20)
             .withObjectCollection(Collections.emptyList())
             .withMapper(o -> new Put(new byte[0]))
             .withTable(testTable)
             .build();
    }

    @Test( description = "Test invalid batch parameter: empty object collection",
            expectedExceptions = NullPointerException.class )
    public void testInvalidObjectCollection() throws Exception {
        Batch.newBuilder().withBatchSize(10)
             .withObjectCollection(( Iterator<Object> ) null)
             .withMapper(o -> new Put(new byte[0]))
             .withTable(testTable)
             .build();
    }

    @Test( description = "Test invalid batch parameter: null mapper",
            expectedExceptions = NullPointerException.class )
    public void testInvalidMapper() throws Exception {
        Batch.newBuilder().withBatchSize(10)
             .withObjectCollection(Collections.emptyList())
             .withMapper(null)
             .withTable(testTable)
             .build();
    }

    @Test( description = "Test invalid batch parameter: HBase table instance",
            expectedExceptions = NullPointerException.class )
    public void testInvalidHbaseTable() throws Exception {
        Batch.newBuilder().withBatchSize(10)
             .withObjectCollection(Collections.emptyList())
             .withMapper(o -> new Put(new byte[0]))
             .withTable(null)
             .build();
    }

    @Test(description = "Test dynamic batch")
    public void testDynamicBatch() throws Exception {
        final byte[] dynBatchRow = Bytes.toBytes("dyn_batch_row");
        final int batchSize = 10;
        final DynamicBatch<Integer> batch
                = DynamicBatch.<Integer>newBuilder()
                              .withBatchSize(batchSize)
                              .withMapper(o -> new Put(dynBatchRow)
                                                       .addImmutable(TEST_CF_BYTES,
                                                                     Bytes.toBytes("dynQf" + o),
                                                                     Bytes.toBytes(o)))
                              .withTable(testTable)
                              .build();
        IntStream.range(0, batchSize).forEach(value -> {
            try {
                batch.append(value);
            }
            catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
        assertThat(testTable.get(new Get(dynBatchRow).addFamily(TEST_CF_BYTES)).listCells())
                .isNotNull()
                .hasSize(batchSize);

        batch.append(batchSize + 1);
        batch.finish();

        assertThat(testTable.get(new Get(dynBatchRow).addFamily(TEST_CF_BYTES)).listCells())
                .isNotNull()
                .hasSize(batchSize + 1);
    }

    /**
     * Object collection data provider for batch operation.
     *
     * @return
     */
    @DataProvider( name = "objCollection" )
    public static Object[][] objCollection() {
        final List<String> list1 = Stream.of("v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8")
                                         .collect(toList());
        final List<String> list2 = Stream.of("v1", "v2", "v3", "v4", "v5", "v6", "v7")
                                         .collect(toList());
        final List<String> list3 = Stream.of("v1", "v2", "v3", "v4", "v5", "v6")
                                         .collect(toList());
        final List<String> list4 = Stream.of("v1")
                                         .collect(toList());
        final List<String> list5 = Stream.of("v1", "v2")
                                         .collect(toList());
        final List<String> list6 = Collections.emptyList();
        return new Object[][] {
                new Object[] { list1 },
                new Object[] { list2 },
                new Object[] { list3 },
                new Object[] { list4 },
                new Object[] { list5 },
                new Object[] { list6 },
                };
    }

    @Test( description = "Create valid batch instance and call batch method",
            dataProvider = "objCollection" )
    public void testCreateValidBatchInstanceAndCallBatch( List<String> strCollection )
            throws Exception {

        final byte[] valQualifier = Bytes.toBytes("val");
        Batch.<String>newBuilder()
                .withBatchSize(3)
                .withObjectCollection(strCollection)
                .withMapper(string -> {
                    final byte[] key = Bytes.toBytes(string);
                    return new Put(key).addImmutable(TEST_CF_BYTES,
                                                     valQualifier,
                                                     key);
                })
                .withTable(testTable)
                .build()
                .call();

        final ResultScanner scanner = testTable.getScanner(TEST_CF_BYTES, valQualifier);
        final Integer rowCount
                = StreamSupport.stream(scanner.spliterator(), false)
                               .reduce(0,
                                       ( curVal, hresult ) -> curVal + hresult.size(),
                                       ( val1, val2 ) -> val1 + val2);
        assertThat(rowCount).isEqualTo(strCollection.size());
    }

    @Test( description = "Create valid batch instance and call batch method on Iterator",
            dataProvider = "objCollection" )
    public void testCreateValidBatchInstanceAndCallBatchOnIterator( List<String> strCollection )
            throws Exception {

        final byte[] valQualifier = Bytes.toBytes("val");
        final Batch<String> batch = Batch.<String>newBuilder()
                                            .withBatchSize(3)
                                            .withObjectCollection(strCollection.iterator())
                                            .withMapper(string -> {
                                                final byte[] key = Bytes.toBytes(string);
                                                return new Put(key).addImmutable(TEST_CF_BYTES,
                                                                                 valQualifier,
                                                                                 key);
                                            })
                                            .withTable(testTable)
                                            .build();
        // batch on iterator instance perform only one time
        // other calls must have no effect(and must not fail)
        batch.call();
        batch.call();
        batch.call();

        final ResultScanner scanner = testTable.getScanner(TEST_CF_BYTES, valQualifier);
        final Integer rowCount
                = StreamSupport.stream(scanner.spliterator(), false)
                               .reduce(0,
                                       ( curVal, hresult ) -> curVal + hresult.size(),
                                       ( val1, val2 ) -> val1 + val2);
        assertThat(rowCount).isEqualTo(strCollection.size());
    }
}