package top.imwork.assist.generator.config;

import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Copyright (C), 2015-2025
 * FileName: DynamicDataSourceManager
 * Author<作者姓名>:   silos
 * CreateTime<创建时间>:   2025/11/9 20:20
 * UpdateTime<修改时间>:   2025/11/9 20:20
 * Description〈功能简述〉: 数据源管理器
 * History<历史描述>:
 * Since<版本号>: 1.0.0
 */
@Component
public class DynamicDataSourceManager {

    private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();

    @Autowired
    private JdbcTemplate primaryJdbcTemplate;

    /**
     * 动态添加数据源
     */
    public void addDataSource(String key, DataSourceProperties properties) {
        DataSource dataSource = createDataSource(properties);
        dataSourceMap.put(key, dataSource);
    }

    /**
     * 动态移除数据源
     */
    public void removeDataSource(String key) {
        DataSource dataSource = dataSourceMap.remove(key);
        if (dataSource != null) {
            // 关闭数据源连接
            if (dataSource instanceof HikariDataSource) {
                ((HikariDataSource) dataSource).close();
            }
        }
    }

    /**
     * 获取数据源
     */
    public DataSource getDataSource(String key) {
        return dataSourceMap.get(key);
    }

    /**
     * 创建数据源
     */
    private DataSource createDataSource(DataSourceProperties properties) {
        return DataSourceBuilder.create()
                .url(properties.getUrl())
                .username(properties.getUsername())
                .password(properties.getPassword())
                .driverClassName(properties.getDriverClassName())
                .build();
    }

    /**
     * 获取特定数据源的 JdbcTemplate
     */
    public JdbcTemplate getJdbcTemplate(String dataSourceKey) {
        DataSource dataSource = dataSourceMap.get(dataSourceKey);
        if (dataSource == null) {
            throw new IllegalArgumentException("数据源不存在: " + dataSourceKey);
        }
        return new JdbcTemplate(dataSource);
    }
}