/*
 * Copyright (c) 京北方信息技术股份有限公司 Corporation 2021 . All rights reserved.
 *
 */
package net.northking.cctp.common.db.mode;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

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

/**
 * 数据源集群
 *
 * @author : xiusong.xie
 * @version : 1.0
 */
public class DataSourceCluster
{
  /**
   * 日志
   */
  private final static Logger logger = LoggerFactory.getLogger(DataSourceCluster.class);
  /**
   * 集群名稱
   */
  private String name;
  /**
   * 只读数据源
   */
  private Map<String, DataSource> readSources = new HashMap<>();

  /**
   * 可读可写数据源
   */
  private Map<String, DataSource> writeSources = new HashMap<>();

  /**
   * 获取只读数据源集合
   *
   * @return 只读数据源集合
   */
  public Collection<DataSource> getReadSources()
  {
    return readSources.values();
  }

  /**
   * 获取可读可写数据源集合
   *
   * @return 可读可写数据源集合
   */
  public Collection<DataSource> getWriteSources()
  {
    return writeSources.values();
  }


  public DataSourceCluster()
  {
  }

  public DataSourceCluster(String name)
  {
    this.name = name;
  }

  public String getName()
  {
    return name;
  }

  public void setName(String name)
  {
    this.name = name;
  }

  public Connection getReadConnection() throws SQLException
  {
    return getReadConnection(null, null);
  }


  public Connection getReadConnection(String username, String password) throws SQLException
  {
    Connection connection = null;
    for (DataSource dataSource : readSources.values())
    {
      try
      {
        connection = getConnection(dataSource, username, password);
        break;
      } catch (SQLException ex)
      {
        logger.error("获取只读数据库连接发生异常", ex);
      }
    }
    if (connection == null)
    {
      logger.warn("无可用的只读数据库连接，切换至可读可写数据源");
      connection = getWriteConnection(username, password);
    }
    return connection;
  }

  public Connection getWriteConnection() throws SQLException
  {
    return getWriteConnection(null, null);
  }

  public Connection getWriteConnection(String username, String password) throws SQLException
  {
    Connection connection = null;
    SQLException sqlException = new SQLException("获取可写数据库连接时发生异常");
    for (DataSource dataSource : writeSources.values())
    {
      try
      {
        connection = getConnection(dataSource, username, password);
        break;
      } catch (SQLException ex)
      {
        logger.error("获取可写数据库连接时发生异常", ex);
        sqlException.addSuppressed(ex);
      }
    }
    if (connection == null)
    {
      throw sqlException;
    }
    return connection;
  }

  private Connection getConnection(DataSource dataSource, String username, String password) throws SQLException
  {
    Connection connection;
    if (StringUtils.hasText(username))
    {
      connection = dataSource.getConnection(username, password);
    } else
    {
      connection = dataSource.getConnection();
    }
    if (logger.isDebugEnabled())
    {
      logger.debug("从数据源[" + dataSource + "]获取数据库连接成功！");
    }

    return connection;
  }

  public void putReadDataSource(String key, DataSource dataSource)
  {
    readSources.put(key, dataSource);
  }

  public DataSource getReadDataSource(String key)
  {
    return readSources.get(key);
  }

  public void putWriteDataSource(String key, DataSource dataSource)
  {
    writeSources.put(key, dataSource);
  }

  public DataSource getWriteDataSource(String key)
  {
    return writeSources.get(key);
  }

  public boolean hasDataSource()
  {
    return writeSources.size() > 0;
  }
}
