/*
 * Copyright 1999-2018 Alibaba Group Holding Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.druid.pool.ha.selector;

import com.alibaba.druid.pool.DruidDataSource;
import com.alibaba.druid.pool.ha.HighAvailableDataSource;
import com.alibaba.druid.support.logging.Log;
import com.alibaba.druid.support.logging.LogFactory;

import javax.sql.DataSource;

import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * A selector which uses java.util.Random to choose DataSource.
 *
 * @author DigitalSonic
 */
public class RandomDataSourceSelector implements DataSourceSelector {
    private static final String PROP_PREFIX = "druid.ha.random.";
    public static final String PROP_CHECKING_INTERVAL = PROP_PREFIX + "checkingIntervalSeconds";
    public static final String PROP_RECOVERY_INTERVAL = PROP_PREFIX + "recoveryIntervalSeconds";
    public static final String PROP_VALIDATION_SLEEP = PROP_PREFIX + "validationSleepSeconds";
    public static final String PROP_BLACKLIST_THRESHOLD = PROP_PREFIX + "blacklistThreshold";

    private static final Log LOG = LogFactory.getLog(RandomDataSourceSelector.class);

    private Random random = new Random();
    private List<DataSource> blacklist = new CopyOnWriteArrayList<DataSource>();
    private HighAvailableDataSource highAvailableDataSource;
    private RandomDataSourceValidateThread validateThread;
    private RandomDataSourceRecoverThread recoverThread;
    private Thread runningValidateThread;
    private Thread runningRecoverThread;

    private int checkingIntervalSeconds = RandomDataSourceValidateThread.DEFAULT_CHECKING_INTERVAL_SECONDS;
    private int recoveryIntervalSeconds = RandomDataSourceRecoverThread.DEFAULT_RECOVER_INTERVAL_SECONDS;
    private int validationSleepSeconds;
    private int blacklistThreshold = RandomDataSourceValidateThread.DEFAULT_BLACKLIST_THRESHOLD;

    public RandomDataSourceSelector(HighAvailableDataSource highAvailableDataSource) {
        this.highAvailableDataSource = highAvailableDataSource;
    }

    @Override
    public void init() {
        if (highAvailableDataSource == null) {
            LOG.warn("highAvailableDataSource is NULL!");
            return;
        }
        if (!highAvailableDataSource.isTestOnBorrow() && !highAvailableDataSource.isTestOnReturn()) {
            loadProperties();
            initThreads();
        } else {
            LOG.info("testOnBorrow or testOnReturn has been set to true, ignore validateThread");
        }
    }

    /**
     * Interrupt Threads if needed.
     */
    @Override
    public void destroy() {
        if (runningValidateThread != null) {
            runningValidateThread.interrupt();
            validateThread.setSelector(null);
        }
        if (runningRecoverThread != null) {
            runningRecoverThread.interrupt();
            recoverThread.setSelector(null);
        }
    }

    @Override
    public String getName() {
        return DataSourceSelectorEnum.RANDOM.getName();
    }

    @Override
    public DataSource get() {
        Map<String, DataSource> dataSourceMap = getDataSourceMap();
        if (dataSourceMap == null || dataSourceMap.isEmpty()) {
            return null;
        }

        Collection<DataSource> targetDataSourceSet = removeBlackList(dataSourceMap);
        removeBusyDataSource(targetDataSourceSet);
        DataSource dataSource = getRandomDataSource(targetDataSourceSet);
        return dataSource;
    }

    @Override
    public void setTarget(String name) {
        // do nothing
    }

    public Map<String, DataSource> getFullDataSourceMap() {
        if (highAvailableDataSource != null) {
            return highAvailableDataSource.getDataSourceMap();
        }
        return new HashMap<String, DataSource>();
    }

    public Map<String, DataSource> getDataSourceMap() {
        if (highAvailableDataSource != null) {
            return highAvailableDataSource.getAvailableDataSourceMap();
        }
        return new HashMap<String, DataSource>();
    }

    public List<DataSource> getBlacklist() {
        return blacklist;
    }

    public boolean containInBlacklist(DataSource dataSource) {
        return dataSource != null && blacklist.contains(dataSource);
    }

    public void addBlacklist(DataSource dataSource) {
        if (dataSource != null && !blacklist.contains(dataSource)) {
            blacklist.add(dataSource);
            if (dataSource instanceof DruidDataSource) {
                ((DruidDataSource) dataSource).setTestOnReturn(true);
            }
        }
    }

    public void removeBlacklist(DataSource dataSource) {
        if (containInBlacklist(dataSource)) {
            blacklist.remove(dataSource);
            if (dataSource instanceof DruidDataSource) {
                ((DruidDataSource) dataSource).setTestOnReturn(highAvailableDataSource.isTestOnReturn());
            }
        }
    }

    private void loadProperties() {
        checkingIntervalSeconds = loadInteger(PROP_CHECKING_INTERVAL, checkingIntervalSeconds);
        recoveryIntervalSeconds = loadInteger(PROP_RECOVERY_INTERVAL, recoveryIntervalSeconds);
        validationSleepSeconds = loadInteger(PROP_VALIDATION_SLEEP, validationSleepSeconds);
        blacklistThreshold = loadInteger(PROP_BLACKLIST_THRESHOLD, blacklistThreshold);
    }

