package com.usoftchina.uas.office.jdbc;

import com.usoftchina.uas.office.util.StringUtils;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.jdbc.datasource.lookup.DataSourceLookup;
import org.springframework.jdbc.datasource.lookup.JndiDataSourceLookup;
import org.springframework.util.Assert;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * 基于{@AbstractRoutingDataSource}，解决不能动态添加新数据源的问题
 *
 * @author yingp
 * @date 2017/7/27
 */
public class DynamicDataSource extends AbstractDataSource implements InitializingBean {

    private Map<Object, Object> targetDataSources;
    private Object defaultTargetDataSource;
    private boolean lenientFallback = true;
    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();
    private Map<Object, DataSource> resolvedDataSources;
    private DataSource resolvedDefaultDataSource;

    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources = targetDataSources;
    }

    public void setDefaultTargetDataSource(Object defaultTargetDataSource) {
        this.defaultTargetDataSource = defaultTargetDataSource;
    }

    public void setLenientFallback(boolean lenientFallback) {
        this.lenientFallback = lenientFallback;
    }

    public void setDataSourceLookup(DataSourceLookup dataSourceLookup) {
        this.dataSourceLookup = (DataSourceLookup) (dataSourceLookup != null ? dataSourceLookup : new JndiDataSourceLookup());
    }

    @Override
    public void afterPropertiesSet() {
        if (this.targetDataSources == null) {
            throw new IllegalArgumentException("Property \'targetDataSources\' is required");
        } else {
            this.resolvedDataSources = new HashMap(this.targetDataSources.size());
            Iterator var1 = this.targetDataSources.entrySet().iterator();

            while (var1.hasNext()) {
                Map.Entry entry = (Map.Entry) var1.next();
                Object lookupKey = this.resolveSpecifiedLookupKey(entry.getKey());
                DataSource dataSource = this.resolveSpecifiedDataSource(entry.getValue());
                this.resolvedDataSources.put(lookupKey, dataSource);
            }

            if (this.defaultTargetDataSource != null) {
                this.resolvedDefaultDataSource = this.resolveSpecifiedDataSource(this.defaultTargetDataSource);
            }

        }
    }

    protected Object resolveSpecifiedLookupKey(Object lookupKey) {
        return lookupKey;
    }

    protected DataSource resolveSpecifiedDataSource(Object dataSource) throws IllegalArgumentException {
        if (dataSource instanceof DataSource) {
            return (DataSource) dataSource;
        } else if (dataSource instanceof String) {
            return this.dataSourceLookup.getDataSource((String) dataSource);
        } else {
            throw new IllegalArgumentException("Illegal data source value - only [javax.sql.DataSource] and String supported: " + dataSource);
        }
    }

    @Override
    public Connection getConnection() throws SQLException {
        return this.determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return this.determineTargetDataSource().getConnection(username, password);
    }

    public String getDriverClassName() throws SQLException{
        HikariDataSource dataSource = unwrap(HikariDataSource.class);
        return dataSource.getDriverClassName();
    }

    public String getJdbcUrl() throws SQLException {
        HikariDataSource dataSource = unwrap(HikariDataSource.class);
        return dataSource.getJdbcUrl();
    }

    public String lookupSchema() throws SQLException{
        HikariDataSource dataSource = unwrap(HikariDataSource.class);
        return StringUtils.nvl(dataSource.getSchema(), dataSource.getUsername()).toUpperCase();
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> T unwrap(Class<T> iface) throws SQLException {
        if (iface.isInstance(this)) {
            return (T) this;
        }
        return determineTargetDataSource().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return (iface.isInstance(this) || determineTargetDataSource().isWrapperFor(iface));
    }

    protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        Object lookupKey = this.determineCurrentLookupKey();
        DataSource dataSource = (DataSource) this.resolvedDataSources.get(lookupKey);
        if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }

        if (dataSource == null) {
            throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
        } else {
            return dataSource;
        }
    }

    protected Object determineCurrentLookupKey() {
        return DataSourceHolder.get();
    }

    public boolean containsDataSource(Object key) {
        Object lookupKey = this.resolveSpecifiedLookupKey(key);
        return this.resolvedDataSources.containsKey(lookupKey);
    }

    public void addDataSource(Object key, Object value) {
        Object lookupKey = this.resolveSpecifiedLookupKey(key);
        if (!containsDataSource(lookupKey)) {
            DataSource dataSource = this.resolveSpecifiedDataSource(value);
            this.resolvedDataSources.put(lookupKey, dataSource);
        }
    }

    public void removeDataSource(Object key) {
        Object lookupKey = this.resolveSpecifiedLookupKey(key);
        if (!containsDataSource(lookupKey)) {
            this.resolvedDataSources.remove(key);
        }
    }

    public void removeAllResolved() {
        Iterator var1 = this.resolvedDataSources.entrySet().iterator();
        while (var1.hasNext()) {
            Map.Entry entry = (Map.Entry) var1.next();
            if (!"defaultDataSource".equals(entry.getKey())) {
                var1.remove();
            }
        }
    }
}
