/*
 * 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 org.apache.shardingsphere.core.optimize.api.statement.InsertOptimizedStatement;
import org.apache.shardingsphere.core.optimize.api.statement.OptimizedStatement;
import org.apache.shardingsphere.core.parse.sql.segment.dml.column.InsertColumnsSegment;
import org.apache.shardingsphere.core.rewrite.builder.ParameterBuilder;
import org.apache.shardingsphere.core.rewrite.token.pojo.InsertAssistedColumnsToken;
import org.apache.shardingsphere.core.rule.EncryptRule;

import java.util.Collection;
import java.util.LinkedList;

/**
 * Insert assisted columns token generator.
 *
 * @author panjuan
 */
public final class InsertAssistedColumnsTokenGenerator implements OptionalSQLTokenGenerator<EncryptRule> {
    
    @Override
    public Optional<InsertAssistedColumnsToken> generateSQLToken(
            final OptimizedStatement optimizedStatement, final ParameterBuilder parameterBuilder, final EncryptRule encryptRule, final boolean isQueryWithCipherColumn) {
        Optional<InsertColumnsSegment> insertColumnsSegment = optimizedStatement.getSQLStatement().findSQLSegment(InsertColumnsSegment.class);
        if (!(optimizedStatement instanceof InsertOptimizedStatement && insertColumnsSegment.isPresent())) {
            return Optional.absent();
        }
        return createInsertAssistedColumnsToken((InsertOptimizedStatement) optimizedStatement, insertColumnsSegment.get(), encryptRule);
    }
    
    private Optional<InsertAssistedColumnsToken> createInsertAssistedColumnsToken(
            final InsertOptimizedStatement optimizedStatement, final InsertColumnsSegment segment, final EncryptRule encryptRule) {
        return encryptRule.getEncryptEngine().getAssistedQueryColumns(optimizedStatement.getTables().getSingleTableName()).isEmpty() ? Optional.<InsertAssistedColumnsToken>absent()
                : Optional.of(new InsertAssistedColumnsToken(segment.getStopIndex(), getQueryAssistedColumns(optimizedStatement, encryptRule), segment.getColumns().isEmpty()));
    }
    
    private Collection<String> getQueryAssistedColumns(final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule) {
        Collection<String> result = new LinkedList<>();
        for (String each : optimizedStatement.getInsertColumns().getRegularColumnNames()) {
            Optional<String> assistedColumnName = encryptRule.getEncryptEngine().getAssistedQueryColumn(optimizedStatement.getTables().getSingleTableName(), each);
            if (assistedColumnName.isPresent()) {
                result.add(assistedColumnName.get());
            }
        }
        return result;
    }
}

@@ -1613,24 +1613,229 @@
 ule%3E %7B%0A    %0A
+    private EncryptRule encryptRule;%0A    %0A    private InsertOptimizedStatement insertOptimizedStatement;%0A    %0A    private InsertColumnsSegment insertColumnsSegment;%0A    %0A    private String tableName;%0A    %0A
     @Overrid
@@ -2068,24 +2068,349 @@
 erColumn) %7B%0A
+        if (isNotNeedToGenerateSQLToken(optimizedStatement)) %7B%0A            return Optional.absent();%0A        %7D%0A        initParameters(encryptRule, optimizedStatement);%0A        return createInsertAssistedColumnsToken();%0A    %7D%0A    %0A    private boolean isNotNeedToGenerateSQLToken(final OptimizedStatement optimizedStatement) %7B%0A
         Opti
@@ -2548,20 +2548,23 @@
         
-if (
+return 
 !(optimi
@@ -2652,108 +2652,206 @@
 t())
-) %7B
+;
 %0A    
+%7D%0A
     
+%0A
     
-return Optional.absent();%0A        %7D%0A        return createInsertAssistedColumnsToken(
+private void initParameters(final EncryptRule encryptRule, final OptimizedStatement optimizedStatement) %7B%0A        this.encryptRule = encryptRule;%0A        insertOptimizedStatement = 
 (Ins
@@ -2891,17 +2891,25 @@
 tatement
-,
+;%0A       
  insertC
@@ -2925,252 +2925,270 @@
 ment
-.get(), encryptRule);%0A    %7D%0A    %0A    private Optional%3CInsertAssistedColumnsToken%3E createInsertAssistedColumnsToken(%0A            final InsertOptimizedStatement optimizedStatement, final InsertColumnsSegment segment, final EncryptRule encryptRule
+ = optimizedStatement.getSQLStatement().findSQLSegment(InsertColumnsSegment.class).get();%0A        tableName = insertOptimizedStatement.getTables().getSingleTableName();%0A    %7D%0A    %0A    private Optional%3CInsertAssistedColumnsToken%3E createInsertAssistedColumnsToken(
 ) %7B%0A
@@ -3261,57 +3261,17 @@
 mns(
-optimizedStatement.getTables().getSingleT
+t
 ableName
 ()).
@@ -3266,18 +3266,16 @@
 ableName
-()
 ).isEmpt
@@ -3387,17 +3387,30 @@
 nsToken(
-s
+insertColumnsS
 egment.g
@@ -3452,43 +3452,25 @@
 mns(
-optimizedStatement, encryptRule), s
+), insertColumnsS
 egme
@@ -3568,88 +3568,8 @@
 mns(
-final InsertOptimizedStatement optimizedStatement, final EncryptRule encryptRule
 ) %7B%0A
@@ -3651,17 +3651,23 @@
  each : 
-o
+insertO
 ptimized
@@ -3830,57 +3830,17 @@
 umn(
-optimizedStatement.getTables().getSingleT
+t
 ableName
 (), 
@@ -3835,18 +3835,16 @@
 ableName
-()
 , each);

reconstruct it