/*
 * 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.statement.OptimizedStatement;
import org.apache.shardingsphere.core.parse.sql.segment.dml.assignment.AssignmentSegment;
import org.apache.shardingsphere.core.parse.sql.segment.dml.assignment.SetAssignmentsSegment;
import org.apache.shardingsphere.core.parse.sql.statement.dml.InsertStatement;
import org.apache.shardingsphere.core.rewrite.token.pojo.InsertSetAddAssistedColumnsToken;
import org.apache.shardingsphere.core.rule.EncryptRule;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;

/**
 * Insert set add items token generator.
 *
 * @author panjuan
 */
public final class InsertSetAddAssistedColumnsTokenGenerator implements OptionalSQLTokenGenerator<EncryptRule> {
    
    @Override
    public Optional<InsertSetAddAssistedColumnsToken> generateSQLToken(final OptimizedStatement optimizedStatement, final List<Object> parameters, final EncryptRule encryptRule) {
        Optional<SetAssignmentsSegment> setAssignmentsSegment = optimizedStatement.getSQLStatement().findSQLSegment(SetAssignmentsSegment.class);
        if (!(optimizedStatement.getSQLStatement() instanceof InsertStatement && setAssignmentsSegment.isPresent())) {
            return Optional.absent();
        }
        return createInsertSetAddItemsToken((InsertStatement) optimizedStatement.getSQLStatement(), encryptRule, setAssignmentsSegment.get());
    }
    
    private Optional<InsertSetAddAssistedColumnsToken> createInsertSetAddItemsToken(final InsertStatement insertStatement, final EncryptRule encryptRule, final SetAssignmentsSegment segment) {
        Collection<String> columnNames = getQueryAssistedColumn(insertStatement, encryptRule);
        if (columnNames.isEmpty()) {
            return Optional.absent();
        }
        List<AssignmentSegment> assignments = new ArrayList<>(segment.getAssignments());
        return Optional.of(new InsertSetAddAssistedColumnsToken(assignments.get(assignments.size() - 1).getStopIndex() + 1, columnNames));
    }
    
    private Collection<String> getQueryAssistedColumn(final InsertStatement insertStatement, final EncryptRule encryptRule) {
        Collection<String> result = new LinkedList<>();
        for (String each : insertStatement.getColumnNames()) {
            Optional<String> assistedColumnName = encryptRule.getEncryptorEngine().getAssistedQueryColumn(insertStatement.getTables().getSingleTableName(), each);
            if (assistedColumnName.isPresent()) {
                result.add(assistedColumnName.get());
            }
        }
        return result;
    }
}

@@ -970,32 +970,143 @@
 mizedStatement;%0A
+import org.apache.shardingsphere.core.optimize.statement.sharding.dml.insert.ShardingInsertOptimizedStatement;%0A
 import org.apach
@@ -1265,32 +1265,116 @@
 gnmentsSegment;%0A
+import org.apache.shardingsphere.core.parse.sql.segment.dml.expr.ExpressionSegment;%0A
 import org.apach
@@ -2434,16 +2434,33 @@
 msToken(
+%0A                
 (InsertS
@@ -2519,16 +2519,71 @@
 yptRule,
+ (ShardingInsertOptimizedStatement) optimizedStatement,
  setAssi
@@ -2696,24 +2696,37 @@
 dItemsToken(
+%0A            
 final Insert
@@ -2782,16 +2782,75 @@
 yptRule,
+ final ShardingInsertOptimizedStatement optimizedStatement,
  final S
@@ -2940,24 +2940,29 @@
 sistedColumn
+Names
 (insertState
@@ -3219,16 +3219,33 @@
 nsToken(
+%0A                
 assignme
@@ -3307,16 +3307,79 @@
 umnNames
+, getQueryAssistedColumnValues(columnNames, optimizedStatement)
 ));%0A    
@@ -3438,16 +3438,21 @@
 edColumn
+Names
 (final I
@@ -3955,10 +3955,424 @@
 ;%0A    %7D%0A
+    %0A    private Collection%3CExpressionSegment%3E getQueryAssistedColumnValues(final Collection%3CString%3E columnNames, final ShardingInsertOptimizedStatement optimizedStatement) %7B%0A        Collection%3CExpressionSegment%3E result = new LinkedList%3C%3E();%0A        for (String each : columnNames) %7B%0A            result.add(optimizedStatement.getUnits().get(0).getColumnSQLExpression(each));%0A        %7D%0A        return result;%0A    %7D%0A
 %7D%0A

add getQueryAssistedColumnValues()