/**
 * 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.apex.malhar.kudu;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.powermock.api.mockito.PowerMockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.apex.malhar.lib.util.KryoCloneUtils;
import org.apache.kudu.ColumnSchema;
import org.apache.kudu.Schema;
import org.apache.kudu.Type;
import org.apache.kudu.client.CreateTableOptions;
import org.apache.kudu.client.KuduClient;
import org.apache.kudu.client.KuduException;
import org.apache.kudu.client.KuduPredicate;
import org.apache.kudu.client.KuduScanToken;
import org.apache.kudu.client.KuduScanner;
import org.apache.kudu.client.KuduSession;
import org.apache.kudu.client.KuduTable;
import org.apache.kudu.client.PartialRow;
import org.apache.kudu.client.RowResult;
import org.apache.kudu.client.RowResultIterator;

import static org.powermock.api.mockito.PowerMockito.when;

/** Common test code for all kudu operators
 */
public class KuduClientTestCommons
{

  private static final transient Logger LOG = LoggerFactory.getLogger(KuduClientTestCommons.class);


  protected static final String tableName = "unittests";

  protected static KuduClient kuduClient;

  protected static KuduTable kuduTable;

  protected static Schema schemaForUnitTests;

  protected static String kuduMasterAddresses = "192.168.1.41:7051";

  public static final int SPLIT_COUNT_FOR_INT_ROW_KEY = 5;

  public static final int HASH_BUCKETS_SIZE_FOR_ALL_HASH_COL = 2;

  public static final int TOTAL_KUDU_TABLETS_FOR_UNITTEST_TABLE = 12;

  public static boolean tableInitialized = false;

  public static Object objectForLocking = new Object();

  protected static Map<String,ColumnSchema> columnDefs = new HashMap<>();


  public static void setup() throws Exception
  {
    kuduClient = getClientHandle();
    if (kuduClient.tableExists(tableName)) {
      kuduClient.deleteTable(tableName);
    }
    createTestTable(tableName,kuduClient);
    kuduTable = kuduClient.openTable(tableName);
    tableInitialized = true;
  }


  public static void shutdown() throws Exception
  {
    kuduClient.close();
  }


  private static KuduClient getClientHandle() throws Exception
  {
    KuduClient.KuduClientBuilder builder = new KuduClient.KuduClientBuilder(kuduMasterAddresses);
    KuduClient client = builder.build();
    return client;
  }

  public static ApexKuduConnection.ApexKuduConnectionBuilder getConnectionConfigForTable()
  {
    ApexKuduConnection.ApexKuduConnectionBuilder connectionBuilder = new ApexKuduConnection.ApexKuduConnectionBuilder();
    return connectionBuilder.withAPossibleMasterHostAs(kuduMasterAddresses).withTableName(tableName);
  }

  public static Schema buildSchemaForUnitTestsTable() throws Exception
  {
    if (schemaForUnitTests != null) {
      return schemaForUnitTests;
    }
    List<ColumnSchema> columns = new ArrayList<>();
    ColumnSchema intRowKeyCol = new ColumnSchema.ColumnSchemaBuilder("introwkey", Type.INT32)
        .key(true)
        .build();
    columns.add(intRowKeyCol);
    columnDefs.put("introwkey",intRowKeyCol);
    ColumnSchema stringRowKeyCol = new ColumnSchema.ColumnSchemaBuilder("stringrowkey", Type.STRING)
        .key(true)
        .build();
    columns.add(stringRowKeyCol);
    columnDefs.put("stringrowkey",stringRowKeyCol);
    ColumnSchema timestampRowKey = new ColumnSchema.ColumnSchemaBuilder("timestamprowkey", Type.UNIXTIME_MICROS)
        .key(true)
        .build();
    columns.add(timestampRowKey);
    columnDefs.put("timestamprowkey",timestampRowKey);
    ColumnSchema longData = new ColumnSchema.ColumnSchemaBuilder("longdata", Type.INT64)
        .nullable(true)
        .build();
    columns.add(longData);
    columnDefs.put("longdata",longData);
    ColumnSchema stringData = new ColumnSchema.ColumnSchemaBuilder("stringdata", Type.STRING)
        .nullable(true)
        .build();
    columns.add(stringData);
    columnDefs.put("stringdata",stringData);
    ColumnSchema timestampdata = new ColumnSchema.ColumnSchemaBuilder("timestampdata", Type.UNIXTIME_MICROS)
        .nullable(true)
        .build();
    columns.add(timestampdata);
    columnDefs.put("timestampdata",timestampdata);
    ColumnSchema binarydata = new ColumnSchema.ColumnSchemaBuilder("binarydata", Type.BINARY)
        .nullable(true)
        .build();
    columns.add(binarydata);
    columnDefs.put("binarydata",binarydata);
    ColumnSchema floatdata = new ColumnSchema.ColumnSchemaBuilder("floatdata", Type.FLOAT)
        .nullable(true)
        .build();
    columns.add(floatdata);
    columnDefs.put("floatdata",floatdata);
    ColumnSchema booldata = new ColumnSchema.ColumnSchemaBuilder("booldata", Type.BOOL)
        .nullable(true)
        .build();
    columns.add(booldata);
    columnDefs.put("booldata",booldata);
    schemaForUnitTests = new Schema(columns);
    return schemaForUnitTests;
  }

