/*
 * 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.builder.ParameterBuilder;
import org.apache.shardingsphere.core.rewrite.token.pojo.InsertSetAddGeneratedKeyToken;
import org.apache.shardingsphere.core.rule.ShardingRule;

import java.util.ArrayList;
import java.util.List;

/**
 * Insert set add items token generator.
 *
 * @author panjuan
 */
public final class InsertSetAddGeneratedKeyTokenGenerator implements OptionalSQLTokenGenerator<ShardingRule> {
    
    @Override
    public Optional<InsertSetAddGeneratedKeyToken> generateSQLToken(final OptimizedStatement optimizedStatement, final ParameterBuilder parameterBuilder, final ShardingRule shardingRule) {
        Optional<SetAssignmentsSegment> setAssignmentsSegment = optimizedStatement.getSQLStatement().findSQLSegment(SetAssignmentsSegment.class);
        if (!(optimizedStatement.getSQLStatement() instanceof InsertStatement && setAssignmentsSegment.isPresent())) {
            return Optional.absent();
        }
        return createInsertSetAddGeneratedKeyToken((InsertStatement) optimizedStatement.getSQLStatement(), shardingRule, setAssignmentsSegment.get());
    }
    
    private Optional<InsertSetAddGeneratedKeyToken> createInsertSetAddGeneratedKeyToken(final InsertStatement insertStatement, final ShardingRule shardingRule, final SetAssignmentsSegment segment) {
        Optional<String> generatedKeyColumn = getGeneratedKeyColumn(insertStatement, shardingRule);
        if (generatedKeyColumn.isPresent()) {
            List<AssignmentSegment> assignments = new ArrayList<>(segment.getAssignments());
            return Optional.of(new InsertSetAddGeneratedKeyToken(assignments.get(assignments.size() - 1).getStopIndex() + 1, generatedKeyColumn.get()));
        }
        return Optional.absent();
    }
    
    private Optional<String> getGeneratedKeyColumn(final InsertStatement insertStatement, final ShardingRule shardingRule) {
        String tableName = insertStatement.getTables().getSingleTableName();
        Optional<String> generateKeyColumn = shardingRule.findGenerateKeyColumnName(tableName);
        return generateKeyColumn.isPresent() && !insertStatement.getColumnNames().contains(generateKeyColumn.get()) ? generateKeyColumn : Optional.<String>absent();
    }
}

@@ -970,32 +970,143 @@
 mizedStatement;%0A
+import org.apache.shardingsphere.core.optimize.statement.sharding.dml.insert.ShardingInsertOptimizedStatement;%0A
 import org.apach
@@ -2377,22 +2377,39 @@
 yToken((
-Insert
+ShardingInsertOptimized
 Statemen
@@ -2429,34 +2429,16 @@
 tatement
-.getSQLStatement()
 , shardi
@@ -2574,28 +2574,58 @@
 eyToken(
-final Insert
+%0A            final ShardingInsertOptimized
 Statemen
@@ -2618,38 +2618,41 @@
 imizedStatement 
-insert
+optimized
 Statement, final
@@ -2786,17 +2786,18 @@
 yColumn(
-i
+(I
 nsertSta
@@ -2794,32 +2794,70 @@
 (InsertStatement
+) optimizedStatement.getSQLStatement()
 , shardingRule);
@@ -2919,45 +2919,108 @@
     
-List%3CAssignmentSegment%3E assignments =
+return Optional.of(createInsertSetAddGeneratedKeyToken(optimizedStatement, generatedKeyColumn.get(),
  new
@@ -3057,22 +3057,30 @@
 ments())
+))
 ;%0A    
+    %7D%0A
         
@@ -3095,19 +3095,280 @@
 ptional.
-of(
+absent();%0A    %7D%0A    %0A    private InsertSetAddGeneratedKeyToken createInsertSetAddGeneratedKeyToken(%0A            final ShardingInsertOptimizedStatement optimizedStatement, final String generatedKeyColumn, final List%3CAssignmentSegment%3E assignments) %7B%0A        return 
 new Inse
@@ -3393,16 +3393,33 @@
 eyToken(
+%0A                
 assignme
@@ -3492,59 +3492,89 @@
 lumn
-.get()));%0A        %7D%0A        return Optional.absent(
+, optimizedStatement.getUnits().get(0).getColumnSQLExpression(generatedKeyColumn)
 );%0A 
@@ -3712,85 +3712,8 @@
 ) %7B%0A
-        String tableName = insertStatement.getTables().getSingleTableName();%0A
     
@@ -3796,17 +3796,54 @@
 ame(
-t
+insertStatement.getTables().getSingleT
 ableName
 );%0A 
@@ -3838,16 +3838,18 @@
 ableName
+()
 );%0A     

add createInsertSetAddGeneratedKeyToken()