package me.chyxion.tigon.mybatis;

import java.util.*;
import java.sql.ResultSet;
import java.sql.Statement;
import java.sql.SQLException;
import lombok.extern.slf4j.Slf4j;
import java.sql.ResultSetMetaData;
import me.chyxion.tigon.model.M0;
import org.apache.ibatis.type.JdbcType;
import org.apache.ibatis.type.TypeHandler;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.type.TypeHandlerRegistry;
import org.apache.ibatis.executor.ExecutorException;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.executor.keygen.KeyGenerator;

/**
 * @author Shaun Chyxion <br>
 * chyxion@163.com <br>
 * 2017/1/9 14:46
 */
@Slf4j
class Jdbc3KeyGen implements KeyGenerator {
    private KeyGenerator originKeyGen;

    public Jdbc3KeyGen(KeyGenerator originKeyGen) {
        this.originKeyGen = originKeyGen;
    }

    public void processBefore(Executor executor,
        MappedStatement ms,
        Statement stmt,
        Object parameter) {
        // noop
    }

    /**
     * {@inheritDoc}
     */
    public void processAfter(Executor executor,
        MappedStatement ms,
        Statement stmt,
        Object params) {
        ResultSet rs = null;
        try {
            rs = stmt.getGeneratedKeys();
            final Configuration configuration = ms.getConfiguration();
            final TypeHandlerRegistry typeHandlerRegistry =
                configuration.getTypeHandlerRegistry();
            String[] keyProps = ms.getKeyProperties();
            if (keyProps == null || keyProps.length == 0) {
                keyProps = new String[] {M0.ID};
            }

            final ResultSetMetaData rsmd = rs.getMetaData();
            if (rsmd.getColumnCount() >= keyProps.length) {
                TypeHandler<?>[] typeHandlers = null;
                for (Object model : getModels(params)) {
                    log.info("Process Model [{}] Key.", model);
                    // there should be one row for each statement (also one for each params)
                    if (rs.next()) {
                        final MetaObject metaModel = configuration.newMetaObject(model);
                        if (typeHandlers == null) {
                            typeHandlers = getTypeHandlers(typeHandlerRegistry, metaModel, keyProps, rsmd);
                        }
                        populateKeys(rs, metaModel, keyProps, typeHandlers);
                        log.info("Populate Model [{}] Key Result.", model);
                    }
                    else {
                        break;
                    }
                }
            }
        }
        catch (Exception e) {
            throw new ExecutorException(
                "Error getting generated key or setting result to params object. Cause: " + e, e);
        }
        finally {
            // set original Key Generator back
            SystemMetaObject.forObject(ms).setValue("keyGenerator", originKeyGen);
            if (rs != null) {
                try {
                    rs.close();
                }
                catch (Exception e) {
                    // ignore
                }
            }
        }
    }

    @SuppressWarnings("unchecked")
    private Collection<Object> getModels(Object objParams) {
        Collection<Object> models = null;
        if (objParams instanceof MapperMethod.ParamMap) {
            Map<String, Object> mapParams = (Map<String, Object>) objParams;
            if (mapParams.containsKey(BaseMapper.PARAM_MODEL_KEY)) {
                models = new ArrayList<Object>(1);
                models.add(mapParams.get(BaseMapper.PARAM_MODEL_KEY));
            }
            else if (mapParams.containsKey(BaseMapper.PARAM_MODELS_KEY)) {
                Object objModels = mapParams.get(BaseMapper.PARAM_MODELS_KEY);
                if (objModels.getClass().isArray()) {
                    models = Arrays.asList((Object[]) objModels);
                }
                else if (objModels instanceof Collection) {
                    models = (Collection<Object>) objModels;
                }
            }
        }
        return models != null ? models : Collections.emptyList();
    }

    private TypeHandler<?>[] getTypeHandlers(TypeHandlerRegistry typeHandlerRegistry,
        MetaObject metaParam,
        String[] keyProperties,
        ResultSetMetaData rsmd) throws SQLException {
        TypeHandler<?>[] typeHandlers = new TypeHandler<?>[keyProperties.length];
        for (int i = 0; i < keyProperties.length; i++) {
            if (metaParam.hasSetter(keyProperties[i])) {
                Class<?> keyPropertyType = metaParam.getSetterType(keyProperties[i]);
                typeHandlers[i] = typeHandlerRegistry.getTypeHandler(keyPropertyType,
                    JdbcType.forCode(rsmd.getColumnType(i + 1)));
            }
        }
        return typeHandlers;
    }

    private void populateKeys(ResultSet rs, MetaObject metaParam, String[] keyProperties, TypeHandler<?>[] typeHandlers) throws SQLException {
        for (int i = 0; i < keyProperties.length; ++i) {
            TypeHandler<?> th = typeHandlers[i];
            if (th != null) {
                metaParam.setValue(keyProperties[i], th.getResult(rs, i + 1));
            }
        }
    }
}
