package com.younike.common.atomikos;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import com.younike.common.utils.string.StringUtils;
import org.mybatis.spring.transaction.SpringManagedTransaction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.DataSourceUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 数据连接池切换代码
 */
public class MyManagedTransaction extends SpringManagedTransaction {
    private static final Logger LOGGER = LoggerFactory.getLogger(SpringManagedTransaction.class);

    DynamicRoutingDataSource routingDataSource;
    ConcurrentHashMap<String, Connection> map = new ConcurrentHashMap<>();
    private boolean isConnectionTransactional;
    private Connection connection;
    private boolean autoCommit;
    private DataSource currentDataSource;

    public MyManagedTransaction(DataSource dataSource) {
        super(dataSource);
        this.routingDataSource = (DynamicRoutingDataSource)dataSource;
    }

    @Override
    public Connection getConnection() throws SQLException {
        String key = DynamicDataSourceContextHolder.getDataSourceLookupKey();
        // 默认的数据源
        if(StringUtils.isEmpty(key)){
            key = DataSourceTypeEnum.Master.getName();
        }
        if (map.containsKey(key)) {
            return map.get(key);
        }
        Map<String, DataSource> currentDataSources = routingDataSource.getCurrentDataSources();
        this.currentDataSource = currentDataSources.get(key);
        Connection con = currentDataSource.getConnection();

        this.connection = DataSourceUtils.getConnection(currentDataSource);
        this.autoCommit = this.connection.getAutoCommit();
        this.isConnectionTransactional = DataSourceUtils.isConnectionTransactional(this.connection, currentDataSource);

//        System.out.println("connection = " +currentDataSource.getUniqueResourceName()+ "数据源连接开启.... [" + connection+ "]");
        map.put(key, con);
        return con;
    }

    @Override
    public void close() throws SQLException {
//        System.out.println("connection = " +currentDataSource.getUniqueResourceName()+ "数据源连接关闭.... [" + connection+ "]");
        DataSourceUtils.releaseConnection(this.connection, this.currentDataSource);
        for (Connection connection : map.values()) {
            DataSourceUtils.releaseConnection(connection, this.currentDataSource);
        }

    }

    @Override
    public void commit() throws SQLException {
        if (this.connection != null && !this.isConnectionTransactional && !this.autoCommit) {
//            if (LOGGER.isDebugEnabled()) {
//                LOGGER.debug("Committing JDBC Connection [" + this.connection + "]");
//            }
//            System.out.println("connection = " +currentDataSource.getUniqueResourceName()+ "数据源事务提交.... [" + connection+ "]");
            this.connection.commit();
            for (Connection connection : map.values()) {
                connection.commit();
            }
        }

    }
//
//    /**
//     * {@inheritDoc}
//     */
    @Override
    public void rollback() throws SQLException {
        if (this.connection != null && !this.isConnectionTransactional && !this.autoCommit) {
//            if (LOGGER.isDebugEnabled()) {
//                LOGGER.debug("Rolling back JDBC Connection [" + this.connection + "]");
//            }
//            System.out.println("connection = " +currentDataSource.getUniqueResourceName()+ "数据源事务回滚.... [" + connection+ "]");
            this.connection.rollback();
            for (Connection connection : map.values()) {
                connection.rollback();
            }
        }

    }
}
