package cn.axzo.msg.center.message.service.todo.mybatis.rowcountwarn;

import cn.axzo.framework.rocketmq.utils.TraceUtils;
import cn.axzo.msg.center.inside.notices.config.DiagnosisProps;
import cn.axzo.msg.center.message.service.todo.mybatis.SQLAccessor;
import cn.axzo.msg.center.utils.DingTalkService;
import cn.hutool.core.thread.NamedThreadFactory;
import cn.hutool.db.sql.StatementWrapper;
import com.mysql.cj.jdbc.ClientPreparedStatement;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.ListableBeanFactory;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author yanglin
 */
@Slf4j
public class ProxyStatement extends StatementWrapper {

    private static final AtomicLong ID_SEQ = new AtomicLong(0);
    private static final Executor EXECUTOR = new ThreadPoolExecutor(1, 3,
            1L, TimeUnit.MINUTES,
            new ArrayBlockingQueue<>(20),
            new NamedThreadFactory(
                    ProxyStatement.class.getSimpleName(), false));

    private final long id = ID_SEQ.incrementAndGet();
    private final PreparedStatement rawStatement;
    private final DiagnosisProps props;
    private final ListableBeanFactory beanFactory;
    private volatile String sql;
    private int rowCount = 0;
    private int periodWarnTimes = 0;
    private boolean closeWarned = false;

    public ProxyStatement(PreparedStatement rawStatement,
                          DiagnosisProps props,
                          ListableBeanFactory beanFactory) {
        super(rawStatement);
        this.rawStatement = rawStatement;
        this.props = props;
        this.beanFactory = beanFactory;
    }

    void rowAdvanced() {
        rowCount++;
        if (rowCountAboveThreshold() && canWarnPeriodically()) {
            warnFor("Threshold");
            periodWarnTimes++;
        }
    }

    private boolean canWarnPeriodically() {
        int maxWarnTimes = props.getRowCount().getPeriodMaxWarnTimes();
        return rowCount % 2000 == 0 && periodWarnTimes <= maxWarnTimes;
    }

    @Override
    public boolean execute() throws SQLException {
        sql = determineSQL();
        return super.execute();
    }

    @Override
    public ResultSet getResultSet() throws SQLException {
        ResultSet resultSet = super.getResultSet();
        if (resultSet == null) return null;
        return new ProxyResultSet(this, resultSet);
    }

    @Override
    public void close() throws SQLException {
        super.close();
        if (rowCountAboveThreshold() && !closeWarned) {
            warnFor("StatementClose");
            closeWarned = true;
        }
    }

    private boolean rowCountAboveThreshold() {
        return rowCount > props.getRowCount().determineWarningThreshold();
    }

    private void warnFor(String phase) {
        try {
            // make a copy
            int rowCount = this.rowCount;
            String traceId = TraceUtils.getOrCreateTraceId();
            EXECUTOR.execute(() -> {
                try {
                    sendWarn(phase, rowCount, traceId);
                } catch (Exception e) {
                    log.warn("[Task] Error sending warning", e);
                }
            });
        } catch (RejectedExecutionException e) {
            log.warn("[Submit] Error sending warning, task rejected");
        } catch (Exception e) {
            log.warn("[Submit] Error sending warning", e);
        }
    }

    private void sendWarn(String phase, int currentCount, String traceId) {
        DingTalkService dingTalkService = beanFactory
                .getBeanProvider(DingTalkService.class).getIfAvailable();
        if (dingTalkService == null) return;
        String title = "查询数据量太大";
        String markdown = "## " + title + "\n" +
                "查询Id: " + id + "\n\n" +
                "环境: " + props.getActiveProfile() + "\n\n" +
                "TraceId: " + traceId + "\n\n" +
                "阶段: " + phase + "\n\n" +
                "阈值数量: " + props.getRowCount().determineWarningThreshold() + "\n\n" +
                "当前数量: " + currentCount + "\n\n" +
                "SQL: " + sql + "\n";
        log.info("[Task] Sending row count warning: {}", markdown);
        dingTalkService.send(title, markdown);
    }

    private String determineSQL() {
        try {
            return getSQLFromStmt();
        } catch (Exception e) {
            log.warn("[SQL] Error building sql", e);
            return getSampleSQL();
        }
    }

    private String getSQLFromStmt() throws Exception {
        if (rawStatement instanceof com.zaxxer.hikari.pool.ProxyStatement) {
            Statement stmt = SQLAccessor.getStmtFromProxy((com.zaxxer.hikari.pool.ProxyStatement) rawStatement);
            if (stmt instanceof ClientPreparedStatement)
                return ((ClientPreparedStatement) stmt).asSql();
        }
        return getSampleSQL();
    }

    private String getSampleSQL() {
        try {
            return StringUtils.substringAfter(rawStatement.toString(), ":");
        } catch (Exception e) {
            return "unknown";
        }
    }

    public Statement unwrap() {
        return rawStatement;
    }
}