/**
 * 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.contrib.cassandra;


import java.util.EnumSet;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.datastax.driver.core.DataType;
import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.Session;

/**
 * Used to generate CQL strings that can be used to generate prepared statements.
 *
 * @since 3.6.0
 */
public class CassandraPreparedStatementGenerator
{

  private Set<String> pkColumnNames;

  private Set<String> counterColumns;

  private Set<String> listColumns;

  private Set<String> mapColumns;

  private Set<String> setColumns;

  private Map<String, DataType> columnDefinitions;

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

  public static final String TTL_PARAM_NAME = "ttl";

  public CassandraPreparedStatementGenerator(Set<String> pkColumnNames, Set<String> counterColumns,
      Set<String> listColumns, Set<String> mapColumns, Set<String> setColumns,
      Map<String, DataType> columnDefinitions)
  {
    this.pkColumnNames = pkColumnNames;
    this.counterColumns = counterColumns;
    this.listColumns = listColumns;
    this.mapColumns = mapColumns;
    this.setColumns = setColumns;
    this.columnDefinitions = columnDefinitions;
  }

  public void generatePreparedStatements(Session session,Map<Long, PreparedStatement> preparedStatementTypes,
      String keyspaceName,String tableName)
  {

    Map<Long, String> stringsWithoutPKAndExistsClauses = generatePreparedStatementsQueryStrings(keyspaceName,tableName);
    String ifExistsClause = " IF EXISTS";
    Map<Long, String> finalSetOfQueryStrings = new HashMap<>();
    for (Long currentIndexPos : stringsWithoutPKAndExistsClauses.keySet()) {
      StringBuilder aQueryStub = new StringBuilder(stringsWithoutPKAndExistsClauses.get(currentIndexPos));
      buildWhereClauseForPrimaryKeys(aQueryStub);
      finalSetOfQueryStrings.put(currentIndexPos +
          getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
          AbstractUpsertOutputOperator.OperationContext.IF_EXISTS_CHECK_ABSENT)),
          aQueryStub.toString());
      if (counterColumns.size() == 0) {
        // IF exists cannot be used for counter column tables
        finalSetOfQueryStrings.put(currentIndexPos +
            getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
            AbstractUpsertOutputOperator.OperationContext.IF_EXISTS_CHECK_PRESENT)),
            aQueryStub.toString() + ifExistsClause);
      }
    }
    for (Long currentIndexPos : finalSetOfQueryStrings.keySet()) {
      String currentQueryStr = finalSetOfQueryStrings.get(currentIndexPos);
      LOG.info("Registering query support for " + currentQueryStr);
      PreparedStatement preparedStatementForThisQuery = session.prepare(currentQueryStr);
      preparedStatementTypes.put(currentIndexPos, preparedStatementForThisQuery);
    }
  }

  private void buildWhereClauseForPrimaryKeys(final StringBuilder queryExpression)
  {
    queryExpression.append(" WHERE ");
    int count = 0;
    for (String pkColName : pkColumnNames) {
      if (count > 0) {
        queryExpression.append(" AND ");
      }
      count += 1;
      queryExpression.append(" ").append(pkColName).append(" = :").append(pkColName);
    }
  }


  private void buildQueryStringForTTLSetCollectionsAppendAndListPrepend(StringBuilder updateQueryRoot,
      String ttlSetString, Map<Long,String> queryStrings)
  {
    // TTL set , Collections Append , List prepend
    StringBuilder queryExpTTLSetCollAppendListPrepend = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLSetCollAppendListPrepend.append(ttlSetString);
    buildNonPKColumnsExpression(queryExpTTLSetCollAppendListPrepend,
        UpsertExecutionContext.ListPlacementStyle.PREPEND_TO_EXISTING_LIST,
        UpsertExecutionContext.CollectionMutationStyle.ADD_TO_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_APPEND,
        AbstractUpsertOutputOperator.OperationContext.LIST_PREPEND
    )), queryExpTTLSetCollAppendListPrepend.toString());
  }

  private void buildQueryStringForTTLSetCollectionsAppendAndListAppend(StringBuilder updateQueryRoot,
      String ttlSetString,Map<Long,String> queryStrings)
  {
    // TTL set , Collections Append , List append
    StringBuilder queryExpTTLSetCollAppendListAppend = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLSetCollAppendListAppend.append(ttlSetString);
    buildNonPKColumnsExpression(queryExpTTLSetCollAppendListAppend,
        UpsertExecutionContext.ListPlacementStyle.APPEND_TO_EXISTING_LIST,
        UpsertExecutionContext.CollectionMutationStyle.ADD_TO_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_APPEND,
        AbstractUpsertOutputOperator.OperationContext.LIST_APPEND
    )), queryExpTTLSetCollAppendListAppend.toString());
  }

  private void buildQueryStringForTTLSetCollectionsRemove(StringBuilder updateQueryRoot,
      String ttlSetString,Map<Long,String> queryStrings)
  {
    // TTL set , Collections Remove
    StringBuilder queryExpTTLSetCollRemove = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLSetCollRemove.append(ttlSetString);
    buildNonPKColumnsExpression(queryExpTTLSetCollRemove,
        UpsertExecutionContext.ListPlacementStyle.APPEND_TO_EXISTING_LIST, // Just in case user sets it
        UpsertExecutionContext.CollectionMutationStyle.REMOVE_FROM_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_REMOVE,
        AbstractUpsertOutputOperator.OperationContext.LIST_APPEND
    )), queryExpTTLSetCollRemove.toString());
  }

  private void buildQueryStringForTTLNotSetCollectionsAppendWithListPrepend(StringBuilder updateQueryRoot,
      Map<Long,String> queryStrings)
  {
    // TTL Not set , Collections Append , List prepend
    StringBuilder queryExpTTLNotSetCollAppendListPrepend = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLNotSetCollAppendListPrepend.append(" SET ");
    buildNonPKColumnsExpression(queryExpTTLNotSetCollAppendListPrepend,
        UpsertExecutionContext.ListPlacementStyle.PREPEND_TO_EXISTING_LIST,
        UpsertExecutionContext.CollectionMutationStyle.ADD_TO_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_NOT_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_APPEND,
        AbstractUpsertOutputOperator.OperationContext.LIST_PREPEND
    )), queryExpTTLNotSetCollAppendListPrepend.toString());
  }

  private void buildQueryStringForTTLNotSetCollectionsAppendWithListAppend(StringBuilder updateQueryRoot,
      Map<Long,String> queryStrings)
  {
    // TTL Not set , Collections Append , List append
    StringBuilder queryExpTTLNotSetCollAppendListAppend = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLNotSetCollAppendListAppend.append(" SET ");
    buildNonPKColumnsExpression(queryExpTTLNotSetCollAppendListAppend,
        UpsertExecutionContext.ListPlacementStyle.APPEND_TO_EXISTING_LIST,
        UpsertExecutionContext.CollectionMutationStyle.ADD_TO_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_NOT_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_APPEND,
        AbstractUpsertOutputOperator.OperationContext.LIST_APPEND
    )), queryExpTTLNotSetCollAppendListAppend.toString());
  }


  private void buildQueryStringForTTLNotSetCollectionsRemove(StringBuilder updateQueryRoot,
      Map<Long,String> queryStrings)
  {
    // TTL Not set , Collections Remove
    StringBuilder queryExpTTLNotSetCollRemove = new StringBuilder(updateQueryRoot.toString());
    queryExpTTLNotSetCollRemove.append(" SET ");
    buildNonPKColumnsExpression(queryExpTTLNotSetCollRemove,
        UpsertExecutionContext.ListPlacementStyle.APPEND_TO_EXISTING_LIST, // Just in case user sets it
        UpsertExecutionContext.CollectionMutationStyle.REMOVE_FROM_EXISTING_COLLECTION);
    queryStrings.put(getSlotIndexForMutationContextPreparedStatement(EnumSet.of(
        AbstractUpsertOutputOperator.OperationContext.TTL_NOT_SET,
        AbstractUpsertOutputOperator.OperationContext.COLLECTIONS_REMOVE,
        AbstractUpsertOutputOperator.OperationContext.LIST_APPEND
    )), queryExpTTLNotSetCollRemove.toString());
  }

  private Map<Long, String> generatePreparedStatementsQueryStrings(String keyspaceName,String tableName)
  {
    Map<Long, String> queryStrings = new HashMap<>();
    //UPDATE keyspace_name.table_name USING option AND option SET assignment, assignment, ... WHERE row_specification
    StringBuilder updateQueryRoot = new StringBuilder(" UPDATE " + keyspaceName +
        "." + tableName + " ");
    String ttlSetString = " USING ttl :" + TTL_PARAM_NAME + " SET ";
    buildQueryStringForTTLSetCollectionsAppendAndListPrepend(updateQueryRoot,ttlSetString,queryStrings);
    buildQueryStringForTTLSetCollectionsAppendAndListAppend(updateQueryRoot,ttlSetString,queryStrings);
    buildQueryStringForTTLSetCollectionsRemove(updateQueryRoot,ttlSetString,queryStrings);
    buildQueryStringForTTLNotSetCollectionsAppendWithListPrepend(updateQueryRoot,queryStrings);
    buildQueryStringForTTLNotSetCollectionsAppendWithListAppend(updateQueryRoot,queryStrings);
    buildQueryStringForTTLNotSetCollectionsRemove(updateQueryRoot,queryStrings);
    return queryStrings;
  }


  public static long getSlotIndexForMutationContextPreparedStatement(
      final EnumSet<AbstractUpsertOutputOperator.OperationContext> context)
  {
    Iterator<AbstractUpsertOutputOperator.OperationContext> itrForContexts = context.iterator();
    long indexValue = 0;
    while (itrForContexts.hasNext()) {
      AbstractUpsertOutputOperator.OperationContext aContext = itrForContexts.next();
      indexValue += Math.pow(10, aContext.ordinal());
    }
    return indexValue;
  }

  private void buildNonPKColumnsExpression(final StringBuilder queryExpression,
      UpsertExecutionContext.ListPlacementStyle listPlacementStyle,
      UpsertExecutionContext.CollectionMutationStyle collectionMutationStyle)
  {
    int count = 0;
    for (String colNameEntry : columnDefinitions.keySet()) {
      if (pkColumnNames.contains(colNameEntry)) {
        continue;
      }
      if (count > 0) {
        queryExpression.append(",");
      }
      count += 1;
      if (counterColumns.contains(colNameEntry)) {
        queryExpression.append(" " + colNameEntry + " = " + colNameEntry + " + :" + colNameEntry);
        continue;
      }
      DataType dataType = columnDefinitions.get(colNameEntry);
      if ((!dataType.isCollection()) && (!counterColumns.contains(colNameEntry))) {
        queryExpression.append(" " + colNameEntry + " = :" + colNameEntry);
        continue;
      }
      if ((dataType.isCollection()) && (!dataType.isFrozen())) {
        if (collectionMutationStyle == UpsertExecutionContext.CollectionMutationStyle.REMOVE_FROM_EXISTING_COLLECTION) {
          queryExpression.append(" " + colNameEntry + " = " + colNameEntry + " - :" + colNameEntry);
        }
        if (collectionMutationStyle == UpsertExecutionContext.CollectionMutationStyle.ADD_TO_EXISTING_COLLECTION) {
          if ((setColumns.contains(colNameEntry)) || (mapColumns.contains(colNameEntry))) {
            queryExpression.append(" " + colNameEntry + " = " + colNameEntry + " + :" + colNameEntry);
          }
          if ((listColumns.contains(colNameEntry)) &&
              (listPlacementStyle == UpsertExecutionContext.ListPlacementStyle.APPEND_TO_EXISTING_LIST)) {
            queryExpression.append(" " + colNameEntry + " = " + colNameEntry + " + :" + colNameEntry);
          }
          if ((listColumns.contains(colNameEntry)) &&
              (listPlacementStyle == UpsertExecutionContext.ListPlacementStyle.PREPEND_TO_EXISTING_LIST)) {
            queryExpression.append(" " + colNameEntry + " = :" + colNameEntry + " + " + colNameEntry);
          }
        }
      } else {
        if ((dataType.isCollection()) && (dataType.isFrozen())) {
          queryExpression.append(" " + colNameEntry + " = :" + colNameEntry);
        }
      }
    }
  }

}