package com.sqx.sharding;

import com.zaxxer.hikari.HikariDataSource;
import com.zaxxer.hikari.metrics.micrometer.MicrometerMetricsTrackerFactory;
import io.micrometer.core.instrument.MeterRegistry;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.actuate.autoconfigure.metrics.MetricsAutoConfiguration;
import org.springframework.boot.actuate.autoconfigure.metrics.export.simple.SimpleMetricsExportAutoConfiguration;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.jdbc.DataSourceUnwrapper;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.util.Collection;

/**
 * @author wxkj0012
 */
@Configuration(proxyBeanMethods = false)
@AutoConfigureAfter({MetricsAutoConfiguration.class, DataSourceAutoConfiguration.class,
        SimpleMetricsExportAutoConfiguration.class})
@ConditionalOnClass({DataSource.class, MeterRegistry.class})
@ConditionalOnBean({DataSource.class, MeterRegistry.class})
@Slf4j
public class ShardingDataSourcePoolMetricsAutoConfiguration {

    @Configuration(proxyBeanMethods = false)
    @ConditionalOnClass({ShardingDataSource.class, HikariDataSource.class})
    static class ShardingDataSourceMetricsConfiguration {

        private final MeterRegistry registry;

        ShardingDataSourceMetricsConfiguration(MeterRegistry registry) {
            this.registry = registry;
        }

        @Autowired
        void bindMetricsRegistryToHikariDataSources(Collection<DataSource> dataSources) {
            for (DataSource dataSource : dataSources) {
                ShardingDataSource shardingDataSource = DataSourceUnwrapper.unwrap(dataSource, ShardingDataSource.class);
                Collection<DataSource> shardDataSources = shardingDataSource.getDataSourceMap().values();
                for (DataSource ds : shardDataSources) {
                    HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(ds, HikariDataSource.class);
                    if (hikariDataSource != null) {
                        bindMetricsRegistryToHikariDataSource(hikariDataSource);
                    }
                }
            }
        }

        private void bindMetricsRegistryToHikariDataSource(HikariDataSource hikari) {
            if (hikari.getMetricRegistry() == null && hikari.getMetricsTrackerFactory() == null) {
                try {
                    hikari.setMetricsTrackerFactory(new MicrometerMetricsTrackerFactory(this.registry));
                } catch (Exception ex) {
                    log.warn("Failed to bind Hikari metrics: {}}", ex.getMessage());
                }
            }
        }

    }
}