  public static void createTestTable(String tableName, KuduClient client) throws Exception
  {
    List<String> rangeKeys = new ArrayList<>();
    rangeKeys.add("introwkey");
    List<String> hashPartitions = new ArrayList<>();
    hashPartitions.add("stringrowkey");
    hashPartitions.add("timestamprowkey");
    CreateTableOptions thisTableOptions = new CreateTableOptions()
        .setNumReplicas(1)
        .addHashPartitions(hashPartitions,HASH_BUCKETS_SIZE_FOR_ALL_HASH_COL)
        .setRangePartitionColumns(rangeKeys);
    int stepsize = Integer.MAX_VALUE / SPLIT_COUNT_FOR_INT_ROW_KEY;
    int splitBoundary = stepsize;
    Schema schema = buildSchemaForUnitTestsTable();
    for ( int i = 0; i < SPLIT_COUNT_FOR_INT_ROW_KEY; i++) {
      PartialRow splitRowBoundary = schema.newPartialRow();
      splitRowBoundary.addInt("introwkey",splitBoundary);
      thisTableOptions = thisTableOptions.addSplitRow(splitRowBoundary);
      splitBoundary += stepsize;
    }
    try {
      client.createTable(tableName, schema,thisTableOptions);
    } catch (KuduException e) {
      LOG.error("Error while creating table for unit tests " + e.getMessage(), e);
      throw e;
    }

  }

  protected void lookUpAndPopulateRecord(UnitTestTablePojo keyInfo) throws Exception
  {
    KuduScanner scanner = kuduClient.newScannerBuilder(kuduTable)
        .addPredicate(KuduPredicate.newComparisonPredicate(columnDefs.get("introwkey"),
        KuduPredicate.ComparisonOp.EQUAL,keyInfo.getIntrowkey()))
        .addPredicate(KuduPredicate.newComparisonPredicate(columnDefs.get("stringrowkey"),
        KuduPredicate.ComparisonOp.EQUAL,keyInfo.getStringrowkey()))
        .addPredicate(KuduPredicate.newComparisonPredicate(columnDefs.get("timestamprowkey"),
        KuduPredicate.ComparisonOp.EQUAL,keyInfo.getTimestamprowkey()))
        .build();
    RowResultIterator rowResultItr = scanner.nextRows();
    while (rowResultItr.hasNext()) {
      RowResult thisRow = rowResultItr.next();
      keyInfo.setFloatdata(thisRow.getFloat("floatdata"));
      keyInfo.setBooldata(thisRow.getBoolean("booldata"));
      keyInfo.setBinarydata(thisRow.getBinary("binarydata"));
      keyInfo.setLongdata(thisRow.getLong("longdata"));
      keyInfo.setTimestampdata(thisRow.getLong("timestampdata"));
      keyInfo.setStringdata("stringdata");
      break;
    }
  }

  public ApexKuduConnection buildMockWiring(AbstractKuduInputOperator abstractKuduInputOperator,
      int numScanTokens) throws Exception
  {
    ApexKuduConnection mockedConnectionHandle = PowerMockito.mock(ApexKuduConnection.class);
    ApexKuduConnection.ApexKuduConnectionBuilder mockedConnectionHandleBuilder = PowerMockito.mock(
        ApexKuduConnection.ApexKuduConnectionBuilder.class);
    KuduClient mockedClient = PowerMockito.mock(KuduClient.class);
    KuduSession mockedKuduSession = PowerMockito.mock(KuduSession.class);
    KuduTable mockedKuduTable = PowerMockito.mock(KuduTable.class);
    KuduScanToken.KuduScanTokenBuilder mockedScanTokenBuilder = PowerMockito.mock(
        KuduScanToken.KuduScanTokenBuilder.class);
    List<KuduScanToken> mockedScanTokens = new ArrayList<>();
    int scanTokensToBuild = numScanTokens;
    for (int i = 0; i < scanTokensToBuild; i++) {
      mockedScanTokens.add(PowerMockito.mock(KuduScanToken.class));
    }
    PowerMockito.mockStatic(KryoCloneUtils.class);
    when(KryoCloneUtils.cloneObject(abstractKuduInputOperator)).thenReturn(abstractKuduInputOperator);
    //wire the mocks
    when(abstractKuduInputOperator.getApexKuduConnectionInfo()).thenReturn(mockedConnectionHandleBuilder);
    when(mockedConnectionHandle.getKuduClient()).thenReturn(mockedClient);
    when(mockedClient.newSession()).thenReturn(mockedKuduSession);
    when(mockedConnectionHandle.getKuduTable()).thenReturn(mockedKuduTable);
    when(mockedConnectionHandle.getKuduSession()).thenReturn(mockedKuduSession);
    when(mockedConnectionHandle.getBuilderForThisConnection()).thenReturn(mockedConnectionHandleBuilder);
    when(mockedClient.openTable(tableName)).thenReturn(mockedKuduTable);
    when(mockedConnectionHandleBuilder.build()).thenReturn(mockedConnectionHandle);
    when(mockedKuduTable.getSchema()).thenReturn(schemaForUnitTests);
    when(mockedClient.newScanTokenBuilder(mockedKuduTable)).thenReturn(mockedScanTokenBuilder);
    when(mockedScanTokenBuilder.build()).thenReturn(mockedScanTokens);
    return mockedConnectionHandle;
  }

  public static String getKuduMasterAddresses()
  {
    return kuduMasterAddresses;
  }

  public static void setKuduMasterAddresses(String kuduMasterAddresses)
  {
    KuduClientTestCommons.kuduMasterAddresses = kuduMasterAddresses;
  }
}