    private int loadInteger(String name, int defaultValue) {
        if (name == null) {
            return defaultValue;
        }

        Properties properties = highAvailableDataSource.getConnectProperties();
        int value = defaultValue;
        try {
            if (properties.containsKey(name)) {
                value = Integer.parseInt(properties.getProperty(name));
            }
        } catch (Exception e) {
            LOG.error("Exception occurred while parsing " + name, e);
        }
        return value;
    }

    private void initThreads() {
        if (validateThread == null) {
            validateThread = new RandomDataSourceValidateThread(this);
            validateThread.setCheckingIntervalSeconds(checkingIntervalSeconds);
            validateThread.setValidationSleepSeconds(validationSleepSeconds);
            validateThread.setBlacklistThreshold(blacklistThreshold);
        } else {
            validateThread.setSelector(this);
        }
        if (runningValidateThread != null) {
            runningValidateThread.interrupt();
        }
        runningValidateThread = new Thread(validateThread, "RandomDataSourceSelector-validate-thread");
        runningValidateThread.start();

        if (recoverThread == null) {
            recoverThread = new RandomDataSourceRecoverThread(this);
            recoverThread.setRecoverIntervalSeconds(recoveryIntervalSeconds);
            recoverThread.setValidationSleepSeconds(validationSleepSeconds);
        } else {
            recoverThread.setSelector(this);
        }
        if (runningRecoverThread != null) {
            runningRecoverThread.interrupt();
        }
        runningRecoverThread = new Thread(recoverThread, "RandomDataSourceSelector-recover-thread");
        runningRecoverThread.start();
    }

    private Collection<DataSource> removeBlackList(Map<String, DataSource> dataSourceMap) {
        Collection<DataSource> dataSourceSet;
        if (blacklist == null || blacklist.isEmpty() || blacklist.size() >= dataSourceMap.size()) {
            dataSourceSet = dataSourceMap.values();
        } else {
            dataSourceSet = new HashSet<DataSource>(dataSourceMap.values());
            for (DataSource b : blacklist) {
                dataSourceSet.remove(b);
            }
            LOG.info(blacklist.size() + " Blacklist DataSource removed, return "
                    + dataSourceSet.size() + " DataSource(s).");
        }
        return dataSourceSet;
    }

    private void removeBusyDataSource(Collection<DataSource> dataSourceSet) {
        Collection<DataSource> busyDataSourceSet = new HashSet<DataSource>();
        for (DataSource ds : dataSourceSet) {
            if (ds instanceof DruidDataSource && ((DruidDataSource) ds).getPoolingCount() <= 0) {
                busyDataSourceSet.add(ds);
            }
        }
        if (!busyDataSourceSet.isEmpty() && busyDataSourceSet.size() < dataSourceSet.size()) {
            LOG.info("Busy DataSouces: " + busyDataSourceSet.size() + "/" + dataSourceSet.size());
            for (DataSource ds : busyDataSourceSet) {
                dataSourceSet.remove(ds);
            }
        }
    }

    private DataSource getRandomDataSource(Collection<DataSource> dataSourceSet) {
        DataSource[] dataSources = dataSourceSet.toArray(new DataSource[]{});
        if (dataSources != null && dataSources.length > 0) {
            return dataSources[random.nextInt(dataSourceSet.size())];
        }
        return null;
    }

    public HighAvailableDataSource getHighAvailableDataSource() {
        return highAvailableDataSource;
    }

    public RandomDataSourceValidateThread getValidateThread() {
        return validateThread;
    }

    public void setValidateThread(RandomDataSourceValidateThread validateThread) {
        this.validateThread = validateThread;
    }

    public RandomDataSourceRecoverThread getRecoverThread() {
        return recoverThread;
    }

    public void setRecoverThread(RandomDataSourceRecoverThread recoverThread) {
        this.recoverThread = recoverThread;
    }

    public int getCheckingIntervalSeconds() {
        return checkingIntervalSeconds;
    }

    public void setCheckingIntervalSeconds(int checkingIntervalSeconds) {
        this.checkingIntervalSeconds = checkingIntervalSeconds;
    }

    public int getRecoveryIntervalSeconds() {
        return recoveryIntervalSeconds;
    }

    public void setRecoveryIntervalSeconds(int recoveryIntervalSeconds) {
        this.recoveryIntervalSeconds = recoveryIntervalSeconds;
    }

    public int getValidationSleepSeconds() {
        return validationSleepSeconds;
    }

    public void setValidationSleepSeconds(int validationSleepSeconds) {
        this.validationSleepSeconds = validationSleepSeconds;
    }

    public int getBlacklistThreshold() {
        return blacklistThreshold;
    }

    public void setBlacklistThreshold(int blacklistThreshold) {
        this.blacklistThreshold = blacklistThreshold;
    }
}
