// DataSourceConfig.java
@Configuration
public class DataSourceConfig {

  @Bean
  public DataSource routingDataSource() {
    AbstractRoutingDataSource routingDataSource = new RoutingDatasource();
    routingDataSource.setDefaultTargetDataSource(createMySqlDataSource("jdbc:mysql://localhost:3306/spring"));

    Map<Object, Object> targetDataSources = new HashMap<>();
    targetDataSources.put(0, createPostgresqlDataSource("jdbc:postgresql://localhost:5433/postgres"));
    targetDataSources.put(1, createMariaDataSource("jdbc:mariadb://localhost:3305/maria"));
//    targetDataSources.put(2, createOracleDataSource("jdbc:oracle:thin:@localhost:1521/xe"));
//    targetDataSources.put(3, createSqlServerDataSource("jdbc:sqlserver://localhost:1433;DatabaseName=mssql"));

    routingDataSource.setTargetDataSources(targetDataSources);

    return routingDataSource;
  }

  private DataSource createMySqlDataSource(String url) {
    AtomikosDataSourceBean dataSource = new AtomikosDataSourceBean();

    Properties properties = new Properties();
    properties.setProperty("user", "root");
    properties.setProperty("password", "root");
    properties.setProperty("url", url);

    dataSource.setXaDataSourceClassName("com.mysql.jdbc.jdbc2.optional.MysqlXADataSource");
    dataSource.setUniqueResourceName(url.substring(0, url.length() > 45 ? 44 : url.length()));
    dataSource.setXaProperties(properties);

    return dataSource;
  }

  private DataSource createMariaDataSource(String url) {
    AtomikosDataSourceBean dataSource = new AtomikosDataSourceBean();

    Properties properties = new Properties();
    properties.setProperty("user", "root");
    properties.setProperty("password", "root");
    properties.setProperty("url", url);

    dataSource.setXaDataSourceClassName("org.mariadb.jdbc.MariaDbDataSource");
    dataSource.setUniqueResourceName(url.substring(0, url.length() > 45 ? 44 : url.length()));
    dataSource.setXaProperties(properties);

    return dataSource;
  }

  private DataSource createPostgresqlDataSource(String url) {
    AtomikosDataSourceBean dataSource = new AtomikosDataSourceBean();

    Properties properties = new Properties();
    properties.setProperty("user", "postgres");
    properties.setProperty("password", "root");
    properties.setProperty("url", url);

    dataSource.setXaDataSourceClassName("org.postgresql.xa.PGXADataSource");
    dataSource.setUniqueResourceName(url.substring(0, url.length() > 45 ? 44 : url.length()));
    dataSource.setXaProperties(properties);

    return dataSource;
  }

  private DataSource createOracleDataSource(String url) {
    AtomikosDataSourceBean dataSource = new AtomikosDataSourceBean();

    Properties properties = new Properties();
    properties.setProperty("user", "root");
    properties.setProperty("password", "root");
    properties.setProperty("URL", url);

    dataSource.setXaDataSourceClassName("oracle.jdbc.xa.client.OracleXADataSource");
    dataSource.setUniqueResourceName(url.substring(0, url.length() > 45 ? 44 : url.length()));
    dataSource.setXaProperties(properties);

    return dataSource;
  }

  private DataSource createSqlServerDataSource(String url) {
    AtomikosDataSourceBean dataSource = new AtomikosDataSourceBean();

    Properties properties = new Properties();
    properties.setProperty("user", "root");
    properties.setProperty("password", "root");
    properties.setProperty("URL", url);

    dataSource.setXaDataSourceClassName("com.microsoft.sqlserver.jdbc.SQLServerXADataSource");
    dataSource.setUniqueResourceName(url.substring(0, url.length() > 45 ? 44 : url.length()));
    dataSource.setXaProperties(properties);

    return dataSource;
  }

  @Bean
  public SqlSessionFactory sqlSessionFactory(DataSource dataSource) throws Exception {
    SqlSessionFactoryBean factory = new SqlSessionFactoryBean();
    factory.setDataSource(dataSource);
    factory.setTransactionFactory(new ManagedTransactionFactory());
    return factory.getObject();
  }
}

// RoutingMapperAspect.java
@Slf4j
@Aspect
@Configuration
public class RoutingMapperAspect {

