/*
 * Copyright 1999-2015 dangdang.com.
 * <p>
 * Licensed 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.
 * </p>
 */

package com.dangdang.ddframe.rdb.sharding.parsing.parser.context;

import com.dangdang.ddframe.rdb.sharding.api.rule.ShardingRule;
import com.dangdang.ddframe.rdb.sharding.api.rule.TableRule;
import com.dangdang.ddframe.rdb.sharding.constant.SQLType;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.expr.SQLExpr;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.expr.SQLNumberExpr;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.expr.SQLPlaceholderExpr;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.token.AutoGeneratedKeysToken;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.token.ItemsToken;
import com.dangdang.ddframe.rdb.sharding.parsing.parser.token.SQLToken;
import com.google.common.base.Optional;
import lombok.Getter;
import lombok.Setter;

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

/**
 * Insert SQL上下文.
 *
 * @author zhangliang
 */
@Getter
@Setter
public final class InsertSQLContext extends AbstractSQLContext {
    
    private final GeneratedKeyContext generatedKeyContext = new GeneratedKeyContext();
    
    private final Collection<ShardingColumnContext> shardingColumnContexts = new LinkedList<>();
    
    private int columnsListLastPosition;
    
    private int valuesListLastPosition;
    
    public InsertSQLContext() {
        super(SQLType.INSERT);
    }
    
    /**
     * 判断是否包含该列.
     *
     * @param columnName 列名称
     * @return 是否包含该列
     */
    public boolean hasColumn(final String columnName) {
        for (ShardingColumnContext shardingColumnContext : shardingColumnContexts) {
            if (shardingColumnContext.getColumnName().equalsIgnoreCase(columnName)) {
                return true;
            }
        }
        return false;
    }
    
    /**
     * 获取自增主键.
     *
     * @param shardingRule 分片规则
     * @return 自增列与主键映射表
     */
    public Map<String, Number> generateKeys(final ShardingRule shardingRule) {
        Optional<TableRule> tableRuleOptional = shardingRule.tryFindTableRule(getTables().iterator().next().getName());
        if (!tableRuleOptional.isPresent()) {
            return Collections.emptyMap();
        }
        TableRule tableRule = tableRuleOptional.get();
        Map<String, Number> result = new LinkedHashMap<>(generatedKeyContext.getColumns().size());
        for (String each : generatedKeyContext.getColumns()) {
            Number generatedKey = tableRule.generateId(each);
            result.put(each, generatedKey);
            generatedKeyContext.putValue(each, generatedKey);
        }
        return result;
    }
    
    /**
     * 追加自增主键标记对象.
     *
     * @param shardingRule 分片规则
     */
    public void appendGenerateKeysToken(final ShardingRule shardingRule) {
        Optional<AutoGeneratedKeysToken> autoGeneratedKeysToken =  findAutoGeneratedKeysToken();
        if (!autoGeneratedKeysToken.isPresent()) {
            return;
        }
        Map<String, Number> generatedKeys = generateKeys(shardingRule);
        String tableName = getTables().get(0).getName();
        ItemsToken valuesToken = new ItemsToken(autoGeneratedKeysToken.get().getBeginPosition());
        for (Map.Entry<String, Number> entry : generatedKeys.entrySet()) {
            valuesToken.getItems().add(entry.getValue().toString());
            addCondition(shardingRule, new ShardingColumnContext(entry.getKey(), tableName, true), new SQLNumberExpr(entry.getValue()));
        }
        getSqlTokens().remove(autoGeneratedKeysToken.get());
        getSqlTokens().add(valuesToken);
    }
    
    /**
     * 追加自增主键标记对象.
     *
     * @param shardingRule 分片规则
     * @param parametersSize 参数个数
     */
    public void appendGenerateKeysToken(final ShardingRule shardingRule, final int parametersSize) {
        Optional<AutoGeneratedKeysToken> autoGeneratedKeysToken =  findAutoGeneratedKeysToken();
        if (!autoGeneratedKeysToken.isPresent()) {
            return;
        }
        String tableName = getTables().get(0).getName();
        List<String> autoGenerateKeyColumns = generatedKeyContext.getColumns();
        ItemsToken valuesToken = new ItemsToken(autoGeneratedKeysToken.get().getBeginPosition());
        int offset = 0;
        for (String each : autoGenerateKeyColumns) {
            valuesToken.getItems().add("?");
            addCondition(shardingRule, new ShardingColumnContext(each, tableName, true), new SQLPlaceholderExpr(parametersSize + offset));
            offset++;
        }
        getSqlTokens().remove(autoGeneratedKeysToken.get());
        getSqlTokens().add(valuesToken);
    }
    
    private Optional<AutoGeneratedKeysToken> findAutoGeneratedKeysToken() {
        for (SQLToken each : getSqlTokens()) {
            if (each instanceof AutoGeneratedKeysToken) {
                return Optional.of((AutoGeneratedKeysToken) each);
            }
        }
        return Optional.absent();
    }
    
    private void addCondition(final ShardingRule shardingRule, final ShardingColumnContext shardingColumnContext, final SQLExpr sqlExpr) {
        if (shardingRule.isShardingColumn(shardingColumnContext)) {
            getConditionContext().add(new ConditionContext.Condition(shardingColumnContext, sqlExpr));
        }
    }
}

@@ -1528,31 +1528,8 @@
 st;%0A
-import java.util.List;%0A
 impo
@@ -4675,88 +4675,8 @@
 ();%0A
-        List%3CString%3E autoGenerateKeyColumns = generatedKeyContext.getColumns();%0A
     
@@ -4824,30 +4824,40 @@
 h : 
-autoG
+g
 enerate
+d
 KeyCo
+ntext.getCo
 lumns
+()
 ) %7B%0A

remove unnecessary codes