/*
 * 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.shardingsphere.core.rewrite.token.generator;

import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import org.apache.shardingsphere.core.optimize.api.statement.ConditionOptimizedStatement;
import org.apache.shardingsphere.core.optimize.api.statement.OptimizedStatement;
import org.apache.shardingsphere.core.optimize.encrypt.segment.condition.EncryptCondition;
import org.apache.shardingsphere.core.rewrite.builder.ParameterBuilder;
import org.apache.shardingsphere.core.rewrite.token.pojo.EncryptColumnToken;
import org.apache.shardingsphere.core.rewrite.token.pojo.WhereEncryptColumnToken;
import org.apache.shardingsphere.core.rule.ColumnNode;
import org.apache.shardingsphere.core.rule.EncryptRule;

import java.util.Collection;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

/**
 * Where encrypt column token generator.
 *
 * @author panjuan
 */
public final class WhereEncryptColumnTokenGenerator implements CollectionSQLTokenGenerator<EncryptRule> {
    
    private boolean isQueryWithCipherColumn;
    
    @Override
    public Collection<EncryptColumnToken> generateSQLTokens(
            final OptimizedStatement optimizedStatement, final ParameterBuilder parameterBuilder, final EncryptRule encryptRule, final boolean isQueryWithCipherColumn) {
        this.isQueryWithCipherColumn = isQueryWithCipherColumn;
        return optimizedStatement instanceof ConditionOptimizedStatement
                ? createEncryptColumnToken((ConditionOptimizedStatement) optimizedStatement, parameterBuilder, encryptRule) : Collections.<EncryptColumnToken>emptyList();
    }
    
    private Collection<EncryptColumnToken> createEncryptColumnToken(final ConditionOptimizedStatement optimizedStatement, final ParameterBuilder parameterBuilder, final EncryptRule encryptRule) {
        Collection<EncryptColumnToken> result = new LinkedList<>();
        for (EncryptCondition each : optimizedStatement.getEncryptConditions().getConditions()) {
            result.add(createWhereEncryptColumnToken(each, parameterBuilder, encryptRule));
        }
        return result;
    }
    
    private WhereEncryptColumnToken createWhereEncryptColumnToken(final EncryptCondition encryptCondition, final ParameterBuilder parameterBuilder, final EncryptRule encryptRule) {
        List<Object> originalValues = encryptCondition.getValues(parameterBuilder.getOriginalParameters());
        if (!isQueryWithCipherColumn) {
            Optional<String> plainColumn = encryptRule.getEncryptEngine().getPlainColumn(encryptCondition.getTableName(), encryptCondition.getColumnName());
            Preconditions.checkArgument(!plainColumn.isPresent(), "Plain column is required.");
            return new WhereEncryptColumnToken(encryptCondition.getStartIndex(), encryptCondition.getStopIndex(), 
                    plainColumn.get(), getPositionValues(encryptCondition.getPositionValueMap().keySet(), originalValues), encryptCondition.getPositionIndexMap().keySet(), encryptCondition.getOperator());
        }
        Optional<String> assistedQueryColumn = encryptRule.getEncryptEngine().getAssistedQueryColumn(encryptCondition.getTableName(), encryptCondition.getColumnName());
        ColumnNode columnNode = new ColumnNode(encryptCondition.getTableName(), encryptCondition.getColumnName());
        List<Object> encryptedValues = assistedQueryColumn.isPresent() ? encryptRule.getEncryptEngine().getEncryptAssistedColumnValues(columnNode, originalValues) : encryptRule.getEncryptEngine().getEncryptColumnValues(columnNode, originalValues);
        String encryptedName = assistedQueryColumn.isPresent() ? assistedQueryColumn.get() : encryptRule.getEncryptEngine().getCipherColumn(encryptCondition.getTableName(), encryptCondition.getColumnName());
        encryptParameters(encryptCondition.getPositionIndexMap(), encryptedValues, parameterBuilder);
        return new WhereEncryptColumnToken(encryptCondition.getStartIndex(), encryptCondition.getStopIndex(), encryptedName,
                getPositionValues(encryptCondition.getPositionValueMap().keySet(), encryptedValues), encryptCondition.getPositionIndexMap().keySet(), encryptCondition.getOperator());
    }
    
    private void encryptParameters(final Map<Integer, Integer> positionIndexes, final List<Object> encryptColumnValues, final ParameterBuilder parameterBuilder) {
        if (!positionIndexes.isEmpty()) {
            for (Entry<Integer, Integer> entry : positionIndexes.entrySet()) {
                parameterBuilder.getOriginalParameters().set(entry.getValue(), encryptColumnValues.get(entry.getKey()));
            }
        }
    }
    
    private Map<Integer, Object> getPositionValues(final Collection<Integer> valuePositions, final List<Object> encryptColumnValues) {
        Map<Integer, Object> result = new LinkedHashMap<>();
        for (int each : valuePositions) {
            result.put(each, encryptColumnValues.get(each));
        }
        return result;
    }
}

@@ -912,44 +912,66 @@
 ort 
-com.google.common.base.Precondi
+org.apache.shardingsphere.core.exception.ShardingExcep
 tion
-s
 ;%0Aim
@@ -1465,63 +1465,8 @@
 en;%0A
-import org.apache.shardingsphere.core.rule.ColumnNode;%0A
 impo
@@ -3351,84 +3351,42 @@
     
-Optional%3CString%3E plainColumn = encryptRule.getEncryptEngine().getPlain
+return new WhereEncrypt
 Column
+Token
 (enc
@@ -3398,33 +3398,34 @@
 ondition.get
-TableName
+StartIndex
 (), encryptC
@@ -3440,230 +3440,65 @@
 .get
-ColumnName());%0A            Preconditions.checkArgument(!plainColumn.isPresent(), %22Plain column is required.%22);%0A            return new WhereEncryptColumnToken(encryptCondition.getStartIndex(), encryptCondition.getStopIndex(
+StopIndex(), getPlainColumn(encryptCondition, encryptRule
 ), %0A
@@ -3520,27 +3520,8 @@
     
- plainColumn.get(),
  get
@@ -3866,123 +3866,8 @@
 ));%0A
-        ColumnNode columnNode = new ColumnNode(encryptCondition.getTableName(), encryptCondition.getColumnName());%0A
     
@@ -3993,34 +3993,106 @@
 olumnValues(
-c
+encryptCondition.getTableName(), %0A                encryptCondition.getC
 olumnN
-ode
+ame()
 , originalVa
@@ -4157,18 +4157,73 @@
 ues(
-columnNode
+encryptCondition.getTableName(), encryptCondition.getColumnName()
 , or
@@ -4300,16 +4300,33 @@
 esent() 
+%0A                
 ? assist
@@ -4875,32 +4875,442 @@
 ());%0A    %7D%0A    %0A
+    private String getPlainColumn(final EncryptCondition encryptCondition, final EncryptRule encryptRule) %7B%0A        Optional%3CString%3E result = encryptRule.getEncryptEngine().getPlainColumn(encryptCondition.getTableName(), encryptCondition.getColumnName());%0A        if (result.isPresent()) %7B%0A            return result.get();%0A        %7D%0A        throw new ShardingException(%22Plain column is required.%22);%0A    %7D%0A    %0A
     private void

modify createWhereEncryptColumnToken()