package com.qiaoba.common.database.monitor;

import cn.hutool.core.io.IoUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.qiaoba.common.base.constants.TenantConstant;
import com.qiaoba.common.base.enums.DataBaseEnum;
import com.qiaoba.common.database.config.DynamicDataSourceConfig;
import com.qiaoba.common.database.entity.DynamicDataSource;
import com.qiaoba.common.database.service.DynamicDatasourceService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 监控-数据源连接监控
 *
 * @author ailanyin
 * @version 1.0
 * @since 2023/6/14 14:48
 */
@Component
@Slf4j
public class DatasourceConnectionMonitor {

    @Resource
    private DynamicDataSourceConfig dynamicDataSourceConfig;
    @Resource
    private DynamicDatasourceService dynamicDatasourceService;

    private static Map<String, String> WAIT_UPDATE_DATASOURCE_STATUS = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        // 1s钟运行一次
        new Timer().schedule(new TimerTask() {
            @Override
            public void run() {
                // 项目加载数据源还未完成
                if (!DynamicDataSourceConfig.COMPLETE_LOAD_DATASOURCE) {
                    return;
                }
                log.trace("开始运行数据源监控线程, 时间: {}", new Date());
                for (String tenantId : DynamicDataSourceConfig.TENANT_IDS) {
                    Object primary = DynamicDataSourceConfig.PRIMARY_DATASOURCE_MAP.get(tenantId);
                    if (Objects.isNull(primary)) {
                        // 说明初始化主要数据源的时候出错
                        log.error("租户[{}]-目前主数据源异常, 开始切换备用数据源", tenantId);
                        // 切换备用数据源
                        changePrimary(tenantId);
                    } else {
                        DruidDataSource dataSource = (DruidDataSource) primary;
                        try {
                            Connection connection = dataSource.getConnection();
                            if (check(connection, tenantId)) {
                                // 说明数据源正常
                                log.trace("租户[{}]-目前主数据源正常, 无需切换数据源", tenantId);
                                // 主数据 处理任务
                                if (TenantConstant.DEFAULT_TENANT_ID.equals(tenantId)) {
                                    handleJob();
                                }
                                IoUtil.close(connection);
                                continue;
                            }

                            log.error("租户[{}]-目前主数据源异常, 开始切换备用数据源", tenantId);
                            IoUtil.close(connection);
                            // 主数据源异常 切换备用数据源
                            if (changePrimary(tenantId)) {
                                // 备用切换成功, 关闭原有异常数据源
                                IoUtil.close(dataSource);
                            }
                        } catch (SQLException e) {
                            //e.printStackTrace();
                        }
                    }
                }

                log.trace("结束运行数据源监控线程, 时间: {}", new Date());
            }
        }, 0, 1000);
    }

    private Boolean check(Connection conn, String tenantId) {
        Statement stmt = null;
        ResultSet rs = null;
        try {
            stmt = conn.createStatement();
            // 允许 2s 延时
            stmt.setQueryTimeout(2);
            rs = stmt.executeQuery(DataBaseEnum.getCheckSql(DynamicDataSourceConfig.TENANT_DATASOURCE_TYPE_MAP.get(tenantId)));
            return true;
        } catch (Exception e) {
            return false;
        } finally {
            IoUtil.close(rs);
            IoUtil.close(stmt);
        }
    }

    /**
     * 切换主数据源
     *
     * @param tenantId 租户ID
     */
    private Boolean changePrimary(String tenantId) {
        // 备用数据源
        List<DynamicDataSource> dataSources = DynamicDataSourceConfig.BACKUP_DATASOURCE_MAP.get(tenantId);
        if (Objects.isNull(dataSources)) {
            log.error("租户:[{}]切换备用数据源失败, 原因: 没有备用数据源", tenantId);
            return false;
        }

        Integer backupIndex = null;

        for (int i = 0; i < dataSources.size(); i++) {
            Object dynamicDataSource = DynamicDataSourceConfig.buildDataSource(tenantId, dataSources.get(i));
            // 不是空,说明备用数据源有用
            if (Objects.nonNull(dynamicDataSource)) {
                dynamicDataSourceConfig.changePrimaryDatasource(tenantId, dynamicDataSource);
                backupIndex = i;
                break;
            }
        }

        if (Objects.nonNull(backupIndex)) {
            // 切换成功
            DynamicDataSource dynamicDataSource = dataSources.get(backupIndex);
            // 更改数据库中该数据源为主要数据源
            if (Objects.nonNull(dynamicDataSource.getTenantId()) && !TenantConstant.DEFAULT_TENANT_ID.equals(dynamicDataSource.getTenantId())) {
                // 添加到待处理任务中
                WAIT_UPDATE_DATASOURCE_STATUS.put(dynamicDataSource.getTenantId(), dynamicDataSource.getDatasourceId());
            }
            // 备用数据源集合删除该数据源
            dataSources.remove((int) backupIndex);
            log.info("租户:[{}]切换备用数据源成功, 现主数据ID: {}", tenantId, dynamicDataSource.getDatasourceId());
            return true;
        } else {
            log.error("租户:[{}]切换备用数据源失败, 原因: 备用数据源均无效", tenantId);
            return false;
        }
    }

    private void handleJob() {
        Set<String> keys = WAIT_UPDATE_DATASOURCE_STATUS.keySet();
        for (String key : keys) {
            try {
                log.info("开始更新数据库中租户数据源状态, 租户ID: {}", key);
                // 防止更新过程中主数据挂了
                dynamicDatasourceService.changePrimaryDatasource(key, WAIT_UPDATE_DATASOURCE_STATUS.get(key));
                // 处理完成 删除任务
                WAIT_UPDATE_DATASOURCE_STATUS.remove(key);
                log.info("更新数据库中租户数据源状态完成, 租户ID: {}", key);
            } catch (Exception e) {
                log.error("更新数据库中租户数据源状态完成, 租户ID: {}, 失败原因: {}", key, e.getMessage());
            }
        }
    }
}
