/*
    *    Copyright 2009-2021 the original author or authors.
    *
    *    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.
   */
  package org.apache.ibatis.executor.keygen;
  
  import java.sql.ResultSet;
  import java.sql.ResultSetMetaData;
  import java.sql.SQLException;
  import java.sql.Statement;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Set;
  
  import org.apache.ibatis.binding.MapperMethod.ParamMap;
  import org.apache.ibatis.executor.Executor;
  import org.apache.ibatis.executor.ExecutorException;
  import org.apache.ibatis.mapping.MappedStatement;
  import org.apache.ibatis.reflection.ArrayUtil;
  import org.apache.ibatis.reflection.MetaObject;
  import org.apache.ibatis.reflection.ParamNameResolver;
  import org.apache.ibatis.session.Configuration;
  import org.apache.ibatis.session.defaults.DefaultSqlSession.StrictMap;
  import org.apache.ibatis.type.JdbcType;
  import org.apache.ibatis.type.TypeHandler;
  import org.apache.ibatis.type.TypeHandlerRegistry;
  import org.apache.ibatis.util.MapUtil;
  
  /**
   * @author Clinton Begin
   * @author Kazuki Shimizu
   */
  public class Jdbc3KeyGenerator implements KeyGenerator {
  
    private static final String SECOND_GENERIC_PARAM_NAME = ParamNameResolver.GENERIC_NAME_PREFIX + "2";
  
    /**
     * A shared instance.
     *
     * @since 3.4.3
     */
    public static final Jdbc3KeyGenerator INSTANCE = new Jdbc3KeyGenerator();
  
    private static final String MSG_TOO_MANY_KEYS = "Too many keys are generated. There are only %d target objects. "
        + "You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.";
  
    @Override
    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
      // do nothing
    }
  
    @Override
    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
      processBatch(ms, stmt, parameter);
    }
  
    public void processBatch(MappedStatement ms, Statement stmt, Object parameter) {
      final String[] keyProperties = ms.getKeyProperties();
      if (keyProperties == null || keyProperties.length == 0) {
        return;
      }
      try (ResultSet rs = stmt.getGeneratedKeys()) {
        final ResultSetMetaData rsmd = rs.getMetaData();
        final Configuration configuration = ms.getConfiguration();
        if (rsmd.getColumnCount() < keyProperties.length) {
          // Error?
        } else {
          assignKeys(configuration, rs, rsmd, keyProperties, parameter);
        }
      } catch (Exception e) {
        throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e, e);
      }
    }
  
    @SuppressWarnings("unchecked")
    private void assignKeys(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties,
        Object parameter) throws SQLException {
      if (parameter instanceof ParamMap || parameter instanceof StrictMap) {
        // Multi-param or single param with @Param
        assignKeysToParamMap(configuration, rs, rsmd, keyProperties, (Map<String, ?>) parameter);
      } else if (parameter instanceof ArrayList && !((ArrayList<?>) parameter).isEmpty()
          && ((ArrayList<?>) parameter).get(0) instanceof ParamMap) {
       // Multi-param or single param with @Param in batch operation
       assignKeysToParamMapList(configuration, rs, rsmd, keyProperties, (ArrayList<ParamMap<?>>) parameter);
     } else {
       // Single param without @Param
       assignKeysToParam(configuration, rs, rsmd, keyProperties, parameter);
     }
   }
 
   private void assignKeysToParam(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
       String[] keyProperties, Object parameter) throws SQLException {
     Collection<?> params = collectionize(parameter);
     if (params.isEmpty()) {
       return;
     }
     List<KeyAssigner> assignerList = new ArrayList<>();
     for (int i = 0; i < keyProperties.length; i++) {
       assignerList.add(new KeyAssigner(configuration, rsmd, i + 1, null, keyProperties[i]));
     }
     Iterator<?> iterator = params.iterator();
     while (rs.next()) {
       if (!iterator.hasNext()) {
         throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, params.size()));
       }
       Object param = iterator.next();
       assignerList.forEach(x -> x.assign(rs, param));
     }
   }
 
   private void assignKeysToParamMapList(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
       String[] keyProperties, ArrayList<ParamMap<?>> paramMapList) throws SQLException {
     Iterator<ParamMap<?>> iterator = paramMapList.iterator();
     List<KeyAssigner> assignerList = new ArrayList<>();
     long counter = 0;
     while (rs.next()) {
       if (!iterator.hasNext()) {
         throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter));
       }
       ParamMap<?> paramMap = iterator.next();
       if (assignerList.isEmpty()) {
         for (int i = 0; i < keyProperties.length; i++) {
           assignerList
               .add(getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, false)
                   .getValue());
         }
       }
       assignerList.forEach(x -> x.assign(rs, paramMap));
       counter++;
     }
   }
 
   private void assignKeysToParamMap(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
       String[] keyProperties, Map<String, ?> paramMap) throws SQLException {
     if (paramMap.isEmpty()) {
       return;
     }
     Map<String, Entry<Iterator<?>, List<KeyAssigner>>> assignerMap = new HashMap<>();
     for (int i = 0; i < keyProperties.length; i++) {
       Entry<String, KeyAssigner> entry = getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i],
           keyProperties, true);
       Entry<Iterator<?>, List<KeyAssigner>> iteratorPair = MapUtil.computeIfAbsent(assignerMap, entry.getKey(),
           k -> MapUtil.entry(collectionize(paramMap.get(k)).iterator(), new ArrayList<>()));
       iteratorPair.getValue().add(entry.getValue());
     }
     long counter = 0;
     while (rs.next()) {
       for (Entry<Iterator<?>, List<KeyAssigner>> pair : assignerMap.values()) {
         if (!pair.getKey().hasNext()) {
           throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter));
         }
         Object param = pair.getKey().next();
         pair.getValue().forEach(x -> x.assign(rs, param));
       }
       counter++;
     }
   }
 
   private Entry<String, KeyAssigner> getAssignerForParamMap(Configuration config, ResultSetMetaData rsmd,
       int columnPosition, Map<String, ?> paramMap, String keyProperty, String[] keyProperties, boolean omitParamName) {
     Set<String> keySet = paramMap.keySet();
     // A caveat : if the only parameter has {@code @Param("param2")} on it,
     // it must be referenced with param name e.g. 'param2.x'.
     boolean singleParam = !keySet.contains(SECOND_GENERIC_PARAM_NAME);
     int firstDot = keyProperty.indexOf('.');
     if (firstDot == -1) {
       if (singleParam) {
         return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
       }
       throw new ExecutorException("Could not determine which parameter to assign generated keys to. "
           + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). "
           + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are "
           + keySet);
     }
     String paramName = keyProperty.substring(0, firstDot);
     if (keySet.contains(paramName)) {
       String argParamName = omitParamName ? null : paramName;
       String argKeyProperty = keyProperty.substring(firstDot + 1);
       return MapUtil.entry(paramName, new KeyAssigner(config, rsmd, columnPosition, argParamName, argKeyProperty));
     } else if (singleParam) {
       return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
     } else {
       throw new ExecutorException("Could not find parameter '" + paramName + "'. "
           + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). "
           + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are "
           + keySet);
     }
   }
 
   private Entry<String, KeyAssigner> getAssignerForSingleParam(Configuration config, ResultSetMetaData rsmd,
       int columnPosition, Map<String, ?> paramMap, String keyProperty, boolean omitParamName) {
     // Assume 'keyProperty' to be a property of the single param.
     String singleParamName = nameOfSingleParam(paramMap);
     String argParamName = omitParamName ? null : singleParamName;
     return MapUtil.entry(singleParamName, new KeyAssigner(config, rsmd, columnPosition, argParamName, keyProperty));
   }
 
   private static String nameOfSingleParam(Map<String, ?> paramMap) {
     // There is virtually one parameter, so any key works.
     return paramMap.keySet().iterator().next();
   }
 
   private static Collection<?> collectionize(Object param) {
     if (param instanceof Collection) {
       return (Collection<?>) param;
     } else if (param instanceof Object[]) {
       return Arrays.asList((Object[]) param);
     } else {
       return Arrays.asList(param);
     }
   }
 
   private class KeyAssigner {
     private final Configuration configuration;
     private final ResultSetMetaData rsmd;
     private final TypeHandlerRegistry typeHandlerRegistry;
     private final int columnPosition;
     private final String paramName;
     private final String propertyName;
     private TypeHandler<?> typeHandler;
 
     protected KeyAssigner(Configuration configuration, ResultSetMetaData rsmd, int columnPosition, String paramName,
         String propertyName) {
       super();
       this.configuration = configuration;
       this.rsmd = rsmd;
       this.typeHandlerRegistry = configuration.getTypeHandlerRegistry();
       this.columnPosition = columnPosition;
       this.paramName = paramName;
       this.propertyName = propertyName;
     }
 
     protected void assign(ResultSet rs, Object param) {
       if (paramName != null) {
         // If paramName is set, param is ParamMap
         param = ((ParamMap<?>) param).get(paramName);
       }
       MetaObject metaParam = configuration.newMetaObject(param);
       try {
         if (typeHandler == null) {
           if (metaParam.hasSetter(propertyName)) {
             Class<?> propertyType = metaParam.getSetterType(propertyName);
             typeHandler = typeHandlerRegistry.getTypeHandler(propertyType,
                 JdbcType.forCode(rsmd.getColumnType(columnPosition)));
           } else {
             throw new ExecutorException("No setter found for the keyProperty '" + propertyName + "' in '"
                 + metaParam.getOriginalObject().getClass().getName() + "'.");
           }
         }
         if (typeHandler == null) {
           // Error?
         } else {
           Object value = typeHandler.getResult(rs, columnPosition);
           metaParam.setValue(propertyName, value);
         }
       } catch (SQLException e) {
         throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e,
             e);
       }
     }
   }
 }