package jnpf.database.plugins;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.enums.DdConstants;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.baomidou.dynamic.datasource.tx.TransactionContext;
import jnpf.database.util.DynamicDataSourceUtil;
import jnpf.util.TenantHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;

/**
 * 主库读写分离
 */
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class})})
@Slf4j
public class MyDefaultMasterSlaveAutoRoutingPlugin implements Interceptor {

    protected DynamicRoutingDataSource dynamicDataSource;

    public MyDefaultMasterSlaveAutoRoutingPlugin(DataSource dataSource){
        this.dynamicDataSource = (DynamicRoutingDataSource) dataSource;
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 连接隔离模式多租户不处理, 当前非主数据源不处理, 未添加从库不处理
        if (TenantHolder.isRemote()
                || !DynamicDataSourceUtil.isPrimaryDataSoure()
                || !dynamicDataSource.getGroupDataSources().containsKey(DdConstants.SLAVE)) {
            return invocation.proceed();
        }
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        String pushedDataSource = null;
        try {
            // 存在事务只使用主库
            boolean hasTrans = TransactionSynchronizationManager.isActualTransactionActive();
            if (!hasTrans) {
                hasTrans = StringUtils.hasText(TransactionContext.getXID());
            }
            String dataSource = !hasTrans && SqlCommandType.SELECT == ms.getSqlCommandType() ? DdConstants.SLAVE : DynamicDataSourceUtil.dynamicDataSourceProperties.getPrimary();
            pushedDataSource = DynamicDataSourceContextHolder.push(dataSource);
            return invocation.proceed();
        } finally {
            if (pushedDataSource != null) {
                DynamicDataSourceContextHolder.poll();
            }
        }
    }
}