package com.ruoyi.framework.mybatis.interceptor;

import cn.hutool.extra.spring.SpringUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.baomidou.mybatisplus.extension.plugins.handler.TenantLineHandler;
import com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor;
import com.ruoyi.common.enums.DataSourceType;
import com.ruoyi.common.utils.SecurityUtils;
import com.ruoyi.framework.datasource.DynamicDataSourceContextHolder;
import com.ruoyi.framework.mybatis.holder.TenantHolder;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.StringValue;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class MyTenantInterceptor extends TenantLineInnerInterceptor {

    private static final String TenantName = "tenant_id";

    private static final Map<String, List<String>> notIncludeTenantIdTableNameMap = new ConcurrentHashMap<>();

    public MyTenantInterceptor() {

        // 获取主库数据源
        DruidDataSource dataSource = SpringUtil.getBean("masterDataSource");
        putTableMap(dataSource, DataSourceType.MASTER);

        // TODO 暂未配置从库

        setTenantLineHandler(new TenantLineHandler() {

            @Override
            public Expression getTenantId() {
                return new StringValue(SecurityUtils.getTenantId());
            }

            @Override
            public String getTenantIdColumn() {
                return TenantName;
            }

            @Override
            public boolean ignoreTable(String tableName) {

                // 存在忽略自动拼接租户注解，不拼接租户字段
                if (TenantHolder.isIgnore()) {
                    return true;
                }

                // 获取当前线程处理的数据源类型
                String dataSourceType = DynamicDataSourceContextHolder.getDataSourceType();
                // 获取当前数据原中不包涵租户字段的表名
                List<String> tableNameList = notIncludeTenantIdTableNameMap.get(dataSourceType);
                return tableNameList.contains(tableName);
            }

        });
    }

    /**
     * 添加不存在租户字段的表map
     *
     * @param dataSource     数据源
     * @param dataSourceType 数据源类型
     */
    private void putTableMap(DruidDataSource dataSource, DataSourceType dataSourceType) {
        // 获取链接url
        String url = dataSource.getUrl();
        // 获取数据库名
        String dbName = url.split("/")[3].split("\\?")[0];
        // 查询不包含租户字段的表名
        String sql = "SELECT DISTINCT table_name FROM information_schema.COLUMNS WHERE table_schema = ? AND column_name != ?";
        // 获取jdbc
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
        List<String> notIncludeTenantIdTableNameList = jdbcTemplate.queryForList(sql, String.class, dbName, TenantName);
        notIncludeTenantIdTableNameMap.put(dataSourceType.name(), notIncludeTenantIdTableNameList);
    }


}
