package me.chyxion.tigon.mybatis;

import java.sql.Connection;
import java.util.Properties;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.mapping.*;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.executor.keygen.KeyGenerator;
import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.executor.statement.RoutingStatementHandler;

/**
 * @author Shaun Chyxion <br>
 * chyxion@163.com <br>
 * 2017/1/9 14:46
 */
@Intercepts({
    @Signature(type = StatementHandler.class,
        method = "prepare",
        args = {Connection.class, Integer.class})
})
@Slf4j
class KeyGenInterceptor implements Interceptor {

    /**
     * {@inheritDoc}
     */
    public Object intercept(Invocation invocation) throws Throwable {
        String method = invocation.getMethod().getName();
        Object objRtn = null;
        if ("prepare".equals(method)) {
            objRtn = injectKeyGen(invocation);
        }
        else {
            objRtn = invocation.proceed();
        }
        return objRtn;
    }

    /**
     * {@inheritDoc}
     */
    public Object plugin(Object target) {
        return target instanceof StatementHandler ?
            Plugin.wrap(target, this) : target;
    }

    /**
     * {@inheritDoc}
     */
    public void setProperties(Properties properties) {
    }

    // --
    // private methods

    StatementHandler getStatementHandler(Invocation invocation) {
        StatementHandler statement = (StatementHandler) invocation.getTarget();
        if (statement instanceof RoutingStatementHandler) {
            statement = (StatementHandler) SystemMetaObject
                .forObject(statement).getValue("delegate");
        }
        return statement;
    }

    /**
     * intercept and set custom key generator
     * @param invocation invocation
     * @return result
     * @throws Throwable exception
     */
    Object injectKeyGen(Invocation invocation) throws Throwable {
        Object objRtn = invocation.proceed();
        MetaObject metaObjStatementHandler =
            SystemMetaObject.forObject(getStatementHandler(invocation));
        MappedStatement mappedStatement =
            (MappedStatement) metaObjStatementHandler.getValue("mappedStatement");

        if (SqlCommandType.INSERT.equals(mappedStatement.getSqlCommandType())) {
            BoundSql boundSql = (BoundSql) metaObjStatementHandler.getValue("boundSql");
            Object paramObj = boundSql.getParameterObject();
            if (paramObj instanceof MapperMethod.ParamMap) {
                @SuppressWarnings("unchecked")
                MapperMethod.ParamMap<Object> paramMap = 
                    (MapperMethod.ParamMap<Object>) paramObj;
                if (paramMap.containsKey(BaseMapper.PARAM_MODEL_KEY) ||
                        paramMap.containsKey(BaseMapper.PARAM_MODELS_KEY)) {
                    KeyGenerator keygen = mappedStatement.getKeyGenerator();
                    if (keygen.getClass().equals(Jdbc3KeyGenerator.class)) {
                        log.info("Replace JDBC3 key generator.");
                        SystemMetaObject.forObject(mappedStatement)
                            .setValue("keyGenerator",
                                new Jdbc3KeyGen(keygen));
                    }
                }
            }
        }
        return objRtn;
    }
}
