package org.jeecg.ctop.material.service.impl;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.sql.Statement;
import java.util.Properties;

/**
 * 按 mapper 维度切换 TiDB 读引擎。
 * 当前报表查询统一走 tiflash,tidb。
 */
@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
@Component
public class SafeTiFlashInterceptor implements Interceptor {

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

    private static final String MATERIAL_MAPPER_PREFIX = "org.jeecg.ctop.material.mapper.";
    private static final String TIFLASH_ENGINE = "tiflash,tidb";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            Connection connection = (Connection) invocation.getArgs()[0];
            if (connection != null) {
                String statementId = resolveStatementId(invocation.getTarget());
                String engine = resolveEngine(statementId);
                if (engine != null) {
                    applyEngine(connection, engine);
                }
            }
        } catch (Exception e) {
            log.warn("set tidb engine failed", e);
        }
        return invocation.proceed();
    }

    private String resolveStatementId(Object target) {
        MetaObject metaObject = SystemMetaObject.forObject(target);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        if (mappedStatement == null) {
            mappedStatement = (MappedStatement) metaObject.getValue("mappedStatement");
        }
        return mappedStatement == null ? null : mappedStatement.getId();
    }

    private String resolveEngine(String statementId) {
        if (statementId == null || !statementId.startsWith(MATERIAL_MAPPER_PREFIX)) {
            return null;
        }
        return TIFLASH_ENGINE;
    }

    private void applyEngine(Connection connection, String engine) throws Exception {
        try (Statement statement = connection.createStatement()) {
            statement.execute("SET SESSION tidb_isolation_read_engines = '" + engine + "'");
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
    }
}