  @Around("execution(* com.seungho..repository..*Mapper.*(..))")
  public Object aroundTargetMethod(ProceedingJoinPoint thisJoinPoint) {
    MethodSignature methodSignature = (MethodSignature) thisJoinPoint.getSignature();
    Class<?> mapperInterface = methodSignature.getDeclaringType();
    Method method = methodSignature.getMethod();
    Parameter[] parameters = method.getParameters();
    Object[] args = thisJoinPoint.getArgs();

    RoutingMapper routingMapper = mapperInterface.getDeclaredAnnotation(RoutingMapper.class);

    if (routingMapper != null) {
      String userId = findRoutingKey(parameters, args);
      Integer index = determineRoutingDataSourceIndex(userId);
      log.debug("index: {}", index);
      ContextHolder.set(index);
    }

    try {
      return thisJoinPoint.proceed();
    } catch (Throwable throwable) {
      throw new RuntimeException(throwable);
    } finally {
      ContextHolder.clear();
    }
  }

  private String findRoutingKey(Parameter[] parameters, Object[] args) {
    for (int i = 0; i < parameters.length; i++) {
      RoutingKey routingKey = parameters[i].getDeclaredAnnotation(RoutingKey.class);
      if (routingKey != null) {
        return String.valueOf(args[i]);
      }
    }
    throw new RuntimeException("can't find RoutingKey");
  }

  private Integer determineRoutingDataSourceIndex(String userId) {
    return Math.abs(userId.hashCode()) % 2;
  }
}

// FooService.java
@Service
public class FooService {

  @Autowired private UserMapper userMapper;

  @Autowired private CommonMapper commonMapper;

  @Transactional
  public void bar(String userId) {
    userMapper.insert(userId);
    commonMapper.insert(userId);
  }

  @Transactional
  public void barWithException(String userId) {
    userMapper.insert(userId);
    commonMapper.insert(userId);

    if ("test5".equals(userId)) {
      throw new IllegalArgumentException("Intended exception");
    }
  }

  @Transactional
  public void deleteAll(String userId) {
    userMapper.deleteAll(userId);
    commonMapper.deleteAll(userId);
  }

  public int getUserCount(String userId) { return userMapper.getCount(userId); }

  public int getCommonCount(String userId) { return commonMapper.getCount(userId); }
}

// UserMapper.java
@RoutingMapper
@Mapper
public interface UserMapper {

  @Insert("insert into user_test(user_id) values(#{userId})")
  void insert(@RoutingKey @Param("userId") String userId);

  @Select("select count(*) from user_test where user_id=#{userId}")
  int getCount(@RoutingKey @Param("userId") String userId);

  @Delete("delete from user_test")
  void deleteAll(@RoutingKey @Param("userId") String userId);
}

// CommonMapper.java
@Mapper
public interface CommonMapper {

  @Insert("insert into common_test(user_id) values(#{userId})")
  void insert(@Param("userId") String userId);

  @Select("select count(*) from common_test where user_Id=#{userId}")
  int getCount(@Param("userId") String userId);

  @Delete("delete from common_test")
  void deleteAll(@Param("userId") String userId);
}

// JtademoApplicationTests.java
@Slf4j
@RunWith(SpringRunner.class)
@SpringBootTest
public class JtademoApplicationTests {

	@Autowired private FooService fooService;

	@Test
	public void addAll() {
		log.debug("addAll test: {}", "=======================================================");
		String userId = "test5";
		fooService.deleteAll(userId);
		fooService.bar(userId);

		Assertions.assertThat(fooService.getUserCount(userId)).isEqualTo(1);
		Assertions.assertThat(fooService.getCommonCount(userId)).isEqualTo(1);
	}

	@Test
	public void addAllWithException() {

			String userId = "test5";
		try {
			log.debug("addAllWithException test: {}", "=======================================================");
			fooService.deleteAll(userId);
			fooService.barWithException(userId);

			Assertions.assertThat(fooService.getUserCount(userId)).isEqualTo(0);
			Assertions.assertThat(fooService.getCommonCount(userId)).isEqualTo(0);
		} catch (IllegalArgumentException e) {
			Assertions.assertThat(fooService.getUserCount(userId)).isEqualTo(0);
			Assertions.assertThat(fooService.getCommonCount(userId)).isEqualTo(0);
		}
	}

}