package org.jeecg.config.saas;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DataSourceProperty;
import com.baomidou.dynamic.datasource.creator.DefaultDataSourceCreator;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.jeecg.modules.system.service.impl.SysBaseApiImpl;
import org.jeecg.modules.system.util.SecurityUtil;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.sql.DataSource;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class TenantInterceptor implements HandlerInterceptor {

    private static final Logger log = LoggerFactory.getLogger(TenantInterceptor.class);

    @Resource
    private DataSource dataSource;

    @Resource(name = "dataSource")
    private DataSource masterDataSource;

    @Resource
    private DefaultDataSourceCreator dataSourceCreator;

    private final Set<String> dsCache = ConcurrentHashMap.newKeySet();
    private final ConcurrentHashMap<String, Long> errorDsCache = new ConcurrentHashMap<>();
    private static final long AUTO_RECOVER_INTERVAL = 60 * 1000;
    @Autowired
    private SysBaseApiImpl sysBaseApiImpl;

    @Override
    public boolean preHandle(HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull Object handler) throws Exception {
        log.info("======== TenantInterceptor preHandle triggered ========");
        String companyCode = request.getHeader("X-Tenant-Code"); // 后端内部依然用header获取，外面都叫“公司编号”
        // 校验公司编号
        if (companyCode == null || companyCode.trim().isEmpty()
                || "undefined".equalsIgnoreCase(companyCode.trim())
                || "null".equalsIgnoreCase(companyCode.trim())) {
//            log.warn("[公司校验] 请求未携带有效公司编号，拒绝访问！");
//            throw new RuntimeException("请填写正确的公司编号，禁止访问！");
            DynamicDataSourceContextHolder.push("master");
            return true;
        }
        // ====== 新增：禁止中文、只允许字母和数字 ======
        String code = companyCode.trim();
// 只允许字母和数字，正则可根据实际需求调整
        if (!code.matches("^[a-zA-Z0-9]+$")) {
            log.warn("[公司校验] 公司编号【{}】格式非法，只允许字母和数字！", code);
            throw new RuntimeException("公司编号只能包含字母和数字，不能有中文！");
        }
        if ("master".equalsIgnoreCase(companyCode.trim()) || "10086".equals(companyCode.trim())) {
            log.warn("[公司校验] 当前请求为平台账号，平台编号：{}", companyCode);
            DynamicDataSourceContextHolder.push("master");
            return true;
        }
        // 检查公司编号是否在错误缓存中，未超时则拒绝
        Long lastFailTime = errorDsCache.get(companyCode);
        if (lastFailTime != null) {
            long now = System.currentTimeMillis();
            if (now - lastFailTime < AUTO_RECOVER_INTERVAL) {
                log.warn("[公司校验] 公司编号【{}】上次校验失败，暂时无法使用！", companyCode);
                throw new RuntimeException("公司编号【" + companyCode + "】暂时无法使用，请稍后再试或联系管理员！");
            } else {
                errorDsCache.remove(companyCode);
                log.warn("[公司校验] 公司编号【{}】恢复可用。", companyCode);
            }
        }

        // 检查是否已存在
        DynamicRoutingDataSource dynamicRoutingDataSource = (DynamicRoutingDataSource) dataSource;
        Set<String> allDs = dynamicRoutingDataSource.getDataSources().keySet();
        if (allDs.contains(companyCode) || dsCache.contains(companyCode)) {
            DynamicDataSourceContextHolder.push(companyCode);
            log.info("[公司校验] 当前公司编号【{}】已校验通过。", companyCode);
            return true;
        }

        // 校验公司编号信息
        JdbcTemplate jdbc = new JdbcTemplate(masterDataSource);
        Map<String, Object> row;
        try {
            row = jdbc.queryForMap(
                    "SELECT code, db_url, db_name, db_username, db_password, db_driver FROM sys_data_source WHERE code=?",
                    companyCode
            );
        } catch (Exception e) {
            errorDsCache.put(companyCode, System.currentTimeMillis());
            log.warn("[公司校验] 公司编号【{}】不存在或信息无效，已被限制访问：{}", companyCode, e.getMessage());
            throw new RuntimeException("公司编号【" + companyCode + "】不存在或不可用，请联系管理员！");
        }

        String dbUrl = String.valueOf(row.get("db_url"));
        String dbName = String.valueOf(row.get("db_name"));
        String url;
        if (dbUrl.contains("?")) {
            url = dbUrl.replaceAll("/+$", "");
        } else {
            url = dbUrl.replaceAll("/+$", "") + "/" + dbName + "?characterEncoding=UTF-8&useSSL=false&serverTimezone=Asia/Shanghai";
        }

        DataSourceProperty dsp = new DataSourceProperty();
        dsp.setUrl(url);
        dsp.setUsername(String.valueOf(row.get("db_username")));
        String db_password = String.valueOf(row.get("db_password"));
        String pwd = SecurityUtil.jiemi(db_password);
        dsp.setPassword(pwd);
        dsp.setDriverClassName(String.valueOf(row.get("db_driver")));

        log.info("[公司校验] 正在校验公司编号【{}】……", companyCode);

        try {
            javax.sql.DataSource realDs = dataSourceCreator.createDataSource(dsp);
            dynamicRoutingDataSource.addDataSource(companyCode, realDs);
            dsCache.add(companyCode);
            errorDsCache.remove(companyCode);
            log.info("[公司校验] 公司编号【{}】校验通过。", companyCode);
        } catch (Exception ex) {
            errorDsCache.put(companyCode, System.currentTimeMillis());
            log.warn("[公司校验] 公司编号【{}】校验失败：{}", companyCode, ex.getMessage());
            throw new RuntimeException("公司编号【" + companyCode + "】暂时不可用，请联系管理员！");
        }

        DynamicDataSourceContextHolder.push(companyCode);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        DynamicDataSourceContextHolder.clear();
    }
}
